/*
 * Copyright (c) 2006-Present, Redis Ltd.
 * All rights reserved.
 *
 * Licensed under your choice of (a) the Redis Source Available License 2.0
 * (RSALv2); or (b) the Server Side Public License v1 (SSPLv1); or (c) the
 * GNU Affero General Public License v3 (AGPLv3).
 */

#include <assert.h>

#include "version.h"
#include "rmutil/util.h"

#include "cmd_info/command_info.h"
#include "topk.h"
#include "rm_topk.h"
#include "rm_cms.h"
#include "common.h"
#include <math.h>

#define TOPK_DEFAULT_WIDTH 8
#define TOPK_DEFAULT_DEPTH 7
#define TOPK_DEFAULT_DECAY 0.9

#include "load_io_error.h"
// clang-format off
#define INNER_ERROR(x) \
    do { \
        RedisModule_ReplyWithError(ctx, x); \
        return REDISMODULE_ERR; \
    } while(0)
// clang-format on

RedisModuleType *TopKType;

static int GetTopKKey(RedisModuleCtx *ctx, RedisModuleString *keyName, TopK **topk, int mode) {
    // All using this function should call RedisModule_AutoMemory to prevent memory leak
    RedisModuleKey *key = RedisModule_OpenKey(ctx, keyName, mode);
    if (RedisModule_KeyType(key) == REDISMODULE_KEYTYPE_EMPTY) {
        RedisModule_CloseKey(key);
        INNER_ERROR("TopK: key does not exist");
    } else if (RedisModule_ModuleTypeGetType(key) != TopKType) {
        RedisModule_CloseKey(key);
        INNER_ERROR(REDISMODULE_ERRORMSG_WRONGTYPE);
    }

    *topk = RedisModule_ModuleTypeGetValue(key);
    RedisModule_CloseKey(key);
    return REDISMODULE_OK;
}

static int createTopK(RedisModuleCtx *ctx, RedisModuleString **argv, int argc, TopK **topk) {
    long long tmp_ll;
    uint32_t k, width, depth;
    double decay;
    if ((RedisModule_StringToLongLong(argv[2], &tmp_ll) != REDISMODULE_OK) || tmp_ll > UINT32_MAX ||
        tmp_ll < 1) {
        INNER_ERROR("TopK: invalid k");
    }
    k = (uint32_t)tmp_ll;
    if (argc == 6) {
        if ((RedisModule_StringToLongLong(argv[3], &tmp_ll) != REDISMODULE_OK) ||
            tmp_ll > UINT32_MAX || tmp_ll < 1) {
            INNER_ERROR("TopK: invalid width");
        }
        width = (uint32_t)tmp_ll;
        if ((RedisModule_StringToLongLong(argv[4], &tmp_ll) != REDISMODULE_OK) ||
            tmp_ll > UINT32_MAX || tmp_ll < 1) {
            INNER_ERROR("TopK: invalid depth");
        }
        depth = (uint32_t)tmp_ll;
        if ((RedisModule_StringToDouble(argv[5], &decay) != REDISMODULE_OK) ||
            (decay <= 0 || decay > 1)) {
            INNER_ERROR("TopK: invalid decay value. must be '<= 1' & '> 0'");
        }
    } else {
        width = TOPK_DEFAULT_WIDTH;
        depth = TOPK_DEFAULT_DEPTH;
        decay = TOPK_DEFAULT_DECAY;
    }
    *topk = TopK_Create(k, width, depth, decay);
    if (!(*topk)) {
        INNER_ERROR("ERR Insufficient memory to create topk data structure");
    }
    return REDISMODULE_OK;
}

static int TopK_Create_Cmd(RedisModuleCtx *ctx, RedisModuleString **argv, int argc) {
    if (argc != 3 && argc != 6) {
        return RedisModule_WrongArity(ctx);
    }

    RedisModuleKey *key = RedisModule_OpenKey(ctx, argv[1], REDISMODULE_READ | REDISMODULE_WRITE);
    if (RedisModule_KeyType(key) != REDISMODULE_KEYTYPE_EMPTY) {
        RedisModule_ReplyWithError(ctx, "TopK: key already exists");
        goto final;
    }

    TopK *topk = NULL;
    if (createTopK(ctx, argv, argc, &topk) != REDISMODULE_OK)
        goto final;

    if (RedisModule_ModuleTypeSetValue(key, TopKType, topk) == REDISMODULE_ERR) {
        goto final;
    }

    RedisModule_ReplicateVerbatim(ctx);
    RedisModule_ReplyWithSimpleString(ctx, "OK");
final:
    RedisModule_CloseKey(key);
    return REDISMODULE_OK;
}

