#define _GNU_SOURCE
#include <stdlib.h>
#include <unistd.h>
#include <getopt.h>
#include <sys/time.h>
#include <time.h>
#include <ctype.h>
#include <assert.h>

#include <event2/event.h>
#include <event2/event_compat.h>
#include <event2/event_struct.h>

#include "memcache.h"
#include "config.h"
#include "common.h"

/*
 * Test:
 *  * zero-length entries
 */

static struct event_base *ev_base;
static struct memcache *mc;
static struct config_endpoint server_endpoint;
static char *data_1 = "rei4quohV8Oocio1ua0co8ni4Ae1re4houcheixahchoh3ioghie0aShooShoh6Ahboequ9eiX5eashuu6Chu1quo"
                            "o0suph7cheiyai1ea0ooh7Aevoo4feihubupohDeephahwee2Ooz7chiediec7neit7keTh6xuheash8chaeKa5vi"
                            "ekooqu7ooj6Eezooroi6Nequ9ca2yi6iSoigh3loowaey9eiphaphaiJ0souy7wohpa9eXo5Ahu2sa";
static char *data_2 = "iefaek7ighi5UpueThageish5ieshohyeil1raiceerahjahng5ui7vuzie9quu4dai5ar2aiXi5ieth4looweigi"
                            "e3fo5ieri1queengaiphuaghaic1xahvoo9joo6baiNaig8puCootheowah4moocohDoiquoh3quieka5ao3aeNg9"
                            "Aimei1soangu4Duch5pho5buu2ohzaich4chahz9iTh3Pei4beep1ongie6au1aafoosh2vierei5E";
                            
#define BENCHMARK_KEY "memcache_benchmark"
#define BENCHMARK_KEY_MAX 256
#define BENCHMARK_DATA_MAX 1024 * 1024

#define MIN(a, b) ((a) < (b) ? (a) : (b))
#define MAX(a, b) ((a) > (b) ? (a) : (b))

struct common {
    unsigned int max_connections;
    char pipeline_requests;
    
    struct timeval start;
} common;

struct benchmark_fetch {
    unsigned int concurrent_ops;
    unsigned int total_ops;
    
    const char *key_prefix;
    unsigned int key_len_min, key_len_max, key_count;
    unsigned int data_len_min, data_len_max;
    
    int keys_stored;
    int cur_ops;
    int op_count;
    
    struct key_buf {
        char buf[BENCHMARK_KEY_MAX];
        struct memcache_key key;
    } *keys;

} benchmark_fetch;

void benchmark_cb (struct memcache_req *req, void *arg);
void benchmark_fetch_fn (void);

enum option_test {
    TEST_INVALID,
    TEST_COMMON,

    TEST_BENCH_FETCH,
};

enum option_code {
    OPT_CODE_INVALID,

    COMMON_CONN_MAX,
    COMMON_PIPELINE,
    BENCH_FETCH_REQ_CONCURRENCY,
    BENCH_FETCH_REQ_AMOUNT,
    BENCH_FETCH_KEY_PREFIX,
    BENCH_FETCH_KEY_LEN_MIN,
    BENCH_FETCH_KEY_LEN_MAX,
    BENCH_FETCH_KEY_COUNT,
    BENCH_FETCH_DATA_LEN_MIN,
    BENCH_FETCH_DATA_LEN_MAX,

    OPT_CODE_MAX,
};

enum option_type {
    OPT_TYPE_NONE,
    OPT_TYPE_BOOL,
    OPT_TYPE_UINT,
    OPT_TYPE_STR,
};

static struct test {
    char *name;
    memcache_cb cb_fn;

    void (*test_fn) (void);

    enum option_test code;

    char *descr;

} test_list[] = {
    { "benchmark_fetch",        &benchmark_cb,  &benchmark_fetch_fn,    TEST_BENCH_FETCH,
        "Measure the speed of fetch requests"                                               },
    { 0,                        0,              0,                      0,  0               }
};

