#include <curl/curl.h>
#include <errno.h>
#include <libssh2.h>
#include <stdio.h>
#include <stdlib.h>
#include <signal.h>


const char* kPublicKey = "key.pub";
const char* kPrivateKey = "key.prv";
const char* kUsername = "cecilia:";
const char* kUrl = "sftp://127.0.0.1/tmp/test.txt";
const char* kData = "hello world";

struct Data
{
  const char* data;
  int length;
};

struct Sockets
{
  curl_socket_t* sockets;
  int count;
};

struct ReadWriteSockets
{
  struct Sockets read, write;
};

/**
 * Initialise curl, and prints versions of ssh & curl.
 */
void initialise()
{
  const char* sshVersion = libssh2_version(0);
  printf("libssh2 version: %s\n", sshVersion);

  if (0 != curl_global_init(CURL_GLOBAL_ALL))
  {
    fprintf(stderr, "Failed to initialise CURL!\n");
  }

  curl_version_info_data *curlVersion = curl_version_info(CURLVERSION_NOW);
  printf("libcurl version: %s\n", curlVersion->version);
}

/**
 * Remove a file descriptor from a sockets array.
 */
void removeFd(struct Sockets* sockets, curl_socket_t fd)
{
  int i;
  for (i = 0; i < sockets->count; ++i)
  {
    if (sockets->sockets[i] == fd)
    {
      memmove(&sockets->sockets[i], &sockets->sockets[i + 1],
            sizeof(curl_socket_t) * (sockets->count - i - 1));
      --sockets->count;
    }
  }
}

/**
 * Add a file descriptor to a sockets array.
 */
void addFd(struct Sockets* sockets, curl_socket_t fd)
{
  // To ensure we only have each file descriptor once, we remove it then add
  // it again.
  removeFd(sockets, fd);
  sockets->sockets = realloc(sockets->sockets,
        sizeof(curl_socket_t) * (sockets->count + 1));
  sockets->sockets[sockets->count] = fd;
  ++sockets->count;
}

/**
 * Callback invoked by curl when user data is to be sent.
 */
int curlReadCallback( void *ptr, size_t size, size_t nmemb, void *readdata)
{
  printf("Curl asked for %i bytes of data.\n", size * nmemb);
  struct Data* data = (struct Data*)readdata;
  int numBytes = size * nmemb;
  if (numBytes > data->length)
  {
    numBytes = data->length;
  }
  memcpy(ptr, data->data, numBytes);
  data->data += numBytes;
  data->length -= numBytes;
  return numBytes;
}

/**
 * Callback invoked by curl when data has been retrieved.
 */
size_t curlRequestCallback(void *ptr, size_t size, size_t nmemb, void *userdata)
{
  printf("Curl retrieved data: ");
  fwrite(ptr, size, nmemb, stdout);
  fputc('\n', stdout);
  return size*nmemb;
}

/**
 * Callback invoked by curl to notify of debug messages.
 */
int curlDebugCallback(CURL *curl, curl_infotype type, char *data, size_t size, void *userdata)
{
  if (type == CURLINFO_TEXT)
  {
    printf("Curl debug: ");
    fwrite(data, size, 1, stdout);
    fputc('\n', stdout);
  }
  return 0;
}

/**
 * Callback invoked by curl to poll reading / writing of a socket.
 */
int curlSocketCallback(CURL *easy, curl_socket_t s, int action,
    void *userp, void *socketp)
{
  struct ReadWriteSockets* sockets = userp;
  switch (action)
  {
  case CURL_POLL_NONE:
    printf("Socket action: poll none (%i)\n", s);
    break;
  case CURL_POLL_IN:
    printf("Socket action: poll in (%i)\n", s);
    break;
  case CURL_POLL_OUT:
    printf("Socket action: poll out (%i)\n", s);
    break;
  case CURL_POLL_INOUT:
    printf("Socket action: poll in out (%i\n)", s);
    break;
  case CURL_POLL_REMOVE:
    printf("Socket action: poll remove (%i)\n", s);
    break;
  }
  if (action == CURL_POLL_IN || action == CURL_POLL_INOUT)
  {
    addFd(&sockets->read, s);
  }
  else
  {
    removeFd(&sockets->read, s);
  }

  if (action == CURL_POLL_OUT || action == CURL_POLL_INOUT)
  {
    addFd(&sockets->write, s);
  }
  else
  {
    removeFd(&sockets->write, s);
  }
  return 0;
}

/**
 * Callback invoked by curl to set a timeout.
 */
int curlTimerCallback(CURLM *multi, long timeout_ms, void *userp)
{
  printf("Curl timer: %ims\n", timeout_ms);
  struct timeval* timeout = userp;
  if (timeout_ms != -1)
  {
    gettimeofday(timeout, 0);
    timeout->tv_usec += timeout_ms * 1000;
  }
  else
  {
    timeout->tv_sec = -1;
  }
  return 0;
}

/**
 * Create the curl multi object.
 */
CURLM* createMulti(struct ReadWriteSockets* sockets, struct timeval* timeout)
{
  CURLM* result = curl_multi_init();
  curl_multi_setopt(result, CURLMOPT_SOCKETFUNCTION, curlSocketCallback);
  curl_multi_setopt(result, CURLMOPT_SOCKETDATA, sockets);

  curl_multi_setopt(result, CURLMOPT_TIMERFUNCTION, curlTimerCallback);
  curl_multi_setopt(result, CURLMOPT_TIMERDATA, timeout);
  return result;
}

/**
 * Start a curl upload.
 */