static int TopK_Add_Cmd(RedisModuleCtx *ctx, RedisModuleString **argv, int argc) {

    if (argc < 3)
        return RedisModule_WrongArity(ctx);

    TopK *topk;
    if (GetTopKKey(ctx, argv[1], &topk, REDISMODULE_READ | REDISMODULE_WRITE) != REDISMODULE_OK) {
        return REDISMODULE_OK;
    }

    int itemCount = argc - 2;
    RedisModule_ReplyWithArray(ctx, itemCount);

    for (int i = 0; i < itemCount; ++i) {
        size_t itemlen;
        const char *item = RedisModule_StringPtrLen(argv[i + 2], &itemlen);
        char *expelledItem = TopK_Add(topk, item, itemlen, 1);

        if (expelledItem == NULL) {
            RedisModule_ReplyWithNull(ctx);
        } else {
            RedisModule_ReplyWithCString(ctx, expelledItem);
            TOPK_FREE(expelledItem);
        }
    }
    RedisModule_ReplicateVerbatim(ctx);
    return REDISMODULE_OK;
}

static int TopK_Incrby_Cmd(RedisModuleCtx *ctx, RedisModuleString **argv, int argc) {

    if (argc < 4 || (argc % 2) == 1)
        return RedisModule_WrongArity(ctx);

    TopK *topk;
    if (GetTopKKey(ctx, argv[1], &topk, REDISMODULE_READ | REDISMODULE_WRITE) != REDISMODULE_OK) {
        return REDISMODULE_OK;
    }

    int itemCount = (argc - 2) / 2;
    RedisModule_ReplyWithArray(ctx, itemCount);

    for (int i = 0; i < itemCount; ++i) {
        size_t itemlen;
        const char *item = RedisModule_StringPtrLen(argv[2 + i * 2], &itemlen);
        long long increment;
        if (RedisModule_StringToLongLong(argv[2 + i * 2 + 1], &increment) || increment < 0 ||
            increment > 100000) {
            RedisModule_ReplyWithError(
                ctx, "TopK: increment must be an integer greater or equal to 0    \
                        and smaller or equal to 100,000");
            goto final;
        }
        char *expelledItem = TopK_Add(topk, item, itemlen, (uint32_t)increment);

        if (expelledItem == NULL) {
            RedisModule_ReplyWithNull(ctx);
        } else {
            RedisModule_ReplyWithCString(ctx, expelledItem);
            TOPK_FREE(expelledItem);
        }
    }
final:
    RedisModule_ReplicateVerbatim(ctx);
    return REDISMODULE_OK;
}

static int TopK_Query_Cmd(RedisModuleCtx *ctx, RedisModuleString **argv, int argc) {
    if (argc < 3)
        return RedisModule_WrongArity(ctx);

    TopK *topk;
    if (GetTopKKey(ctx, argv[1], &topk, REDISMODULE_READ) != REDISMODULE_OK)
        return REDISMODULE_ERR;

    size_t itemlen;
    long long res;
    RedisModule_ReplyWithArray(ctx, argc - 2);
    for (int i = 2; i < argc; ++i) {
        const char *item = RedisModule_StringPtrLen(argv[i], &itemlen);
        res = TopK_Query(topk, item, itemlen);
        if (_is_resp3(ctx)) {
            RedisModule_ReplyWithBool(ctx, !!res);
        } else {
            RedisModule_ReplyWithLongLong(ctx, res);
        }
    }

    return REDISMODULE_OK;
}

