/* A module that implements defrag callback mechanisms.
 */

#include "redismodule.h"
#include <stdlib.h>
#include <string.h>

#define UNUSED(V) ((void) V)

static RedisModuleType *FragType;

struct FragObject {
    unsigned long len;
    void **values;
    int maxstep;
};

/* Make sure we get the expected cursor */
unsigned long int last_set_cursor = 0;

unsigned long int datatype_attempts = 0;
unsigned long int datatype_defragged = 0;
unsigned long int datatype_raw_defragged = 0;
unsigned long int datatype_resumes = 0;
unsigned long int datatype_wrong_cursor = 0;
unsigned long int defrag_started = 0;
unsigned long int defrag_ended = 0;
unsigned long int global_strings_attempts = 0;
unsigned long int global_strings_defragged = 0;
unsigned long int global_dicts_resumes = 0;  /* Number of dict defragmentation resumed from a previous break */
unsigned long int global_subdicts_resumes = 0;  /* Number of subdict defragmentation resumed from a previous break */
unsigned long int global_dicts_attempts = 0; /* Number of attempts to defragment dictionary */
unsigned long int global_dicts_defragged = 0; /* Number of dictionaries successfully defragmented */
unsigned long int global_dicts_items_defragged = 0; /* Number of dictionaries items successfully defragmented */

unsigned long global_strings_len = 0;
RedisModuleString **global_strings = NULL;

unsigned long global_dicts_len = 0;
RedisModuleDict **global_dicts = NULL;

static void createGlobalStrings(RedisModuleCtx *ctx, unsigned long count)
{
    global_strings_len = count;
    global_strings = RedisModule_Alloc(sizeof(RedisModuleString *) * count);

    for (unsigned long i = 0; i < count; i++) {
        global_strings[i] = RedisModule_CreateStringFromLongLong(ctx, i);
    }
}

static int defragGlobalStrings(RedisModuleDefragCtx *ctx)
{
    unsigned long cursor = 0;
    RedisModule_DefragCursorGet(ctx, &cursor);

    RedisModule_Assert(cursor < global_strings_len);
    for (; cursor < global_strings_len; cursor++) {
        RedisModuleString *str = global_strings[cursor];
        if (!str) continue;
        RedisModuleString *new = RedisModule_DefragRedisModuleString(ctx, str);
        global_strings_attempts++;
        if (new != NULL) {
            global_strings[cursor] = new;
            global_strings_defragged++;
        }

        if (RedisModule_DefragShouldStop(ctx)) {
            RedisModule_DefragCursorSet(ctx, cursor);
            return 1;
        }
    }
    return 0;
}

static void createFragGlobalStrings(RedisModuleCtx *ctx) {
    for (unsigned long i = 0; i < global_strings_len; i++) {
        if (i % 2 == 1) {
            RedisModule_FreeString(ctx, global_strings[i]);
            global_strings[i] = NULL;
        }
    }
}

static void createGlobalDicts(RedisModuleCtx *ctx, unsigned long count) {
    global_dicts_len = count;
    global_dicts = RedisModule_Alloc(sizeof(RedisModuleDict *) * count);

    /* Create some nested dictionaries:
     * - Each main dict contains some subdicts.
     * - Each sub-dict contains some strings. */
    for (unsigned long i = 0; i < count; i++) {
        RedisModuleDict *dict = RedisModule_CreateDict(ctx);
        for (unsigned long j = 0; j < 10; j++) {
            /* Create sub dict. */
            RedisModuleDict *subdict = RedisModule_CreateDict(ctx);
            for (unsigned long k = 0; k < 10; k++) {
                RedisModuleString *str = RedisModule_CreateStringFromULongLong(ctx, k);
                RedisModule_DictSet(subdict, str, str);
            }

            RedisModuleString *key = RedisModule_CreateStringFromULongLong(ctx, j);
            RedisModule_DictSet(dict, key, subdict);
            RedisModule_FreeString(ctx, key);
        }
        global_dicts[i] = dict;
    }
}

