// SPDX-FileCopyrightText: 2024 Redict Contributors
// SPDX-FileCopyrightText: 2024 Salvatore Sanfilippo <antirez at gmail dot com>
//
// SPDX-License-Identifier: BSD-3-Clause
// SPDX-License-Identifier: LGPL-3.0-only

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

#include "redictmodule.h"
#include <stdlib.h>

static RedictModuleType *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_resumes = 0;
unsigned long int datatype_wrong_cursor = 0;
unsigned long int global_attempts = 0;
unsigned long int global_defragged = 0;

int global_strings_len = 0;
RedictModuleString **global_strings = NULL;

static void createGlobalStrings(RedictModuleCtx *ctx, int count)
{
    global_strings_len = count;
    global_strings = RedictModule_Alloc(sizeof(RedictModuleString *) * count);

    for (int i = 0; i < count; i++) {
        global_strings[i] = RedictModule_CreateStringFromLongLong(ctx, i);
    }
}

static void defragGlobalStrings(RedictModuleDefragCtx *ctx)
{
    for (int i = 0; i < global_strings_len; i++) {
        RedictModuleString *new = RedictModule_DefragRedictModuleString(ctx, global_strings[i]);
        global_attempts++;
        if (new != NULL) {
            global_strings[i] = new;
            global_defragged++;
        }
    }
}

static void FragInfo(RedictModuleInfoCtx *ctx, int for_crash_report) {
    REDICTMODULE_NOT_USED(for_crash_report);

    RedictModule_InfoAddSection(ctx, "stats");
    RedictModule_InfoAddFieldLongLong(ctx, "datatype_attempts", datatype_attempts);
    RedictModule_InfoAddFieldLongLong(ctx, "datatype_defragged", datatype_defragged);
    RedictModule_InfoAddFieldLongLong(ctx, "datatype_resumes", datatype_resumes);
    RedictModule_InfoAddFieldLongLong(ctx, "datatype_wrong_cursor", datatype_wrong_cursor);
    RedictModule_InfoAddFieldLongLong(ctx, "global_attempts", global_attempts);
    RedictModule_InfoAddFieldLongLong(ctx, "global_defragged", global_defragged);
}

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

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

    return o;
}

/* FRAG.RESETSTATS */
static int fragResetStatsCommand(RedictModuleCtx *ctx, RedictModuleString **argv, int argc) {
    REDICTMODULE_NOT_USED(argv);
    REDICTMODULE_NOT_USED(argc);

    datatype_attempts = 0;
    datatype_defragged = 0;
    datatype_resumes = 0;
    datatype_wrong_cursor = 0;
    global_attempts = 0;
    global_defragged = 0;

    RedictModule_ReplyWithSimpleString(ctx, "OK");
    return REDICTMODULE_OK;
}

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

    RedictModuleKey *key = RedictModule_OpenKey(ctx,argv[1],
                                              REDICTMODULE_READ|REDICTMODULE_WRITE);
    int type = RedictModule_KeyType(key);
    if (type != REDICTMODULE_KEYTYPE_EMPTY)
    {
        return RedictModule_ReplyWithError(ctx, "ERR key exists");
    }

    long long len;
    if ((RedictModule_StringToLongLong(argv[2], &len) != REDICTMODULE_OK)) {
        return RedictModule_ReplyWithError(ctx, "ERR invalid len");
    }

    long long size;
    if ((RedictModule_StringToLongLong(argv[3], &size) != REDICTMODULE_OK)) {
        return RedictModule_ReplyWithError(ctx, "ERR invalid size");
    }

    long long maxstep;
    if ((RedictModule_StringToLongLong(argv[4], &maxstep) != REDICTMODULE_OK)) {
        return RedictModule_ReplyWithError(ctx, "ERR invalid maxstep");
    }

    struct FragObject *o = createFragObject(len, size, maxstep);
    RedictModule_ModuleTypeSetValue(key, FragType, o);
    RedictModule_ReplyWithSimpleString(ctx, "OK");
    RedictModule_CloseKey(key);

    return REDICTMODULE_OK;
}

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

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

size_t FragFreeEffort(RedictModuleString *key, const void *value) {
    REDICTMODULE_NOT_USED(key);

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

int FragDefrag(RedictModuleDefragCtx *ctx, RedictModuleString *key, void **value) {
    REDICTMODULE_NOT_USED(key);
    unsigned long i = 0;
    int steps = 0;

    int dbid = RedictModule_GetDbIdFromDefragCtx(ctx);
    RedictModule_Assert(dbid != -1);

    /* Attempt to get cursor, validate it's what we're exepcting */
    if (RedictModule_DefragCursorGet(ctx, &i) == REDICTMODULE_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 = RedictModule_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 = RedictModule_DefragAlloc(ctx, o->values[i]);
        if (new) {
            o->values[i] = new;
            datatype_defragged++;
        }

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

    last_set_cursor = 0;
    return 0;
}

int RedictModule_OnLoad(RedictModuleCtx *ctx, RedictModuleString **argv, int argc) {
    REDICTMODULE_NOT_USED(argv);
    REDICTMODULE_NOT_USED(argc);

    if (RedictModule_Init(ctx, "defragtest", 1, REDICTMODULE_APIVER_1)
        == REDICTMODULE_ERR) return REDICTMODULE_ERR;

    if (RedictModule_GetTypeMethodVersion() < REDICTMODULE_TYPE_METHOD_VERSION) {
        return REDICTMODULE_ERR;
    }

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

    createGlobalStrings(ctx, glen);

    RedictModuleTypeMethods tm = {
            .version = REDICTMODULE_TYPE_METHOD_VERSION,
            .free = FragFree,
            .free_effort = FragFreeEffort,
            .defrag = FragDefrag
    };

    FragType = RedictModule_CreateDataType(ctx, "frag_type", 0, &tm);
    if (FragType == NULL) return REDICTMODULE_ERR;

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

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

    RedictModule_RegisterInfoFunc(ctx, FragInfo);
    RedictModule_RegisterDefragFunc(ctx, defragGlobalStrings);

    return REDICTMODULE_OK;
}