/*
  Compile with:

  c++ -o evstress evstress.cpp -lcurl -levent_core
*/

#include <string.h>		// strdup()
#include <stdlib.h>
#include <sys/time.h>
#include <sys/stat.h>		// mkdir()
#include <curl/curl.h>
#include <event2/event.h>
#include <event2/event_struct.h> // struct event
#include <zlib.h>
#include <fstream>
#include <iostream>
#include <deque>
#include <map>

static char const* CacheDir = "/tmp/Cache/";

int MaxConcur = 1000;		// total concurrent connectins
int MaxUrl = 30;		// per connnection

#define CRAWL_TIMEOUT 30 
/// Bytes/s
#define CRAWL_LOW_SPEED_LIMIT 500
/// Seconds for the conn to be below LOW_SPEED_LIMIT
#define CRAWL_LOW_SPEED_TIME 20

using namespace std;

struct Retriever;
struct HostQueue;

static void multi_timer_cb(CURLM* multi, long timeout_ms, Retriever* retr);
static void sock_cb(CURL* e, curl_socket_t s, int what, Retriever* retr, void* sockp);

struct lessStr {
  bool operator()(const char* a, const char* b)
  {
    return strcmp(a, b) < 0;
  }
};

/* Retriever, dealing with multiple connections */
struct Retriever
{
  Retriever(event_base* base, char const* fileList) :
    base(base),
    multi(curl_multi_init()),
    ifs(fileList),
    uid(0),
    queued(0)
  {
    /* setup the generic multi interface options we want */
    curl_multi_setopt(multi, CURLMOPT_SOCKETFUNCTION, sock_cb);
    curl_multi_setopt(multi, CURLMOPT_SOCKETDATA, this);
    curl_multi_setopt(multi, CURLMOPT_TIMERFUNCTION, multi_timer_cb);
    curl_multi_setopt(multi, CURLMOPT_TIMERDATA, this);
  }

  ~Retriever() {
    curl_multi_cleanup(multi);
  }

  void	add(HostQueue* conn);

  void	action(int action, int fd);

  void	activate(event* ev, curl_socket_t s, int action);

  void	completed(HostQueue* conn);

  void	moreUrls();

  // key is service name <scheme, host, port>
  typedef std::map<char const*, HostQueue*, lessStr> HostQueues;
  HostQueues connections;

  event_base*	base;
  event		timer_event;
  CURLM*	multi;
  int		uid;		///< URL ID
  int		queued;		///< to connections
  ifstream	ifs;
};

static size_t header_cb(char* hdr, size_t size, size_t nmemb, HostQueue* conn);
static size_t write_cb(void* ptr, size_t size, size_t nmemb, HostQueue* conn);

/* Wrapper for cURL easy handle */
struct HostQueue
{
  /* Create an easy handle, and add it to the curl_multi of the Retriever */
  HostQueue(Retriever* retriever) :
    retr(retriever),
    easy(curl_easy_init()),
    status(0),
    ev(),			// clear for event_initialized() to work
    encoded(false)
  {
    stream.file = 0;
    if (!easy) {
      cerr << "curl_easy_init() failed, exiting!\n";
      exit(2);
    }
    curl_easy_setopt(easy, CURLOPT_HEADERFUNCTION, header_cb);
    curl_easy_setopt(easy, CURLOPT_HEADERDATA, this);
    curl_easy_setopt(easy, CURLOPT_WRITEFUNCTION, write_cb);
    curl_easy_setopt(easy, CURLOPT_WRITEDATA, this);
    curl_easy_setopt(easy, CURLOPT_PRIVATE, this);
    //curl_easy_setopt(easy, CURLOPT_VERBOSE, 1L);
    curl_easy_setopt(easy, CURLOPT_CONNECTTIMEOUT, CRAWL_TIMEOUT);
    curl_easy_setopt(easy, CURLOPT_TIMEOUT, CRAWL_TIMEOUT*10);
    curl_easy_setopt(easy, CURLOPT_LOW_SPEED_LIMIT, CRAWL_LOW_SPEED_LIMIT);
    curl_easy_setopt(easy, CURLOPT_LOW_SPEED_TIME, CRAWL_LOW_SPEED_TIME);
    curl_easy_setopt(easy, CURLOPT_HTTPAUTH, CURLAUTH_NONE);
    curl_easy_setopt(easy, CURLOPT_ACCEPT_ENCODING, "gzip");
    curl_easy_setopt(easy, CURLOPT_HTTP_CONTENT_DECODING, 0);
    curl_easy_setopt(easy, CURLOPT_HTTP_TRANSFER_DECODING, 1);
    curl_easy_setopt(easy, CURLOPT_DNS_CACHE_TIMEOUT, 60*5);
    curl_easy_setopt(easy, CURLOPT_COOKIEFILE, "");
  }

