#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <stdint.h>
#include <hiredis/hiredis.h>

#define TIME_SUB_MS(tv1, tv2)  ((tv1.tv_sec - tv2.tv_sec) * 1000 + (tv1.tv_usec - tv2.tv_usec) / 1000)
// #define PRINT   printf
#define PRINT   

static void die(redisContext *c, const char *msg) {
    fprintf(stderr, "%s: %s\n", msg, c && c->err ? c->errstr : "unknown");
    exit(1);
}

static void must_ok(redisReply *r, const char *what) {
    if (!r) { fprintf(stderr, "%s: reply null\n", what); exit(1); }
    if (!(r->type == REDIS_REPLY_STATUS && r->str && strcasecmp(r->str, "OK") == 0)) {
        fprintf(stderr, "%s: expect +OK, got type=%d str=%s\n",
                what, r->type, r->str ? r->str : "(null)");
        freeReplyObject(r);
        exit(1);
    }
    freeReplyObject(r);
}

static void must_int(redisReply *r, long long expect, const char *what) {
    if (!r) { fprintf(stderr, "%s: reply null\n", what); exit(1); }
    if (r->type != REDIS_REPLY_INTEGER || r->integer != expect) {
        fprintf(stderr, "%s: expect :%lld, got type=%d int=%lld\n",
                what, expect, r->type, (long long)r->integer);
        freeReplyObject(r);
        exit(1);
    }
    freeReplyObject(r);
}

static void must_bulk_eq(redisReply *r, const void *buf, size_t n, const char *what) {
    if (!r) { fprintf(stderr, "%s: reply null\n", what); exit(1); }
    if (r->type != REDIS_REPLY_STRING || r->len != n || memcmp(r->str, buf, n) != 0) {
        fprintf(stderr, "%s: bulk mismatch. type=%d len=%zu\n", what, r->type, r->len);
        fprintf(stderr, "expect:%s, truely:%s\n", (const char*)buf, r->str);
        freeReplyObject(r);
        exit(1);
    }
    freeReplyObject(r);
}

static void must_nil(redisReply *r, const char *what) {
    if (!r) { fprintf(stderr, "%s: reply null\n", what); exit(1); }
    if (r->type != REDIS_REPLY_NIL) {
        fprintf(stderr, "%s: expect nil, got type=%d\n", what, r->type);
        freeReplyObject(r);
        exit(1);
    }
    freeReplyObject(r);
}

void basic_command_test(redisContext *c){
    /* ---------- 1) 基本命令测试 ---------- */
    const char *k = "k1";
    const char *v1 = "v1";
    const char *v2 = "v2";

    must_ok((redisReply*)redisCommand(c, "SET %b %b", k, strlen(k), v1, strlen(v1)), "SET");
    must_bulk_eq((redisReply*)redisCommand(c, "GET %b", k, strlen(k)), v1, strlen(v1), "GET");
    must_int((redisReply*)redisCommand(c, "EXIST %b", k, strlen(k)), 1, "EXIST=1");

    must_ok((redisReply*)redisCommand(c, "MOD %b %b", k, strlen(k), v2, strlen(v2)), "MOD");
    must_bulk_eq((redisReply*)redisCommand(c, "GET %b", k, strlen(k)), v2, strlen(v2), "GET after MOD");

    must_int((redisReply*)redisCommand(c, "DEL %b", k, strlen(k)), 1, "DEL=1");
    must_int((redisReply*)redisCommand(c, "EXIST %b", k, strlen(k)), 0, "EXIST=0");
    must_nil((redisReply*)redisCommand(c, "GET %b", k, strlen(k)), "GET nil");

    printf("[OK] basic SET/GET/MOD/DEL/EXIST\n");
}

void special_char_test(redisContext *c){
    /* ---------- 2) 特殊字符/二进制测试 ---------- */
    uint8_t key2[] = { 'b','i','n',':',0x00,'K','\r','\n',0xFF };
    uint8_t val2[] = { 0x00,'A','\r','\n','B',0x00,0xFF };

    must_ok((redisReply*)redisCommand(c, "SET %b %b", key2, sizeof(key2), val2, sizeof(val2)),
            "SET binary");
    must_bulk_eq((redisReply*)redisCommand(c, "GET %b", key2, sizeof(key2)),
                 val2, sizeof(val2), "GET binary");
    (void)redisCommand(c, "DEL %b", key2, sizeof(key2)); /* 不强制检查返回 */

    printf("[OK] binary/special chars\n");
}

void save(redisContext *c){
    must_ok((redisReply*)redisCommand(c, "SAVE"), "SET binary");

    printf("[OK] SAVE\n");
}

void pipline_set_test(redisContext *c, int start, int countN, const char *op){
    /* ---------- 3) Pipeline 批处理测试 ---------- */
    const int N = countN;

    /* 一次塞 N 个 SET */
    int end = start + N;
    for (int i = start; i < end; i++) {
        char kk[64], vv[64];
        int kn = snprintf(kk, sizeof(kk), "p:%d", i);
        int vn = snprintf(vv, sizeof(vv), "v:%d", i);
        if (redisAppendCommand( c, "%s %b %b",
                                op,
                                kk, (size_t)kn, 
                                vv, (size_t)vn) != REDIS_OK) {
            die(c, "redisAppendCommand SET failed");            
        }

        if(i%10000 == 0) PRINT("SEND: %d\n", i);
    }
    /* 再一次性把 N 个回复读出来 */
    for (int i = start; i < end; i++) {
        redisReply *r = NULL;
        if (redisGetReply(c, (void**)&r) != REDIS_OK || !r) die(c, "redisGetReply SET failed");
        must_ok(r, "pipeline SET reply");

        if(i%10000 == 0) PRINT("RECV: %d\n", i);
    }

    PRINT("[OK] SET pipeline batch %d\n", N);
}