static struct option options[] = {
    { "conn-max",           required_argument,  NULL,   COMMON_CONN_MAX             },
    { "pipeline",           required_argument,  NULL,   COMMON_PIPELINE             },
    { "req-concurrency",    required_argument,  NULL,   BENCH_FETCH_REQ_CONCURRENCY },
    { "req-amount",         required_argument,  NULL,   BENCH_FETCH_REQ_AMOUNT      },
    { "key-prefix",         required_argument,  NULL,   BENCH_FETCH_KEY_PREFIX      },
    { "key-len-min",        required_argument,  NULL,   BENCH_FETCH_KEY_LEN_MIN     },
    { "key-len-max",        required_argument,  NULL,   BENCH_FETCH_KEY_LEN_MAX     },
    { "key-count",          required_argument,  NULL,   BENCH_FETCH_KEY_COUNT       },
    { "data-len-min",       required_argument,  NULL,   BENCH_FETCH_DATA_LEN_MIN    },
    { "data-len-max",       required_argument,  NULL,   BENCH_FETCH_DATA_LEN_MAX    },
    { 0,                    0,                  0,      0                           },
};

static struct opt {
    enum option_test test;
    enum option_code code;
    enum option_type type;
    
    union opt_type_data {
        struct  {
            char *value;
            char default_value;
        } bool;

        struct {
            unsigned int *value;
            unsigned int default_value;
        } uint;

        struct {
            const char **value;
            const char *default_value;
        } str;
    } data;
    
    const char *name;
    const char *descr;
} option_info[OPT_CODE_MAX] = {
    {   TEST_INVALID,       OPT_CODE_INVALID,               OPT_TYPE_NONE                                                           },

    {   TEST_COMMON,        COMMON_CONN_MAX,                OPT_TYPE_UINT,  { .uint = { &common.max_connections,          1     }},
            "conn-max",             "number of connections to use"                                                                  },

    {   TEST_COMMON,        COMMON_PIPELINE,                OPT_TYPE_BOOL,  { .bool = { &common.pipeline_requests,        1     }},  
            "pipeline",             "pipeline requests"                                                                             },

    {   TEST_BENCH_FETCH,   BENCH_FETCH_REQ_CONCURRENCY,    OPT_TYPE_UINT,  { .uint = { &benchmark_fetch.concurrent_ops,  1     }},
            "req-concurrency",      "number of requests to have running"                                                            },

    {   TEST_BENCH_FETCH,   BENCH_FETCH_REQ_AMOUNT,         OPT_TYPE_UINT,  { .uint = { &benchmark_fetch.total_ops,       500   }},
            "req-amount",           "number of requests to issue"                                                                   },

    {   TEST_BENCH_FETCH,   BENCH_FETCH_KEY_PREFIX,         OPT_TYPE_STR,   { .str  = { &benchmark_fetch.key_prefix,      "bf_" }},
            "key-prefix",           "key prefix"                                                                                    },

    {   TEST_BENCH_FETCH,   BENCH_FETCH_KEY_LEN_MIN,        OPT_TYPE_UINT,  { .uint = { &benchmark_fetch.key_len_min,     8     }},  
            "key-len-min",          "minimum key length"                                                                            },

    {   TEST_BENCH_FETCH,   BENCH_FETCH_KEY_LEN_MAX,        OPT_TYPE_UINT,  { .uint = { &benchmark_fetch.key_len_max,     8     }},
            "key-len-max",          "maximum key length"                                                                            },

    {   TEST_BENCH_FETCH,   BENCH_FETCH_KEY_COUNT,          OPT_TYPE_UINT,  { .uint = { &benchmark_fetch.key_count,       1     }},
            "key-count",            "how many keys to use"                                                                          },

    {   TEST_BENCH_FETCH,   BENCH_FETCH_DATA_LEN_MIN,       OPT_TYPE_UINT,  { .uint = { &benchmark_fetch.data_len_min,    64    }},
            "data-len-min",         "minimum data length"                                                                           },

    {   TEST_BENCH_FETCH,   BENCH_FETCH_DATA_LEN_MAX,       OPT_TYPE_UINT,  { .uint = { &benchmark_fetch.data_len_max,    64    }},
            "data-len-max",         "maximum data length"                                                                            },
};