static void freeFragGlobalSubDict(RedisModuleCtx *ctx, RedisModuleDict *subdict) {
    char *key;
    size_t keylen;
    RedisModuleString *str;
    RedisModuleDictIter *iter = RedisModule_DictIteratorStartC(subdict, "^", NULL, 0);
    while ((key = RedisModule_DictNextC(iter, &keylen, (void**)&str))) {
        RedisModule_FreeString(ctx, str);
    }
    RedisModule_FreeDict(ctx, subdict);
    RedisModule_DictIteratorStop(iter);
}

static void createFragGlobalDicts(RedisModuleCtx *ctx) {
    char *key;
    size_t keylen;
    RedisModuleDict *subdict;

    for (unsigned long i = 0; i < global_dicts_len; i++) {
        RedisModuleDict *dict = global_dicts[i];
        if (!dict) continue;

        /* Handle dictionaries differently based on their index in global_dicts array:
         * 1. For odd indices (i % 2 == 1): Remove the entire dictionary.
         * 2. For even indices: Keep the dictionary but remove half of its items. */
        if (i % 2 == 1) {
            RedisModuleDictIter *iter = RedisModule_DictIteratorStartC(dict, "^", NULL, 0);
            while ((key = RedisModule_DictNextC(iter, &keylen, (void**)&subdict))) {
                freeFragGlobalSubDict(ctx, subdict);
            }
            RedisModule_FreeDict(ctx, dict);
            global_dicts[i] = NULL;
            RedisModule_DictIteratorStop(iter);
        } else {
            int key_index = 0;
            RedisModuleDictIter *iter = RedisModule_DictIteratorStartC(dict, "^", NULL, 0);
            while ((key = RedisModule_DictNextC(iter, &keylen, (void**)&subdict))) {
                if (key_index++ % 2 == 1) {
                    freeFragGlobalSubDict(ctx, subdict);
                    RedisModule_DictReplaceC(dict, key, keylen, NULL);
                }
            }
            RedisModule_DictIteratorStop(iter);
        }
    }
}

static int defragGlobalSubDictValueCB(RedisModuleDefragCtx *ctx, void *data, unsigned char *key, size_t keylen, void **newptr) {
    REDISMODULE_NOT_USED(key);
    REDISMODULE_NOT_USED(keylen);
    if (!data) return 0;
    *newptr = RedisModule_DefragAlloc(ctx, data);
    return 0;
}

static int defragGlobalDictValueCB(RedisModuleDefragCtx *ctx, void *data, unsigned char *key, size_t keylen, void **newptr) {
    REDISMODULE_NOT_USED(key);
    REDISMODULE_NOT_USED(keylen);
    static RedisModuleString *seekTo = NULL;
    RedisModuleDict *subdict = data;
    if (!subdict) return 0;
    if (seekTo != NULL) global_subdicts_resumes++;

    *newptr = RedisModule_DefragRedisModuleDict(ctx, subdict, defragGlobalSubDictValueCB, &seekTo);
    if (*newptr) global_dicts_items_defragged++;
    /* Return 1 if seekTo is not NULL, indicating this node needs more defrag work. */
    return seekTo != NULL;
}