void pipline_get_test(redisContext *c, int start, int countN, const char *op){
    const int N = countN;

    /* pipeline GET + 校验 */
    int end = start + N;
    for (int i = start; i < end; i++) {
        char kk[64];
        int kn = snprintf(kk, sizeof(kk), "p:%d", i);
        if (redisAppendCommand( c, "%s %b", 
                                op, 
                                kk, (size_t)kn) != REDIS_OK) {
            die(c, "redisAppendCommand GET failed");
        }
        
        if(i%10000 == 0) PRINT("SEND: %d\n", i);
    }
    for (int i = start; i < end; i++) {
        redisReply *r = NULL;
        if (redisGetReply(c, (void**)&r) != REDIS_OK || !r) die(c, "redisGetReply GET failed");
        char expect[64];
        int en = snprintf(expect, sizeof(expect), "v:%d", i);
        must_bulk_eq(r, expect, (size_t)en, "pipeline GET reply");
        
        if(i%10000 == 0) PRINT("RECV: %d\n", i);
    }

    PRINT("[OK] GET pipeline batch %d\n", N);
}

void pipline_del_test(redisContext *c, int start, int countN, const char *op){
    const int N = countN;

    /* cleanup：pipeline DEL */
    int end = start + N;
    for (int i = start; i < end; i++) {
        char kk[64];
        int kn = snprintf(kk, sizeof(kk), "p:%d", i);
        if (redisAppendCommand( c, "%s %b",
                                op, 
                                kk, (size_t)kn) != REDIS_OK) {
            die(c, "redisAppendCommand DEL failed");
        }

        if(i%10000 == 0) PRINT("SEND: %d\n", i);
    }
    for (int i = start; i < end; i++) {
        redisReply *r = NULL;
        if (redisGetReply(c, (void**)&r) != REDIS_OK || !r) die(c, "redisGetReply DEL failed");
        freeReplyObject(r); /* DEL 返回 int，这里不强制检查 */

        if(i%10000 == 0) PRINT("RECV: %d\n", i);
    }

    PRINT("[OK] DEL pipeline batch %d\n", N);
}

long long test_nopersist_noreplica(redisContext *c, int rounds, long long batch_size){
    struct timeval tv_begin, tv_end;
    gettimeofday(&tv_begin, NULL);

    long long total_ops  = batch_size*rounds;

    for(int i = 0;i < total_ops ; i += batch_size){
        pipline_set_test(c, i, batch_size, "RSET");
    }

    for(int i = 0;i < total_ops ; i += batch_size){
        pipline_get_test(c, i, batch_size, "RGET");
    }

    for(int i = 0;i < total_ops ; i += batch_size){
        pipline_del_test(c, i, batch_size, "RDEL");
    }

    gettimeofday(&tv_end, NULL);
    int time_used = TIME_SUB_MS(tv_end, tv_begin);
    long long qps = total_ops *3*1000/time_used;
    printf("BATCH (N=%lld) --> time_used=%d ms, qps=%lld\n", total_ops *3, time_used, qps);

    return qps;
}

int main(int argc, char **argv) {
    if(argc < 4) {
        printf("invalid input\n");
        return -1;
    }

    const char *host = argv[1];
    int port = atoi(argv[2]);
    int mode = atoi(argv[3]);

    redisContext *c = redisConnect(host, port);
    if (!c || c->err) die(c, "connect failed");

    printf("Connected to %s:%d\n", host, port);

    if(mode == 0){
        save(c);
    }else if(mode == 1){
        basic_command_test(c);
    }else if(mode == 3){
        int rounds = 10;
        long long batch_size = 100000;
        int testrounds = 50;
        long long total_qps = 0;

        for(int i = 0;i < testrounds; ++ i){
            total_qps += test_nopersist_noreplica(c, rounds, batch_size);
        }
        printf("average qps:%lld\n", total_qps/testrounds);
    }else if(mode == 4){

    }else if(mode == 5){

    }else if(mode == 10){
        pipline_set_test(c, 0, 1000, "SET");
    }else if(mode == 11){
        pipline_get_test(c, 0, 1000, "GET");
    }else if(mode == 12){
        pipline_del_test(c, 0, 1000, "DEL");

    }else if(mode == 20){
        pipline_set_test(c, 0, 1000, "RSET");
    }else if(mode == 21){
        pipline_get_test(c, 0, 1000, "RGET");
    }else if(mode == 22){
        pipline_del_test(c, 0, 1000, "RDEL");
    
    }else if(mode == 30){
        pipline_set_test(c, 0, 1000, "HSET");
    }else if(mode == 31){
        pipline_get_test(c, 0, 1000, "HGET");
    }else if(mode == 32){
        pipline_del_test(c, 0, 1000, "HDEL");
    }

    redisFree(c);
    printf("ALL TESTS PASSED.\n");
    return 0;
}