void time_reset () {
    // start timing
    assert(gettimeofday(&common.start, NULL) == 0);
}

double time_offset () {
    struct timeval time;

    assert(gettimeofday(&time, NULL) == 0);
    
    return ((double) (time.tv_sec - common.start.tv_sec)) + ((double) (time.tv_usec - common.start.tv_usec)) / 1000000;
}

int ratelimit (int count, int total) {
    return ((total / 10) ? (count % (total / 10) == 0) : 1);
}



void mc_init (int max_connections, memcache_cb cb_fn) {
    // memcache init    
    if ((mc = memcache_alloc(cb_fn, common.pipeline_requests)) == NULL)
        ERROR("memcache_alloc");
    
    // fix up the endpoint
    endpoint_init(&server_endpoint, 11211);
    
    if (endpoint_parse(&server_endpoint, "localhost"))
        ERROR("config_endpoint_parse");
    
    // add the server
    if (memcache_add_server(mc, &server_endpoint, common.max_connections))
        ERROR("memcache_add_server");
    
    INFO("[memcache] initialized with pipeline_requests=%c max_connections=%d", common.pipeline_requests ? 't' : 'f', common.max_connections);

    return;

error:
    assert(0);
}

void dump_req (struct memcache_req *req, void *unused) {
    const struct memcache_obj *obj;
    const struct memcache_buf *buf;

    INFO("[%.*s]: cmd=%s state=%s reply=%s",
        (int) memcache_req_key(req)->len, memcache_req_key(req)->buf,
        memcache_command_str(memcache_req_cmd(req)),
        memcache_state_str(memcache_req_state(req)),
        memcache_reply_str(memcache_req_reply(req))
    );
    
    if ((obj = memcache_req_obj(req)))
        INFO("\tobj: flags=0x%04X exptime=%zu bytes=%zu cas=%llu", obj->flags, obj->exptime, obj->bytes, obj->cas);

    if ((buf = memcache_req_buf(req)))
        INFO("\tbuf: data=%p len=%zu offset=%zu", buf->data, buf->len, buf->offset);
    
    INFO("%s", "");
}

size_t random_value (size_t min, size_t max) {
    return ((max == min) ? min : (random() % (max - min)) + min);
}

size_t random_data (char *buf, size_t min, size_t max) {
#define CHAR_TABLE_MAX (('z' - 'a' + 1) + ('Z' - 'A' + 1) + ('9' - '0' + 1))

    static char char_table[CHAR_TABLE_MAX] = {
        'a', 'b', 'c', 'd', 'e', 'f', 'g', 'h', 'i', 'j', 'k', 'l', 'm', 'n', 'o', 'p', 'q', 'r', 's', 't', 'u', 'v', 'w', 'x', 'y', 'z',
        'A', 'B', 'C', 'D', 'E', 'F', 'G', 'H', 'I', 'J', 'K', 'L', 'M', 'N', 'O', 'P', 'Q', 'R', 'S', 'T', 'U', 'V', 'W', 'X', 'Y', 'Z',
        '0', '1', '2', '3', '4', '5', '6', '7', '8', '9'
    };

    assert(max >= min);

    size_t size = random_value(min, max), i;

    for (i = 0; i < size; i++) {
        buf[i] = char_table[random_value(0, CHAR_TABLE_MAX)];

        assert(isalnum(buf[i]));
/*    
        switch (MIN((size - i), 4)) {
            case 4: * ((u_int32_t*) &buf[i]) = random(); i += 4; break;
            case 3:
            case 2: * ((u_int16_t*) &buf[i]) = random(); i += 2; break;
            case 1: * ((u_int8_t*)  &buf[i]) = random(); i += 1; break;
            default: assert(0);
        }
*/        
    }

    return size;
}