static int defragGlobalDicts(RedisModuleDefragCtx *ctx) {
    static RedisModuleString *seekTo = NULL;
    static unsigned long dict_index = 0;
    unsigned long cursor = 0;

    RedisModule_DefragCursorGet(ctx, &cursor);
    if (cursor == 0) { /* Start a new defrag. */
        if (seekTo) {
            RedisModule_FreeString(NULL, seekTo);
            seekTo = NULL;
        }
        dict_index = 0;
    } else {
        global_dicts_resumes++;
    }

    RedisModule_Assert(dict_index < global_dicts_len);
    for (; dict_index < global_strings_len; dict_index++) {
        RedisModuleDict *dict = global_dicts[dict_index];
        if (!dict) continue;
        RedisModuleDict *new = RedisModule_DefragRedisModuleDict(ctx, dict, defragGlobalDictValueCB, &seekTo);
        global_dicts_attempts++;
        if (new != NULL) {
            global_dicts[dict_index] = new;
            global_dicts_defragged++;
        }

        if (seekTo != NULL) {
            /* Set cursor to 1 to indicate defragmentation is not finished. */
            RedisModule_DefragCursorSet(ctx, 1);
            return 1;
        }
    }

    /* Set cursor to 0 to indicate completion. */
    dict_index = 0;
    RedisModule_DefragCursorSet(ctx, 0);
    return 0;
}

typedef enum { DEFRAG_NOT_START, DEFRAG_STRING, DEFRAG_DICT } defrag_module_stage;
static int defragGlobal(RedisModuleDefragCtx *ctx) {
    static defrag_module_stage stage = DEFRAG_NOT_START;
    if (stage == DEFRAG_NOT_START) {
        stage = DEFRAG_STRING; /* Start a new global defrag. */
    }

    if (stage == DEFRAG_STRING) {
        if (defragGlobalStrings(ctx) != 0) return 1;
        stage = DEFRAG_DICT;
    }
    if (stage == DEFRAG_DICT) {
        if (defragGlobalDicts(ctx) != 0) return 1;
        stage = DEFRAG_NOT_START;
    }
    return 0;
}

static void defragStart(RedisModuleDefragCtx *ctx) {
    REDISMODULE_NOT_USED(ctx);
    defrag_started++;
}

static void defragEnd(RedisModuleDefragCtx *ctx) {
    REDISMODULE_NOT_USED(ctx);
    defrag_ended++;
}

static void FragInfo(RedisModuleInfoCtx *ctx, int for_crash_report) {
    REDISMODULE_NOT_USED(for_crash_report);

    RedisModule_InfoAddSection(ctx, "stats");
    RedisModule_InfoAddFieldLongLong(ctx, "datatype_attempts", datatype_attempts);
    RedisModule_InfoAddFieldLongLong(ctx, "datatype_defragged", datatype_defragged);
    RedisModule_InfoAddFieldLongLong(ctx, "datatype_raw_defragged", datatype_raw_defragged);
    RedisModule_InfoAddFieldLongLong(ctx, "datatype_resumes", datatype_resumes);
    RedisModule_InfoAddFieldLongLong(ctx, "datatype_wrong_cursor", datatype_wrong_cursor);
    RedisModule_InfoAddFieldLongLong(ctx, "global_strings_attempts", global_strings_attempts);
    RedisModule_InfoAddFieldLongLong(ctx, "global_strings_defragged", global_strings_defragged);
    RedisModule_InfoAddFieldLongLong(ctx, "global_dicts_resumes", global_dicts_resumes);
    RedisModule_InfoAddFieldLongLong(ctx, "global_subdicts_resumes", global_subdicts_resumes);
    RedisModule_InfoAddFieldLongLong(ctx, "global_dicts_attempts", global_dicts_attempts);
    RedisModule_InfoAddFieldLongLong(ctx, "global_dicts_defragged", global_dicts_defragged);
    RedisModule_InfoAddFieldLongLong(ctx, "global_dicts_items_defragged", global_dicts_items_defragged);
    RedisModule_InfoAddFieldLongLong(ctx, "defrag_started", defrag_started);
    RedisModule_InfoAddFieldLongLong(ctx, "defrag_ended", defrag_ended);
}

struct FragObject *createFragObject(unsigned long len, unsigned long size, int maxstep) {
    struct FragObject *o = RedisModule_Alloc(sizeof(*o));
    o->len = len;
    o->values = RedisModule_Alloc(sizeof(RedisModuleString*) * len);
    o->maxstep = maxstep;

