#include <errno.h>
#include <inttypes.h>
#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <strings.h>
#include <time.h>

#include <hiredis/hiredis.h>

typedef enum {
    MODE_SET = 0,
    MODE_GET = 1,
    MODE_MIXED = 2,
} bench_mode_t;

typedef struct {
    const char *host;
    int port;
    bench_mode_t mode;
    uint64_t requests;
    uint32_t pipeline;
    uint32_t keyspace;
    uint32_t value_size;
    uint32_t set_ratio;
    const char *set_cmd;
    const char *get_cmd;
    const char *key_prefix;
    uint64_t seed;
    int verify_get;
} bench_opts_t;

typedef struct {
    uint64_t set_ops;
    uint64_t get_ops;
    uint64_t errors;
    double elapsed_sec;
} bench_result_t;

static void usage(const char *prog) {
    fprintf(stderr,
            "Usage: %s [options]\n"
            "  --host <ip>             default: 127.0.0.1\n"
            "  --port <port>           default: 6379\n"
            "  --mode <set|get|mixed>  default: mixed\n"
            "  --requests <n>          default: 1000000\n"
            "  --pipeline <n>          default: 64\n"
            "  --keyspace <n>          default: 100000\n"
            "  --value-size <n>        default: 32\n"
            "  --set-ratio <0..100>    default: 50 (mixed mode only)\n"
            "  --set-cmd <cmd>         default: SET\n"
            "  --get-cmd <cmd>         default: GET\n"
            "  --key-prefix <prefix>   default: bench:key:\n"
            "  --seed <n>              default: time-based\n"
            "  --verify-get            verify GET value content\n"
            "\nExamples:\n"
            "  # Benchmark Redis\n"
            "  %s --host 127.0.0.1 --port 6379 --mode mixed --requests 2000000\n"
            "\n"
            "  # Benchmark kvstore with Redis-compatible commands\n"
            "  %s --host 127.0.0.1 --port 8888 --mode mixed --requests 2000000\n"
            "\n"
            "  # Benchmark kvstore RBTree path\n"
            "  %s --host 127.0.0.1 --port 8888 --mode mixed --set-cmd RSET --get-cmd RGET\n",
            prog, prog, prog, prog);
}

static void opts_init(bench_opts_t *o) {
    memset(o, 0, sizeof(*o));
    o->host = "127.0.0.1";
    o->port = 6379;
    o->mode = MODE_MIXED;
    o->requests = 1000000;
    o->pipeline = 64;
    o->keyspace = 100000;
    o->value_size = 32;
    o->set_ratio = 50;
    o->set_cmd = "SET";
    o->get_cmd = "GET";
    o->key_prefix = "bench:key:";
    o->seed = (uint64_t)time(NULL);
    o->verify_get = 0;
}

static int parse_u64(const char *s, uint64_t *out) {
    char *end = NULL;
    unsigned long long v;
    errno = 0;
    v = strtoull(s, &end, 10);
    if (errno != 0 || end == s || *end != 0) {
        return -1;
    }
    *out = (uint64_t)v;
    return 0;
}

static int parse_u32(const char *s, uint32_t *out) {
    uint64_t v = 0;
    if (parse_u64(s, &v) != 0 || v > UINT32_MAX) {
        return -1;
    }
    *out = (uint32_t)v;
    return 0;
}

static int parse_args(int argc, char **argv, bench_opts_t *o) {
    int i;
    for (i = 1; i < argc; i++) {
        if (strcmp(argv[i], "--host") == 0 && i + 1 < argc) {
            o->host = argv[++i];
        } else if (strcmp(argv[i], "--port") == 0 && i + 1 < argc) {
            uint32_t p = 0;
            if (parse_u32(argv[++i], &p) != 0 || p == 0 || p > 65535) {
                return -1;
            }
            o->port = (int)p;
        } else if (strcmp(argv[i], "--mode") == 0 && i + 1 < argc) {
            const char *m = argv[++i];
            if (strcmp(m, "set") == 0) {
                o->mode = MODE_SET;
            } else if (strcmp(m, "get") == 0) {
                o->mode = MODE_GET;
            } else if (strcmp(m, "mixed") == 0) {
                o->mode = MODE_MIXED;
            } else {
                return -1;
            }
        } else if (strcmp(argv[i], "--requests") == 0 && i + 1 < argc) {
            if (parse_u64(argv[++i], &o->requests) != 0 || o->requests == 0) {
                return -1;
            }
        } else if (strcmp(argv[i], "--pipeline") == 0 && i + 1 < argc) {
            if (parse_u32(argv[++i], &o->pipeline) != 0 || o->pipeline == 0) {
                return -1;
            }
        } else if (strcmp(argv[i], "--keyspace") == 0 && i + 1 < argc) {
            if (parse_u32(argv[++i], &o->keyspace) != 0 || o->keyspace == 0) {
                return -1;
            }
        } else if (strcmp(argv[i], "--value-size") == 0 && i + 1 < argc) {
            if (parse_u32(argv[++i], &o->value_size) != 0 || o->value_size == 0) {
                return -1;
            }
        } else if (strcmp(argv[i], "--set-ratio") == 0 && i + 1 < argc) {
            if (parse_u32(argv[++i], &o->set_ratio) != 0 || o->set_ratio > 100) {
                return -1;
            }
        } else if (strcmp(argv[i], "--set-cmd") == 0 && i + 1 < argc) {
            o->set_cmd = argv[++i];
        } else if (strcmp(argv[i], "--get-cmd") == 0 && i + 1 < argc) {
            o->get_cmd = argv[++i];
        } else if (strcmp(argv[i], "--key-prefix") == 0 && i + 1 < argc) {
            o->key_prefix = argv[++i];
        } else if (strcmp(argv[i], "--seed") == 0 && i + 1 < argc) {
            if (parse_u64(argv[++i], &o->seed) != 0) {
                return -1;
            }
        } else if (strcmp(argv[i], "--verify-get") == 0) {
            o->verify_get = 1;
        } else if (strcmp(argv[i], "--help") == 0 || strcmp(argv[i], "-h") == 0) {
            return 1;
        } else {
            return -1;
        }
    }
    return 0;
}