void test_cb (struct memcache_req *req, void *arg) {
    dump_req(req, arg);
}

void begin_test () {
    struct memcache_key key_1, key_2;
    struct memcache_obj obj_1, obj_2;
    struct memcache_buf buf_1, buf_2;
    struct memcache_req *req_1s, *req_2s, *req_1f, *req_2f;
    
    mc_init(1, &test_cb);

    // add a request or two
    key_1.buf = "memcache_test_k1";
    key_2.buf = "memcache_test_k2";
    key_1.len = key_2.len = 0;

    obj_1.flags = 0x1A;
    obj_2.flags = 0x2B;

    obj_1.exptime = 0;
    obj_2.exptime = 3600;

    obj_1.bytes = strlen(data_1);
    obj_2.bytes = strlen(data_2);

    buf_1.data = data_1;
    buf_1.len = strlen(data_1);
    buf_1.offset = buf_1.len;

    buf_2.data = data_2;
    buf_2.len = strlen(data_2);
    buf_2.offset = buf_2.len;

    if ((req_1s = memcache_store(mc, MEMCACHE_CMD_STORE_SET, &key_1, &obj_1, &buf_1, key_1.buf)) == NULL)
        ERROR("memcache_store: key_1");
    
    if ((req_2s = memcache_store(mc, MEMCACHE_CMD_STORE_ADD, &key_2, &obj_2, &buf_2, key_2.buf)) == NULL)
        ERROR("memcache_store: key_2");
    
    if ((req_1f = memcache_fetch(mc, &key_1, key_1.buf)) == NULL)
        ERROR("memcache_fetch: key_1");
    
    if ((req_2f = memcache_fetch(mc, &key_2, key_2.buf)) == NULL)
        ERROR("memcache_fetch: key_2");

error:
    return;
}

void benchmark_continue () {
    while (benchmark_fetch.cur_ops < benchmark_fetch.concurrent_ops && (benchmark_fetch.op_count + benchmark_fetch.cur_ops) < benchmark_fetch.total_ops) {
        // launch
        assert(memcache_fetch(mc, &benchmark_fetch.keys[random_value(0, benchmark_fetch.key_count)].key, NULL) != NULL);

        benchmark_fetch.cur_ops++;
        
        if (ratelimit(benchmark_fetch.op_count + benchmark_fetch.cur_ops, benchmark_fetch.total_ops))
            INFO("[benchmark] %0.6f: %d+%d/%d requests", 
                time_offset(),
                benchmark_fetch.op_count, benchmark_fetch.cur_ops, benchmark_fetch.total_ops
            );
    }

    if (benchmark_fetch.op_count == benchmark_fetch.total_ops) {
        // done
        assert(event_base_loopexit(ev_base, NULL) == 0);

        INFO("[benchmark] %.6f: %.6f req/s", 
            time_offset(),
            benchmark_fetch.op_count / time_offset()
        );
    }
}

void benchmark_fetch_start () {    
    INFO(
        "[benchmark] %0.6f: starting\n"
        "\tconcurrent_ops   = %u\n"
        "\ttotal_ops        = %u\n"
        "\tkey_prefix       = %s\n"
        "\tkey_len_min      = %u\n"
        "\tkey_len_max      = %u\n"
        "\tkey_count        = %u\n"
        "\tdata_len_min     = %u\n"
        "\tdata_len_max     = %u\n"
        , time_offset(),
        benchmark_fetch.concurrent_ops, benchmark_fetch.total_ops,
        benchmark_fetch.key_prefix, benchmark_fetch.key_len_min, benchmark_fetch.key_len_max, benchmark_fetch.key_count,
        benchmark_fetch.data_len_min, benchmark_fetch.data_len_max
    );
    
    time_reset();

    benchmark_continue();
    
    INFO("[benchmark] %0.6f: running", 
        time_offset()
    );
}