CURL* startUpload(struct Data* data)
{
  CURL* curl = curl_easy_init();
  curl_easy_setopt(curl, CURLOPT_INFILESIZE, data->length);
  curl_easy_setopt(curl, CURLOPT_UPLOAD, 1);

  curl_easy_setopt(curl, CURLOPT_SSH_PRIVATE_KEYFILE, kPrivateKey);
  curl_easy_setopt(curl, CURLOPT_SSH_PUBLIC_KEYFILE, kPublicKey);
  curl_easy_setopt(curl, CURLOPT_VERBOSE, 1);
  curl_easy_setopt(curl, CURLOPT_FAILONERROR, 1);
  curl_easy_setopt(curl, CURLOPT_USERPWD, kUsername);
  curl_easy_setopt(curl, CURLOPT_URL, kUrl);

  curl_easy_setopt(curl, CURLOPT_READFUNCTION, curlReadCallback);
  curl_easy_setopt(curl, CURLOPT_READDATA, data);
  curl_easy_setopt(curl, CURLOPT_WRITEFUNCTION, curlRequestCallback);
  curl_easy_setopt(curl, CURLOPT_DEBUGFUNCTION, curlDebugCallback);
  return curl;
}

/**
 * Update a fd_set with all of the sockets in use.
 */
void updateFdSet(struct Sockets* sockets, fd_set* fdset, int* maxFd)
{
  int i;
  for (i = 0; i < sockets->count; ++i)
  {
    FD_SET(sockets->sockets[i], fdset);
    if (*maxFd < sockets->sockets[i] + 1)
    {
      *maxFd = sockets->sockets[i] + 1;
    }
  }
}

void notifyCurl(CURL* curl, curl_socket_t socket, int evBitmask, const char* info)
{
  printf("Got a %s (%i)\n", info, socket);
  int numhandles = 0;
  CURLMcode result = curl_multi_socket_action(curl, socket, evBitmask, &numhandles);
  if (result != CURLM_OK && result != CURLM_CALL_MULTI_PERFORM)
  {
    fprintf(stderr, "Curl error on %s: %i (%s)\n", info, result, curl_multi_strerror(result));
  }
}

/**
 * Invoke curl when a file descriptor is set.
 */
void checkFdSet(CURL* curl, struct Sockets* sockets, fd_set* fdset,
      int evBitmask, const char* name)
{
  int i;
  for (i = 0; i < sockets->count; ++i)
  {
    if (FD_ISSET(sockets->sockets[i], fdset))
    {
      notifyCurl(curl, sockets->sockets[i], evBitmask, name);
    }
  }
}

/**
 * Check for curl completion.
 */
int checkForCompletion(CURLM* curl, int* success)
{
  printf("Checking for completion...\n");
  int numMessages;
  CURLMsg* message;
  int result = 0;
  *success = 0;
  while ((message = curl_multi_info_read(curl, &numMessages)) != 0)
  {
    if (message->msg == CURLMSG_DONE)
    {
      printf("Curl done: %i\n", message->data.result);
      result = 1;
      if (message->data.result == CURLE_OK)
      {
        printf("Request successful.\n");
        *success = 1;
      }
      else
      {
        printf("Request failed: %s\n", curl_easy_strerror(message->data.result));
        *success = 0;
      }
    }
    else
    {
      fprintf(stderr, "Got an unexpected message from curl: %i\n", message->msg);
    }
  }
  return result;
}

int getUsTimeout(struct timeval* timeout)
{
  struct timeval now;
  gettimeofday(&now, 0);
  int result = (timeout->tv_sec - now.tv_sec) * 1000000 + timeout->tv_usec - now.tv_usec;
  if (result < 0)
  {
    result = 0;
  }
  return result;
}

int run(CURLM* curl, struct ReadWriteSockets* sockets, struct timeval* timeout)
{
  int success = 0;
  fd_set readSet, writeSet;
  while (!checkForCompletion(curl, &success))
  {
    FD_ZERO(&readSet);
    FD_ZERO(&writeSet);
    int maxFd = 0;
    updateFdSet(&sockets->read, &readSet, &maxFd);
    updateFdSet(&sockets->write, &writeSet, &maxFd);

    struct timeval tv = {10, 0};
    if (timeout->tv_sec != -1)
    {
      int usTimeout = getUsTimeout(timeout);
      tv.tv_sec = usTimeout / 1000000;
      tv.tv_usec = usTimeout % 1000000;
    }

    printf("Running select (maxfd %i).\n", maxFd);
    if (select(maxFd, &readSet, &writeSet, NULL, &tv) == -1)
    {
      fprintf(stderr, "Select failed: %s\n", strerror(errno));
    }

    // Check the sockets for reading / writing
    checkFdSet(curl, &sockets->read, &readSet, CURL_CSELECT_IN, "read");
    checkFdSet(curl, &sockets->write, &writeSet, CURL_CSELECT_OUT, "write");

    if (timeout->tv_sec != -1 && getUsTimeout(timeout) == 0)
    {
      // Curl's timer has elapsed.
      notifyCurl(curl, CURL_SOCKET_TIMEOUT, 0, "timeout");
    }
  }
  return success;
}

int main()
{
  initialise();
  
  struct ReadWriteSockets sockets = {{0, 0}, {0, 0}};
  struct timeval timeout = {-1, 0};
  CURLM* curlMulti = createMulti(&sockets, &timeout);

  struct Data data = {kData, strlen(kData)};
  CURL* curl = startUpload(&data);
  curl_multi_add_handle(curlMulti, curl);

  if (!run(curlMulti, &sockets, &timeout))
  {
    fprintf(stderr, "Test failed.\n");
    return EXIT_FAILURE;
  }
  return EXIT_SUCCESS;
}