static int TopK_Count_Cmd(RedisModuleCtx *ctx, RedisModuleString **argv, int argc) {
    if (argc < 3)
        return RedisModule_WrongArity(ctx);

    TopK *topk;
    if (GetTopKKey(ctx, argv[1], &topk, REDISMODULE_READ) != REDISMODULE_OK)
        return REDISMODULE_ERR;

    size_t itemlen;
    long long res;
    RedisModule_ReplyWithArray(ctx, argc - 2);
    for (int i = 2; i < argc; ++i) {
        const char *item = RedisModule_StringPtrLen(argv[i], &itemlen);
        res = TopK_Count(topk, item, itemlen);
        RedisModule_ReplyWithLongLong(ctx, res);
    }
    return REDISMODULE_OK;
}

static int TopK_List_Cmd(RedisModuleCtx *ctx, RedisModuleString **argv, int argc) {
    if (argc != 2 && argc != 3)
        return RedisModule_WrongArity(ctx);

    bool withCount = 0;
    if (argc == 3) {
        size_t len;
        const char *countStr = RedisModule_StringPtrLen(argv[2], &len);
        if (strncasecmp(countStr, "WITHCOUNT", len) == 0) {
            withCount = 1;
        } else {
            return RedisModule_ReplyWithError(ctx, "WITHCOUNT keyword expected");
        }
    }

    TopK *topk = NULL;
    if (GetTopKKey(ctx, argv[1], &topk, REDISMODULE_READ) != REDISMODULE_OK) {
        return REDISMODULE_OK;
    }
    HeapBucket *heapList = TopK_List(topk);
    RedisModule_ReplyWithArray(ctx, REDISMODULE_POSTPONED_ARRAY_LEN);
    long arrlen = 0;
    for (int i = 0; i < topk->k; ++i) {
        if (heapList[i].count != 0) {
            RedisModule_ReplyWithCString(ctx, heapList[i].item);
            if (withCount) {
                RedisModule_ReplyWithLongLong(ctx, heapList[i].count);
            }
            arrlen += 1 + withCount;
        } else {
            break;
        }
    }
    RedisModule_ReplySetArrayLength(ctx, arrlen);

    TOPK_FREE(heapList);

    return REDISMODULE_OK;
}

static int TopK_Info_Cmd(RedisModuleCtx *ctx, RedisModuleString **argv, int argc) {

    if (argc != 2)
        return RedisModule_WrongArity(ctx);

    TopK *topk = NULL;
    if (GetTopKKey(ctx, argv[1], &topk, REDISMODULE_READ) != REDISMODULE_OK) {
        return REDISMODULE_OK;
    }

    RedisModule_ReplyWithMapOrArray(ctx, 4 * 2, true);
    RedisModule_ReplyWithSimpleString(ctx, "k");
    RedisModule_ReplyWithLongLong(ctx, topk->k);
    RedisModule_ReplyWithSimpleString(ctx, "width");
    RedisModule_ReplyWithLongLong(ctx, topk->width);
    RedisModule_ReplyWithSimpleString(ctx, "depth");
    RedisModule_ReplyWithLongLong(ctx, topk->depth);
    RedisModule_ReplyWithSimpleString(ctx, "decay");
    RedisModule_ReplyWithDouble(ctx, topk->decay);

    return REDISMODULE_OK;
}

/**************** Module functions *********************************/

static void TopKRdbSave(RedisModuleIO *io, void *obj) {
    TopK *topk = obj;
    RedisModule_SaveUnsigned(io, topk->k);
    RedisModule_SaveUnsigned(io, topk->width);
    RedisModule_SaveUnsigned(io, topk->depth);
    RedisModule_SaveDouble(io, topk->decay);
    RedisModule_SaveStringBuffer(io, (const char *)topk->data,
                                 ((size_t)topk->width) * topk->depth * sizeof(Bucket));
    RedisModule_SaveStringBuffer(io, (const char *)topk->heap, topk->k * sizeof(HeapBucket));
    for (uint32_t i = 0; i < topk->k; ++i) {
        if (topk->heap[i].item != NULL) {
            RedisModule_SaveStringBuffer(io, topk->heap[i].item, strlen(topk->heap[i].item) + 1);
        } else {
            RedisModule_SaveStringBuffer(io, "", 1);
        }
    }
}