void benchmark_cb (struct memcache_req *req, void *arg) {
    enum memcache_command cmd = memcache_req_cmd(req);
    enum memcache_state state = memcache_req_state(req);

    if (state == MEMCACHE_STATE_ERROR) {
        dump_req(req, arg);
        FATAL("request failed");
    }

    if (state == MEMCACHE_STATE_DONE || state == MEMCACHE_STATE_DONE_DATA) {
        if (cmd == MEMCACHE_CMD_FETCH_GET) {
            benchmark_fetch.cur_ops--;
            benchmark_fetch.op_count++;

            benchmark_continue();

        } else if (cmd == MEMCACHE_CMD_STORE_SET) {
            if (memcache_req_reply(req) != MEMCACHE_RPL_STORED) {
                dump_req(req, arg);
                WARNING("value was not stored");
            }

            benchmark_fetch.keys_stored++;
            
            if (ratelimit(benchmark_fetch.keys_stored, benchmark_fetch.key_count))
                INFO("[benchmark] %.6f: key %u/%u stored: %.*s", 
                    time_offset(), benchmark_fetch.keys_stored, benchmark_fetch.key_count, (int) memcache_req_key(req)->len, memcache_req_key(req)->buf
                );

            if (benchmark_fetch.keys_stored == benchmark_fetch.key_count)
                benchmark_fetch_start();
        }
        
        memcache_req_free(req);
    }
}

/*
 * Run <concurrent_op> ops in parrallel, until we have completed total_ops, at which point we shut down.
 */
void benchmark_fetch_fn () {
    static char data[BENCHMARK_DATA_MAX];
    char key_postfix[BENCHMARK_KEY_MAX];
    int i, key_len, data_len;
    struct memcache_obj obj;
    struct memcache_buf buf;

    assert(benchmark_fetch.key_len_min > 0 && (strlen(benchmark_fetch.key_prefix) + benchmark_fetch.key_len_max) < BENCHMARK_KEY_MAX);
    assert(benchmark_fetch.data_len_min > 0 && benchmark_fetch.data_len_max < BENCHMARK_DATA_MAX);

    benchmark_fetch.cur_ops = benchmark_fetch.op_count = benchmark_fetch.keys_stored = 0;

    if ((benchmark_fetch.keys = calloc(benchmark_fetch.key_count, sizeof(struct key_buf))) == NULL)
        FATAL("calloc");

    // pregenerate the data
    data_len = random_data(data, benchmark_fetch.data_len_min, benchmark_fetch.data_len_max);

    obj.flags = 0x1234;
    obj.exptime = 0;
    obj.bytes = data_len;
    buf.data = data;
    buf.len = buf.offset = data_len;

    // insert keys
    INFO("[benchmark] %0.6f: inserting %u keys with prefix=%s and len=(%u -> %u)",
        time_offset(),
        benchmark_fetch.key_count, benchmark_fetch.key_prefix, benchmark_fetch.key_len_min, benchmark_fetch.key_len_max
    );
    
    for (i = 0; i < benchmark_fetch.key_count; i++) {
        struct key_buf *keybuf = &benchmark_fetch.keys[i];

        key_len = random_data(key_postfix, benchmark_fetch.key_len_min, benchmark_fetch.key_len_max);

        key_postfix[key_len] = '\0';

        assert((keybuf->key.len = snprintf(keybuf->buf, BENCHMARK_KEY_MAX, "%s%*s", benchmark_fetch.key_prefix, key_len, key_postfix)) < BENCHMARK_KEY_MAX);
        
        keybuf->key.buf = keybuf->buf;

        assert(memcache_store(mc, MEMCACHE_CMD_STORE_SET, &keybuf->key, &obj, &buf, NULL) != NULL);
    }
}