  ~HostQueue() {
    curl_multi_remove_handle(retr->multi, easy);
    curl_easy_cleanup(easy);
  }

  void	add(char const* url, int uid);

  void	fetch();

  void	completed(CURLcode code);

  Retriever*	retr;
  CURL*		easy;
  int		status;		///< HTTP status
  event		ev;		///< event monitoring connection

  char		contentsFile[PATH_MAX];	///< file where to store contents
  union {
    gzFile	gzip;		///< compressed file
    FILE*	file;		///< plain file
  }		stream;		///< stream where to store file

  bool		encoded;

  std::deque<std::pair<char const*, int> >	urls;
};

void Retriever::completed(HostQueue* conn)
{
  queued--;
  moreUrls();
  if (conn->urls.size())
    conn->fetch();			// get next
  else {
    // remove from connections
    for (HostQueues::iterator cit = connections.begin(); cit != connections.end(); cit++)
      if (cit->second == conn) {
	connections.erase(cit);
	break;
      }
    delete conn;
  }
}

void HostQueue::add(char const* url, int uid)
{
  if (urls.size() < MaxUrl)
    urls.push_back(make_pair(strdup(url), uid));
  // drop the others
}

void HostQueue::fetch()
{
  pair<char const*, int> url_id = urls.front();
  urls.pop_front();
  snprintf(contentsFile, sizeof(contentsFile), "%s%.4x",
	   CacheDir, url_id.second);
  curl_easy_setopt(easy, CURLOPT_URL, url_id.first);
  curl_multi_add_handle(retr->multi, easy);
  free((void*)url_id.first);
}

////////////////////////////////////////////////////////////////////////
// libevent callbacks

/* Update the event timer after curl_multi library calls */
static void multi_timer_cb(CURLM* multi, long timeout_ms, Retriever* retr)
{
  if (timeout_ms == -1L)
    return;			// no timeout

  struct timeval timeout;
  timeout.tv_sec = timeout_ms / 1000;
  timeout.tv_usec = (timeout_ms % 1000) * 1000;
  evtimer_add(&retr->timer_event, &timeout);
}

/* Called by libevent when action is detected on a socket handled by libcurl */
static void action_cb(int fd, short kind, void* userp)
{
  Retriever* retr = (Retriever*)userp;
  int action =
    (kind & EV_READ ? CURL_CSELECT_IN : 0) |
    (kind & EV_WRITE ? CURL_CSELECT_OUT : 0);

  retr->action(action, fd);
}

void Retriever::action(int action, int fd)
{
  int running;
  CURLMcode code = curl_multi_socket_action(multi, fd, action, &running);
  // find out which operation completed
  int msgs;
  do {
    CURLMsg* msg = curl_multi_info_read(multi, &msgs);
    if (msg && msg->msg == CURLMSG_DONE) {
      // find the HostQueue from the returning handle
      HostQueue* connection;
      curl_easy_getinfo(msg->easy_handle, CURLINFO_PRIVATE, &connection);
      connection->completed(msg->data.result);
    }
  } while (msgs);
}