static uint64_t mono_ns(void) {
    struct timespec ts;
    clock_gettime(CLOCK_MONOTONIC, &ts);
    return (uint64_t)ts.tv_sec * 1000000000ull + (uint64_t)ts.tv_nsec;
}

static uint64_t xorshift64(uint64_t *state) {
    uint64_t x = *state;
    x ^= x << 13;
    x ^= x >> 7;
    x ^= x << 17;
    *state = x;
    return x;
}

static int append_set(redisContext *c, const char *cmd,
                      const char *key, size_t key_len,
                      const char *val, size_t val_len) {
    const char *argv[3];
    size_t argvlen[3];
    argv[0] = cmd;
    argvlen[0] = strlen(cmd);
    argv[1] = key;
    argvlen[1] = key_len;
    argv[2] = val;
    argvlen[2] = val_len;
    return redisAppendCommandArgv(c, 3, argv, argvlen);
}

static int append_get(redisContext *c, const char *cmd,
                      const char *key, size_t key_len) {
    const char *argv[2];
    size_t argvlen[2];
    argv[0] = cmd;
    argvlen[0] = strlen(cmd);
    argv[1] = key;
    argvlen[1] = key_len;
    return redisAppendCommandArgv(c, 2, argv, argvlen);
}

static int consume_set_reply(redisContext *c) {
    redisReply *r = NULL;
    if (redisGetReply(c, (void **)&r) != REDIS_OK || r == NULL) {
        return -1;
    }
    if (r->type != REDIS_REPLY_STATUS || r->str == NULL || strcasecmp(r->str, "OK") != 0) {
        freeReplyObject(r);
        return -1;
    }
    freeReplyObject(r);
    return 0;
}

static int consume_get_reply(redisContext *c, const char *expect, size_t expect_len, int verify) {
    redisReply *r = NULL;
    if (redisGetReply(c, (void **)&r) != REDIS_OK || r == NULL) {
        return -1;
    }

    if (r->type == REDIS_REPLY_STRING) {
        if (verify && ((size_t)r->len != expect_len || memcmp(r->str, expect, expect_len) != 0)) {
            freeReplyObject(r);
            return -1;
        }
        freeReplyObject(r);
        return 0;
    }

    if (r->type == REDIS_REPLY_NIL) {
        freeReplyObject(r);
        return verify ? -1 : 0;
    }

    freeReplyObject(r);
    return -1;
}

static int prefill(redisContext *c, const bench_opts_t *o, const char *value, size_t value_len) {
    uint64_t i = 0;
    char key[256];

    while (i < o->keyspace) {
        uint32_t batch = o->pipeline;
        if ((uint64_t)batch > (uint64_t)o->keyspace - i) {
            batch = (uint32_t)((uint64_t)o->keyspace - i);
        }

        for (uint32_t j = 0; j < batch; j++) {
            int klen = snprintf(key, sizeof(key), "%s%" PRIu64, o->key_prefix, i + j);
            if (klen <= 0 || (size_t)klen >= sizeof(key)) {
                return -1;
            }
            if (append_set(c, o->set_cmd, key, (size_t)klen, value, value_len) != REDIS_OK) {
                return -1;
            }
        }

        for (uint32_t j = 0; j < batch; j++) {
            if (consume_set_reply(c) != 0) {
                return -1;
            }
        }

        i += batch;
    }

    return 0;
}

/*
 * Keep append/reply operation choices consistent in each batch by building an op mask.
 * This wrapper keeps implementation simple and avoids per-op heap allocation.
 */