void usage_opt (enum option_test test) {
    struct opt *opt;

    for (opt = option_info + 1; opt->code; opt++) {
        if (opt->test == test) {
            switch (opt->type) {
                case OPT_TYPE_BOOL:
                    INFO("\t%-20s %c     %s", opt->name, opt->data.bool.default_value ? 't' : 'f', opt->descr);
                    break;

                case OPT_TYPE_UINT:
                    INFO("\t%-20s %-6d %s", opt->name, opt->data.uint.default_value, opt->descr);
                    break;

                case OPT_TYPE_STR:
                    INFO("\t%-20s %-6s %s", opt->name, opt->data.str.default_value, opt->descr);
                    break;
                
                default:
                    assert(0);
            }
        }
    }
}

void usage (char *cmd) {
    struct test *test;

    INFO("Usage: %s <cmd> [<args> ... ]", cmd);

    INFO("\nCOMMON OPTIONS\n");
    
    usage_opt(TEST_COMMON);

    INFO("\nCOMMANDS\n");

    for (test = test_list; test->name; test++) {
        INFO("%s:", test->name);
        INFO("\t%s", test->descr);
        INFO("\t");
        
        usage_opt(test->code);

        INFO("\t");
    }

    exit(1);
}

int main (int argc, char **argv) {
    char *name, *invalid;
    struct test *test;
    int c, option_index;

    // argument-parsing
    if (argc < 2) {
        WARNING("No command given");
        usage(argv[0]);
    }
    
    // look up the test
    name = argv[1];
    test = test_list;
    
    while (test->name && strcmp(test->name, name) != 0)
        test++;
    
    if (!test->name) {
        WARNING("Unknown cmd '%s'", name);
        usage(argv[0]);
    }
    
    // default values
    for (c = OPT_CODE_INVALID; c < OPT_CODE_MAX; c++) {
        switch (option_info[c].type) {
            case OPT_TYPE_NONE:
                break;
            
            case OPT_TYPE_BOOL:
                *option_info[c].data.bool.value = option_info[c].data.bool.default_value;

                break;

            case OPT_TYPE_UINT:
                *option_info[c].data.uint.value = option_info[c].data.uint.default_value;

                break;

            case OPT_TYPE_STR:
                *option_info[c].data.str.value = option_info[c].data.str.default_value;

                break;

            default:
                assert(0);
        }
    }

    while ((c = getopt_long(argc, argv, "", options, &option_index)) != -1) {
        if (c <= OPT_CODE_INVALID || c >= OPT_CODE_MAX)
            FATAL("invalid argument %s", options[option_index].name);

        if (option_info[c].test != test->code && option_info[c].test != TEST_COMMON)
            FATAL("invalid option %s for test %s", options[option_index].name, test->name);
        
        switch (option_info[c].type) {
            case OPT_TYPE_BOOL:
                assert(optarg);

                switch (optarg[0]) {
                    case 't':
                    case '1':
                        *option_info[c].data.bool.value = 1;
                        break;

                    case 'f':
                    case '0':
                        *option_info[c].data.bool.value = 0;
                        break;

                    default:
                        FATAL("invalid true/false value: %s: %s", options[option_index].name, optarg);
                }

                break;

            case OPT_TYPE_UINT:
                assert(optarg);

                *option_info[c].data.uint.value = strtol(optarg, &invalid, 10);

                if (*invalid)
                    FATAL("invalid argument value: %s: %s (%s)", options[option_index].name, optarg, invalid);
                
                break;

            case OPT_TYPE_STR:
                assert(optarg);

                *option_info[c].data.str.value = optarg;

                break;
    
            case OPT_TYPE_NONE:
            default:
                FATAL("invalid argument type %s", options[option_index].name);
                
                break;
        }
    }

    // libevent init 
    ev_base = event_init();

    if (!ev_base)
        FATAL("event_init");
    
    // set up the memcache
    mc_init(common.max_connections, test->cb_fn);

    // start timing
    time_reset();
    
    // start the test
    test->test_fn();
    
    INFO("[libevent] run");

    // run the libevent mainloop
    if (event_base_dispatch(ev_base))
        WARNING("event_dispatch");

    INFO("[libevent] shutdown");
    
    // clean up
    event_base_free(ev_base);
    
    // successfull exit
    return 0;
}