static void *TopKRdbLoad(RedisModuleIO *io, int encver) {
    if (encver > TOPK_ENC_VER) {
        return NULL;
    }
    TopK *topk = TOPK_CALLOC(1, sizeof *topk);
    bool err = false;
    errdefer(err, TopK_Destroy(topk));
    topk->k = LoadUnsigned_IOError(io, err, NULL);
    topk->width = LoadUnsigned_IOError(io, err, NULL);
    topk->depth = LoadUnsigned_IOError(io, err, NULL);
    topk->decay = LoadDouble_IOError(io, err, NULL);

    size_t size;
    topk->data = (Bucket *)LoadStringBuffer_IOError(io, &size, err, NULL);
    assert(size == sizeof *topk->data * topk->width * topk->depth);
    topk->heap = (HeapBucket *)LoadStringBuffer_IOError(io, &size, err, NULL);
    assert(size == sizeof *topk->heap * topk->k);

    for (HeapBucket *bucket = topk->heap; bucket < topk->heap + topk->k; ++bucket) {
        char *it = LoadStringBuffer_IOError(io, &size, err, NULL);
        bucket->item = size == 1 ? RedisModule_Free(it), NULL : it;
    }

    /* Initialize lookupTable */
    for (uint32_t i = 0; i < TOPK_DECAY_LOOKUP_TABLE; ++i) {
        topk->lookupTable[i] = i == 0 ? 1 : topk->lookupTable[i - 1] * topk->decay;
    }

    return topk;
}

static void TopKFree(void *value) { TopK_Destroy(value); }

static int TopKDefrag(RedisModuleDefragCtx *ctx, RedisModuleString *key, void **value) {
    *value = defragPtr(ctx, *value);
    TopK *topk = *value;
    topk->data = defragPtr(ctx, topk->data);
    topk->heap = defragPtr(ctx, topk->heap);
    for (uint32_t i = 0; i < topk->k; ++i) {
        if (topk->heap[i].item)
            topk->heap[i].item = defragPtr(ctx, topk->heap[i].item);
    }
}

static size_t TopKMemUsage(const void *value) {
    const TopK *topk = value;
    size_t size = sizeof *topk;
    size += sizeof *topk->data * topk->width * topk->depth;
    size += sizeof *topk->heap * topk->k;
    return size;
}

int TopKModule_onLoad(RedisModuleCtx *ctx, RedisModuleString **argv, int argc) {
    // TODO: add option to set defaults from command line and in program
    RedisModuleTypeMethods tm = {
        .version = REDISMODULE_TYPE_METHOD_VERSION,
        .rdb_load = TopKRdbLoad,
        .rdb_save = TopKRdbSave,
        .aof_rewrite = RMUtil_DefaultAofRewrite,
        .mem_usage = TopKMemUsage,
        .free = TopKFree,
        .defrag = TopKDefrag,
    };

    TopKType = RedisModule_CreateDataType(ctx, "TopK-TYPE", TOPK_ENC_VER, &tm);
    if (TopKType == NULL) {
        return REDISMODULE_ERR;
    }

#define RegisterCommand(ctx, name, cmd, mode, acl)                                                 \
    RegisterCommandWithModesAndAcls(ctx, name, cmd, mode, acl " topk")

    RegisterAclCategory(ctx, "topk");
    RegisterCommand(ctx, "topk.reserve", TopK_Create_Cmd, "write deny-oom", "write fast");
    RegisterCommand(ctx, "topk.add", TopK_Add_Cmd, "write deny-oom", "write");
    RegisterCommand(ctx, "topk.incrby", TopK_Incrby_Cmd, "write deny-oom", "write");
    RegisterCommand(ctx, "topk.query", TopK_Query_Cmd, "readonly", "read");
    RegisterCommand(ctx, "topk.count", TopK_Count_Cmd, "readonly", "read");
    RegisterCommand(ctx, "topk.list", TopK_List_Cmd, "readonly", "read");
    RegisterCommand(ctx, "topk.info", TopK_Info_Cmd, "readonly", "read fast");

#undef RegisterCommand
    if (RegisterTopKCommandInfos(ctx) != REDISMODULE_OK)
        return REDISMODULE_ERR;

    return REDISMODULE_OK;
}