static int run_bench_with_mask(redisContext *c, const bench_opts_t *o,
                               const char *value, size_t value_len,
                               bench_result_t *res) {
    uint64_t done = 0;
    uint64_t rng = o->seed ? o->seed : 1;
    char key[256];
    uint8_t *opmask = NULL;
    uint64_t begin_ns;

    memset(res, 0, sizeof(*res));
    opmask = (uint8_t *)malloc(o->pipeline);
    if (!opmask) {
        return -1;
    }

    begin_ns = mono_ns();

    while (done < o->requests) {
        uint32_t batch = o->pipeline;
        if ((uint64_t)batch > o->requests - done) {
            batch = (uint32_t)(o->requests - done);
        }

        for (uint32_t i = 0; i < batch; i++) {
            uint64_t rnd = xorshift64(&rng);
            uint64_t key_id = rnd % o->keyspace;
            int is_set = 0;
            int klen;

            if (o->mode == MODE_SET) {
                is_set = 1;
            } else if (o->mode == MODE_GET) {
                is_set = 0;
            } else {
                is_set = (rnd % 100) < o->set_ratio;
            }
            opmask[i] = (uint8_t)is_set;

            klen = snprintf(key, sizeof(key), "%s%" PRIu64, o->key_prefix, key_id);
            if (klen <= 0 || (size_t)klen >= sizeof(key)) {
                free(opmask);
                return -1;
            }

            if (is_set) {
                if (append_set(c, o->set_cmd, key, (size_t)klen, value, value_len) != REDIS_OK) {
                    free(opmask);
                    return -1;
                }
                res->set_ops++;
            } else {
                if (append_get(c, o->get_cmd, key, (size_t)klen) != REDIS_OK) {
                    free(opmask);
                    return -1;
                }
                res->get_ops++;
            }
        }

        for (uint32_t i = 0; i < batch; i++) {
            int rc = opmask[i] ? consume_set_reply(c)
                               : consume_get_reply(c, value, value_len, o->verify_get);
            if (rc != 0) {
                res->errors++;
                free(opmask);
                return -1;
            }
        }

        done += batch;
    }

    res->elapsed_sec = (double)(mono_ns() - begin_ns) / 1e9;
    free(opmask);
    return 0;
}

int main(int argc, char **argv) {
    bench_opts_t opts;
    bench_result_t result;
    redisContext *ctx;
    struct timeval timeout;
    char *value;
    size_t value_len;
    int parse_rc;

    opts_init(&opts);
    parse_rc = parse_args(argc, argv, &opts);
    if (parse_rc == 1) {
        usage(argv[0]);
        return 0;
    }
    if (parse_rc != 0) {
        usage(argv[0]);
        return 2;
    }

    timeout.tv_sec = 3;
    timeout.tv_usec = 0;
    ctx = redisConnectWithTimeout(opts.host, opts.port, timeout);
    if (!ctx || ctx->err) {
        fprintf(stderr, "connect %s:%d failed: %s\n", opts.host, opts.port,
                ctx ? ctx->errstr : "oom");
        if (ctx) {
            redisFree(ctx);
        }
        return 1;
    }

    value_len = opts.value_size;
    value = (char *)malloc(value_len);
    if (!value) {
        fprintf(stderr, "malloc value buffer failed\n");
        redisFree(ctx);
        return 1;
    }
    for (size_t i = 0; i < value_len; i++) {
        value[i] = (char)('a' + (int)(i % 26));
    }

    if (opts.mode != MODE_SET) {
        fprintf(stdout, "[prefill] keyspace=%u using %s\n", opts.keyspace, opts.set_cmd);
        if (prefill(ctx, &opts, value, value_len) != 0) {
            fprintf(stderr, "prefill failed, err=%s\n", ctx->err ? ctx->errstr : "unknown");
            free(value);
            redisFree(ctx);
            return 1;
        }
    }

    fprintf(stdout,
            "[bench] target=%s:%d mode=%s requests=%" PRIu64
            " pipeline=%u keyspace=%u value_size=%u set_cmd=%s get_cmd=%s\n",
            opts.host, opts.port,
            opts.mode == MODE_SET ? "set" : (opts.mode == MODE_GET ? "get" : "mixed"),
            opts.requests, opts.pipeline, opts.keyspace, opts.value_size,
            opts.set_cmd, opts.get_cmd);

    if (run_bench_with_mask(ctx, &opts, value, value_len, &result) != 0) {
        fprintf(stderr, "benchmark failed, err=%s\n", ctx->err ? ctx->errstr : "reply mismatch");
        free(value);
        redisFree(ctx);
        return 1;
    }

    {
        double qps = result.elapsed_sec > 0 ? (double)(result.set_ops + result.get_ops) / result.elapsed_sec : 0.0;
        double avg_us = (result.set_ops + result.get_ops) > 0
                            ? (result.elapsed_sec * 1e6) / (double)(result.set_ops + result.get_ops)
                            : 0.0;
        fprintf(stdout,
                "[result] elapsed=%.3fs total=%" PRIu64 " set=%" PRIu64 " get=%" PRIu64
                " errors=%" PRIu64 " qps=%.0f avg=%.2fus/op\n",
                result.elapsed_sec,
                result.set_ops + result.get_ops,
                result.set_ops,
                result.get_ops,
                result.errors,
                qps,
                avg_us);
    }

    free(value);
    redisFree(ctx);
    return 0;
}