    for (unsigned long i = 0; i < len; i++) {
        o->values[i] = RedisModule_Calloc(1, size);
    }

    return o;
}

/* FRAG.RESETSTATS */
static int fragResetStatsCommand(RedisModuleCtx *ctx, RedisModuleString **argv, int argc) {
    REDISMODULE_NOT_USED(argv);
    REDISMODULE_NOT_USED(argc);

    datatype_attempts = 0;
    datatype_defragged = 0;
    datatype_raw_defragged = 0;
    datatype_resumes = 0;
    datatype_wrong_cursor = 0;
    global_strings_attempts = 0;
    global_strings_defragged = 0;
    global_dicts_resumes = 0;
    global_subdicts_resumes = 0;
    global_dicts_attempts = 0;
    global_dicts_defragged = 0;
    global_dicts_items_defragged = 0;
    defrag_started = 0;
    defrag_ended = 0;

    RedisModule_ReplyWithSimpleString(ctx, "OK");
    return REDISMODULE_OK;
}

/* FRAG.CREATE key len size maxstep */
static int fragCreateCommand(RedisModuleCtx *ctx, RedisModuleString **argv, int argc) {
    if (argc != 5)
        return RedisModule_WrongArity(ctx);

    RedisModuleKey *key = RedisModule_OpenKey(ctx,argv[1],
                                              REDISMODULE_READ|REDISMODULE_WRITE);
    int type = RedisModule_KeyType(key);
    if (type != REDISMODULE_KEYTYPE_EMPTY)
    {
        return RedisModule_ReplyWithError(ctx, "ERR key exists");
    }

    long long len;
    if ((RedisModule_StringToLongLong(argv[2], &len) != REDISMODULE_OK)) {
        return RedisModule_ReplyWithError(ctx, "ERR invalid len");
    }

    long long size;
    if ((RedisModule_StringToLongLong(argv[3], &size) != REDISMODULE_OK)) {
        return RedisModule_ReplyWithError(ctx, "ERR invalid size");
    }

    long long maxstep;
    if ((RedisModule_StringToLongLong(argv[4], &maxstep) != REDISMODULE_OK)) {
        return RedisModule_ReplyWithError(ctx, "ERR invalid maxstep");
    }

    struct FragObject *o = createFragObject(len, size, maxstep);
    RedisModule_ModuleTypeSetValue(key, FragType, o);
    RedisModule_ReplyWithSimpleString(ctx, "OK");
    RedisModule_CloseKey(key);

    return REDISMODULE_OK;
}

/* FRAG.create_frag_global */
static int fragCreateGlobalCommand(RedisModuleCtx *ctx, RedisModuleString **argv, int argc) {
    UNUSED(argv);
    if (argc != 1)
        return RedisModule_WrongArity(ctx);

    createFragGlobalStrings(ctx);
    createFragGlobalDicts(ctx);
    RedisModule_ReplyWithSimpleString(ctx, "OK");
    return REDISMODULE_OK;
}

void FragFree(void *value) {
    struct FragObject *o = value;

    for (unsigned long i = 0; i < o->len; i++)
        RedisModule_Free(o->values[i]);
    RedisModule_Free(o->values);
    RedisModule_Free(o);
}

size_t FragFreeEffort(RedisModuleString *key, const void *value) {
    REDISMODULE_NOT_USED(key);

    const struct FragObject *o = value;
    return o->len;
}