void HostQueue::completed(CURLcode code)
{
  if (stream.file) {
    if (encoded)
      fclose(stream.file);
    else
      gzclose(stream.gzip);
    stream.file = 0;
  }
  if (code != CURLE_OK) {
    clog << curl_easy_strerror(code) << endl;
    unlink(contentsFile);
  }
  curl_multi_remove_handle(retr->multi, easy);
  retr->completed(this);
}

/* Called by libevent when our timeout expires */
static void timer_cb(int fd, short kind, void* userp)
{
  Retriever* retr = (Retriever*)userp;
  retr->action(0, CURL_SOCKET_TIMEOUT);
}

////////////////////////////////////////////////////////////////////////
// cURL callbacks

/* CURLMOPT_SOCKETFUNCTION */
static void sock_cb(CURL* e, curl_socket_t s, int what, Retriever* retr, void* sockp)
{
  HostQueue* connection;
  curl_easy_getinfo(e, CURLINFO_PRIVATE, &connection);
  event* ev = &connection->ev;
  if (what == CURL_POLL_REMOVE)
    event_del(ev);
  else
    retr->activate(ev, s, what);
}

void Retriever::activate(event* ev, curl_socket_t s, int action)
{
  int kind = (action & CURL_POLL_IN ? EV_READ : 0)
    | (action & CURL_POLL_OUT ? EV_WRITE : 0)
    | EV_PERSIST;

  if (event_initialized(ev))
    event_del(ev);
  event_assign(ev, base, s, kind, action_cb, this);
  event_add(ev, NULL);
}

size_t header_cb(char* hdr, size_t size, size_t nmemb, HostQueue* conn)
{
  if (!strncasecmp(hdr, "HTTP/", 5))
    conn->status = atoi(hdr + 9);
  else if (!strncasecmp(hdr, "Content-Encoding: ", 18))
    // we must disable curl --without-zlib, or it will uncompress
    conn->encoded = !strncmp(hdr + 18, "gzip", 4);
  return size * nmemb;
}

/* CURLOPT_WRITEFUNCTION */
static size_t write_cb(void* ptr, size_t size, size_t nmemb, HostQueue *conn)
{
  size_t blocksize = size * nmemb;

  // delayed open of file, since we may not get anything, if
  // for instance file is not-modified-since or contents is too big
  if (!conn->stream.file) {
    char const* file = conn->contentsFile;

    // check magic number
    if (conn->encoded &&
	(*(ushort*)ptr != 0x8b1f && *(ushort*)ptr != 0x1f8b))
      conn->encoded = false; // http server cheated

    if (conn->encoded)
      conn->stream.file = fopen(file, "w+b");
    else
      conn->stream.gzip = gzopen(file, "wb");

    if (conn->stream.file == NULL) {
      clog << "Could not open cache file: " << conn->contentsFile << endl;
      return 0;
    }
  }
  if (conn->encoded)
    return fwrite(ptr, size, nmemb, conn->stream.file);
  else
    return gzwrite(conn->stream.gzip, ptr, blocksize);
}

void Retriever::moreUrls()
{
  char line[200];
  while (queued < MaxConcur && ifs.getline(line, sizeof(line))) {
    char* host = strchr(line+7, '/');
    char const* service = strndup(line, host-line);
    HostQueues::const_iterator cit = connections.find(service);
    if (cit == connections.end()) {
      HostQueue* conn = new HostQueue(this);
      connections[service] = conn;
      conn->add(line, uid);
      conn->fetch();
    } else {
      HostQueue* conn = cit->second;
      conn->add(line, uid);
    }
    uid++;
    queued++;
  }
}

int main(int argc, char **argv)
{
  if (argc != 2) {
    cerr << "Usage: evstress url-list-file" << endl;
    return -1;
  }
  char const* fileList = argv[1];
  mkdir(CacheDir, 0777);
  event_base* base = event_base_new();
  Retriever retr(base, fileList);

  evtimer_assign(&retr.timer_event, base, timer_cb, &retr);

  retr.moreUrls();

  event_base_dispatch(base);
  return 0;
}