int FragDefrag(RedisModuleDefragCtx *ctx, RedisModuleString *key, void **value) {
    unsigned long i = 0;
    int steps = 0;

    int dbid = RedisModule_GetDbIdFromDefragCtx(ctx);
    RedisModule_Assert(dbid != -1);

    RedisModule_Log(NULL, "notice", "Defrag key: %s", RedisModule_StringPtrLen(key, NULL));

    /* Attempt to get cursor, validate it's what we're exepcting */
    if (RedisModule_DefragCursorGet(ctx, &i) == REDISMODULE_OK) {
        if (i > 0) datatype_resumes++;

        /* Validate we're expecting this cursor */
        if (i != last_set_cursor) datatype_wrong_cursor++;
    } else {
        if (last_set_cursor != 0) datatype_wrong_cursor++;
    }

    /* Attempt to defrag the object itself */
    datatype_attempts++;
    struct FragObject *o = RedisModule_DefragAlloc(ctx, *value);
    if (o == NULL) {
        /* Not defragged */
        o = *value;
    } else {
        /* Defragged */
        *value = o;
        datatype_defragged++;
    }

    /* Deep defrag now */
    for (; i < o->len; i++) {
        datatype_attempts++;
        void *new = RedisModule_DefragAlloc(ctx, o->values[i]);
        if (new) {
            o->values[i] = new;
            datatype_defragged++;
        }

        if ((o->maxstep && ++steps > o->maxstep) ||
            ((i % 64 == 0) && RedisModule_DefragShouldStop(ctx)))
        {
            RedisModule_DefragCursorSet(ctx, i);
            last_set_cursor = i;
            return 1;
        }
    }

    /* Defrag the values array itself using RedisModule_DefragAllocRaw
     * and RedisModule_DefragFreeRaw for testing purposes. */
    void *new_values = RedisModule_DefragAllocRaw(ctx, o->len * sizeof(void*));
    memcpy(new_values, o->values, o->len * sizeof(void*));
    RedisModule_DefragFreeRaw(ctx, o->values);
    o->values = new_values;
    datatype_raw_defragged++;

    last_set_cursor = 0;
    return 0;
}

int RedisModule_OnLoad(RedisModuleCtx *ctx, RedisModuleString **argv, int argc) {
    REDISMODULE_NOT_USED(argv);
    REDISMODULE_NOT_USED(argc);

    if (RedisModule_Init(ctx, "defragtest", 1, REDISMODULE_APIVER_1)
        == REDISMODULE_ERR) return REDISMODULE_ERR;

    if (RedisModule_GetTypeMethodVersion() < REDISMODULE_TYPE_METHOD_VERSION) {
        return REDISMODULE_ERR;
    }

    long long glen;
    if (argc != 1 || RedisModule_StringToLongLong(argv[0], &glen) == REDISMODULE_ERR) {
        return REDISMODULE_ERR;
    }

    createGlobalStrings(ctx, glen);
    createGlobalDicts(ctx, glen);

    RedisModuleTypeMethods tm = {
            .version = REDISMODULE_TYPE_METHOD_VERSION,
            .free = FragFree,
            .free_effort = FragFreeEffort,
            .defrag = FragDefrag
    };

    FragType = RedisModule_CreateDataType(ctx, "frag_type", 0, &tm);
    if (FragType == NULL) return REDISMODULE_ERR;

    if (RedisModule_CreateCommand(ctx, "frag.create",
                                  fragCreateCommand, "write deny-oom", 1, 1, 1) == REDISMODULE_ERR)
        return REDISMODULE_ERR;

    if (RedisModule_CreateCommand(ctx, "frag.create_frag_global",
        fragCreateGlobalCommand, "write deny-oom", 1, 1, 1) == REDISMODULE_ERR)
    return REDISMODULE_ERR;

    if (RedisModule_CreateCommand(ctx, "frag.resetstats",
                                  fragResetStatsCommand, "write deny-oom", 1, 1, 1) == REDISMODULE_ERR)
        return REDISMODULE_ERR;

    RedisModule_RegisterInfoFunc(ctx, FragInfo);
    RedisModule_RegisterDefragFunc2(ctx, defragGlobal);
    RedisModule_RegisterDefragCallbacks(ctx, defragStart, defragEnd);

    return REDISMODULE_OK;
}