// Copyright (c) 2020, Meir Shpilraien <meir at redislabs dot com>
// 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

#define _BSD_SOURCE
#define _DEFAULT_SOURCE /* For usleep */

#include "redictmodule.h"
#include <stdio.h>
#include <string.h>
#include <strings.h>
#include <unistd.h>

ustime_t cached_time = 0;

/** stores all the keys on which we got 'loaded' keyspace notification **/
RedictModuleDict *loaded_event_log = NULL;
/** stores all the keys on which we got 'module' keyspace notification **/
RedictModuleDict *module_event_log = NULL;

/** Counts how many deleted KSN we got on keys with a prefix of "count_dels_" **/
static size_t dels = 0;

static int KeySpace_NotificationLoaded(RedictModuleCtx *ctx, int type, const char *event, RedictModuleString *key){
    REDICTMODULE_NOT_USED(ctx);
    REDICTMODULE_NOT_USED(type);

    if(strcmp(event, "loaded") == 0){
        const char* keyName = RedictModule_StringPtrLen(key, NULL);
        int nokey;
        RedictModule_DictGetC(loaded_event_log, (void*)keyName, strlen(keyName), &nokey);
        if(nokey){
            RedictModule_DictSetC(loaded_event_log, (void*)keyName, strlen(keyName), RedictModule_HoldString(ctx, key));
        }
    }

    return REDICTMODULE_OK;
}

static int KeySpace_NotificationGeneric(RedictModuleCtx *ctx, int type, const char *event, RedictModuleString *key) {
    REDICTMODULE_NOT_USED(type);
    const char *key_str = RedictModule_StringPtrLen(key, NULL);
    if (strncmp(key_str, "count_dels_", 11) == 0 && strcmp(event, "del") == 0) {
        if (RedictModule_GetContextFlags(ctx) & REDICTMODULE_CTX_FLAGS_MASTER) {
            dels++;
            RedictModule_Replicate(ctx, "keyspace.incr_dels", "");
        }
        return REDICTMODULE_OK;
    }
    if (cached_time) {
        RedictModule_Assert(cached_time == RedictModule_CachedMicroseconds());
        usleep(1);
        RedictModule_Assert(cached_time != RedictModule_Microseconds());
    }

    if (strcmp(event, "del") == 0) {
        RedictModuleString *copykey = RedictModule_CreateStringPrintf(ctx, "%s_copy", RedictModule_StringPtrLen(key, NULL));
        RedictModuleCallReply* rep = RedictModule_Call(ctx, "DEL", "s!", copykey);
        RedictModule_FreeString(ctx, copykey);
        RedictModule_FreeCallReply(rep);

        int ctx_flags = RedictModule_GetContextFlags(ctx);
        if (ctx_flags & REDICTMODULE_CTX_FLAGS_LUA) {
            RedictModuleCallReply* rep = RedictModule_Call(ctx, "INCR", "c", "lua");
            RedictModule_FreeCallReply(rep);
        }
        if (ctx_flags & REDICTMODULE_CTX_FLAGS_MULTI) {
            RedictModuleCallReply* rep = RedictModule_Call(ctx, "INCR", "c", "multi");
            RedictModule_FreeCallReply(rep);
        }
    }

    return REDICTMODULE_OK;
}

static int KeySpace_NotificationExpired(RedictModuleCtx *ctx, int type, const char *event, RedictModuleString *key) {
    REDICTMODULE_NOT_USED(type);
    REDICTMODULE_NOT_USED(event);
    REDICTMODULE_NOT_USED(key);

    RedictModuleCallReply* rep = RedictModule_Call(ctx, "INCR", "c!", "testkeyspace:expired");
    RedictModule_FreeCallReply(rep);

    return REDICTMODULE_OK;
}

/* This key miss notification handler is performing a write command inside the notification callback.
 * Notice, it is discourage and currently wrong to perform a write command inside key miss event.
 * It can cause read commands to be replicated to the replica/aof. This test is here temporary (for coverage and
 * verification that it's not crashing). */
static int KeySpace_NotificationModuleKeyMiss(RedictModuleCtx *ctx, int type, const char *event, RedictModuleString *key) {
    REDICTMODULE_NOT_USED(type);
    REDICTMODULE_NOT_USED(event);
    REDICTMODULE_NOT_USED(key);

    int flags = RedictModule_GetContextFlags(ctx);
    if (!(flags & REDICTMODULE_CTX_FLAGS_MASTER)) {
        return REDICTMODULE_OK; // ignore the event on replica
    }

    RedictModuleCallReply* rep = RedictModule_Call(ctx, "incr", "!c", "missed");
    RedictModule_FreeCallReply(rep);

    return REDICTMODULE_OK;
}

static int KeySpace_NotificationModuleString(RedictModuleCtx *ctx, int type, const char *event, RedictModuleString *key) {
    REDICTMODULE_NOT_USED(type);
    REDICTMODULE_NOT_USED(event);
    RedictModuleKey *redict_key = RedictModule_OpenKey(ctx, key, REDICTMODULE_READ);

    size_t len = 0;
    /* RedictModule_StringDMA could change the data format and cause the old robj to be freed.
     * This code verifies that such format change will not cause any crashes.*/
    char *data = RedictModule_StringDMA(redict_key, &len, REDICTMODULE_READ);
    int res = strncmp(data, "dummy", 5);
    REDICTMODULE_NOT_USED(res);

    RedictModule_CloseKey(redict_key);

    return REDICTMODULE_OK;
}

static void KeySpace_PostNotificationStringFreePD(void *pd) {
    RedictModule_FreeString(NULL, pd);
}

static void KeySpace_PostNotificationString(RedictModuleCtx *ctx, void *pd) {
    REDICTMODULE_NOT_USED(ctx);
    RedictModuleCallReply* rep = RedictModule_Call(ctx, "incr", "!s", pd);
    RedictModule_FreeCallReply(rep);
}

static int KeySpace_NotificationModuleStringPostNotificationJob(RedictModuleCtx *ctx, int type, const char *event, RedictModuleString *key) {
    REDICTMODULE_NOT_USED(ctx);
    REDICTMODULE_NOT_USED(type);
    REDICTMODULE_NOT_USED(event);

    const char *key_str = RedictModule_StringPtrLen(key, NULL);

    if (strncmp(key_str, "string1_", 8) != 0) {
        return REDICTMODULE_OK;
    }

    RedictModuleString *new_key = RedictModule_CreateStringPrintf(NULL, "string_changed{%s}", key_str);
    RedictModule_AddPostNotificationJob(ctx, KeySpace_PostNotificationString, new_key, KeySpace_PostNotificationStringFreePD);
    return REDICTMODULE_OK;
}

static int KeySpace_NotificationModule(RedictModuleCtx *ctx, int type, const char *event, RedictModuleString *key) {
    REDICTMODULE_NOT_USED(ctx);
    REDICTMODULE_NOT_USED(type);
    REDICTMODULE_NOT_USED(event);

    const char* keyName = RedictModule_StringPtrLen(key, NULL);
    int nokey;
    RedictModule_DictGetC(module_event_log, (void*)keyName, strlen(keyName), &nokey);
    if(nokey){
        RedictModule_DictSetC(module_event_log, (void*)keyName, strlen(keyName), RedictModule_HoldString(ctx, key));
    }
    return REDICTMODULE_OK;
}

static int cmdNotify(RedictModuleCtx *ctx, RedictModuleString **argv, int argc){
    if(argc != 2){
        return RedictModule_WrongArity(ctx);
    }

    RedictModule_NotifyKeyspaceEvent(ctx, REDICTMODULE_NOTIFY_MODULE, "notify", argv[1]);
    RedictModule_ReplyWithNull(ctx);
    return REDICTMODULE_OK;
}

static int cmdIsModuleKeyNotified(RedictModuleCtx *ctx, RedictModuleString **argv, int argc){
    if(argc != 2){
        return RedictModule_WrongArity(ctx);
    }

    const char* key  = RedictModule_StringPtrLen(argv[1], NULL);

    int nokey;
    RedictModuleString* keyStr = RedictModule_DictGetC(module_event_log, (void*)key, strlen(key), &nokey);

    RedictModule_ReplyWithArray(ctx, 2);
    RedictModule_ReplyWithLongLong(ctx, !nokey);
    if(nokey){
        RedictModule_ReplyWithNull(ctx);
    }else{
        RedictModule_ReplyWithString(ctx, keyStr);
    }
    return REDICTMODULE_OK;
}

static int cmdIsKeyLoaded(RedictModuleCtx *ctx, RedictModuleString **argv, int argc){
    if(argc != 2){
        return RedictModule_WrongArity(ctx);
    }

    const char* key  = RedictModule_StringPtrLen(argv[1], NULL);

    int nokey;
    RedictModuleString* keyStr = RedictModule_DictGetC(loaded_event_log, (void*)key, strlen(key), &nokey);

    RedictModule_ReplyWithArray(ctx, 2);
    RedictModule_ReplyWithLongLong(ctx, !nokey);
    if(nokey){
        RedictModule_ReplyWithNull(ctx);
    }else{
        RedictModule_ReplyWithString(ctx, keyStr);
    }
    return REDICTMODULE_OK;
}

static int cmdDelKeyCopy(RedictModuleCtx *ctx, RedictModuleString **argv, int argc) {
    if (argc != 2)
        return RedictModule_WrongArity(ctx);

    cached_time = RedictModule_CachedMicroseconds();

    RedictModuleCallReply* rep = RedictModule_Call(ctx, "DEL", "s!", argv[1]);
    if (!rep) {
        RedictModule_ReplyWithError(ctx, "NULL reply returned");
    } else {
        RedictModule_ReplyWithCallReply(ctx, rep);
        RedictModule_FreeCallReply(rep);
    }
    cached_time = 0;
    return REDICTMODULE_OK;
}

/* Call INCR and propagate using RM_Call with `!`. */
static int cmdIncrCase1(RedictModuleCtx *ctx, RedictModuleString **argv, int argc) {
    if (argc != 2)
        return RedictModule_WrongArity(ctx);

    RedictModuleCallReply* rep = RedictModule_Call(ctx, "INCR", "s!", argv[1]);
    if (!rep) {
        RedictModule_ReplyWithError(ctx, "NULL reply returned");
    } else {
        RedictModule_ReplyWithCallReply(ctx, rep);
        RedictModule_FreeCallReply(rep);
    }
    return REDICTMODULE_OK;
}

/* Call INCR and propagate using RM_Replicate. */
static int cmdIncrCase2(RedictModuleCtx *ctx, RedictModuleString **argv, int argc) {
    if (argc != 2)
        return RedictModule_WrongArity(ctx);

    RedictModuleCallReply* rep = RedictModule_Call(ctx, "INCR", "s", argv[1]);
    if (!rep) {
        RedictModule_ReplyWithError(ctx, "NULL reply returned");
    } else {
        RedictModule_ReplyWithCallReply(ctx, rep);
        RedictModule_FreeCallReply(rep);
    }
    RedictModule_Replicate(ctx, "INCR", "s", argv[1]);
    return REDICTMODULE_OK;
}

/* Call INCR and propagate using RM_ReplicateVerbatim. */
static int cmdIncrCase3(RedictModuleCtx *ctx, RedictModuleString **argv, int argc) {
    if (argc != 2)
        return RedictModule_WrongArity(ctx);

    RedictModuleCallReply* rep = RedictModule_Call(ctx, "INCR", "s", argv[1]);
    if (!rep) {
        RedictModule_ReplyWithError(ctx, "NULL reply returned");
    } else {
        RedictModule_ReplyWithCallReply(ctx, rep);
        RedictModule_FreeCallReply(rep);
    }
    RedictModule_ReplicateVerbatim(ctx);
    return REDICTMODULE_OK;
}

static int cmdIncrDels(RedictModuleCtx *ctx, RedictModuleString **argv, int argc) {
    REDICTMODULE_NOT_USED(argv);
    REDICTMODULE_NOT_USED(argc);
    dels++;
    return RedictModule_ReplyWithSimpleString(ctx, "OK");
}

static int cmdGetDels(RedictModuleCtx *ctx, RedictModuleString **argv, int argc) {
    REDICTMODULE_NOT_USED(argv);
    REDICTMODULE_NOT_USED(argc);
    return RedictModule_ReplyWithLongLong(ctx, dels);
}

/* This function must be present on each Redict module. It is used in order to
 * register the commands into the Redict server. */
int RedictModule_OnLoad(RedictModuleCtx *ctx, RedictModuleString **argv, int argc) {
    if (RedictModule_Init(ctx,"testkeyspace",1,REDICTMODULE_APIVER_1) == REDICTMODULE_ERR){
        return REDICTMODULE_ERR;
    }

    loaded_event_log = RedictModule_CreateDict(ctx);
    module_event_log = RedictModule_CreateDict(ctx);

    int keySpaceAll = RedictModule_GetKeyspaceNotificationFlagsAll();

    if (!(keySpaceAll & REDICTMODULE_NOTIFY_LOADED)) {
        // REDICTMODULE_NOTIFY_LOADED event are not supported we can not start
        return REDICTMODULE_ERR;
    }

    if(RedictModule_SubscribeToKeyspaceEvents(ctx, REDICTMODULE_NOTIFY_LOADED, KeySpace_NotificationLoaded) != REDICTMODULE_OK){
        return REDICTMODULE_ERR;
    }

    if(RedictModule_SubscribeToKeyspaceEvents(ctx, REDICTMODULE_NOTIFY_GENERIC, KeySpace_NotificationGeneric) != REDICTMODULE_OK){
        return REDICTMODULE_ERR;
    }

    if(RedictModule_SubscribeToKeyspaceEvents(ctx, REDICTMODULE_NOTIFY_EXPIRED, KeySpace_NotificationExpired) != REDICTMODULE_OK){
        return REDICTMODULE_ERR;
    }

    if(RedictModule_SubscribeToKeyspaceEvents(ctx, REDICTMODULE_NOTIFY_MODULE, KeySpace_NotificationModule) != REDICTMODULE_OK){
        return REDICTMODULE_ERR;
    }

    if(RedictModule_SubscribeToKeyspaceEvents(ctx, REDICTMODULE_NOTIFY_KEY_MISS, KeySpace_NotificationModuleKeyMiss) != REDICTMODULE_OK){
        return REDICTMODULE_ERR;
    }

    if(RedictModule_SubscribeToKeyspaceEvents(ctx, REDICTMODULE_NOTIFY_STRING, KeySpace_NotificationModuleString) != REDICTMODULE_OK){
        return REDICTMODULE_ERR;
    }

    if(RedictModule_SubscribeToKeyspaceEvents(ctx, REDICTMODULE_NOTIFY_STRING, KeySpace_NotificationModuleStringPostNotificationJob) != REDICTMODULE_OK){
        return REDICTMODULE_ERR;
    }

    if (RedictModule_CreateCommand(ctx,"keyspace.notify", cmdNotify,"",0,0,0) == REDICTMODULE_ERR){
        return REDICTMODULE_ERR;
    }

    if (RedictModule_CreateCommand(ctx,"keyspace.is_module_key_notified", cmdIsModuleKeyNotified,"",0,0,0) == REDICTMODULE_ERR){
        return REDICTMODULE_ERR;
    }

    if (RedictModule_CreateCommand(ctx,"keyspace.is_key_loaded", cmdIsKeyLoaded,"",0,0,0) == REDICTMODULE_ERR){
        return REDICTMODULE_ERR;
    }

    if (RedictModule_CreateCommand(ctx, "keyspace.del_key_copy", cmdDelKeyCopy,
                                  "write", 0, 0, 0) == REDICTMODULE_ERR){
        return REDICTMODULE_ERR;
    }

    if (RedictModule_CreateCommand(ctx, "keyspace.incr_case1", cmdIncrCase1,
                                  "write", 0, 0, 0) == REDICTMODULE_ERR){
        return REDICTMODULE_ERR;
    }

    if (RedictModule_CreateCommand(ctx, "keyspace.incr_case2", cmdIncrCase2,
                                  "write", 0, 0, 0) == REDICTMODULE_ERR){
        return REDICTMODULE_ERR;
    }

    if (RedictModule_CreateCommand(ctx, "keyspace.incr_case3", cmdIncrCase3,
                                  "write", 0, 0, 0) == REDICTMODULE_ERR){
        return REDICTMODULE_ERR;
    }

    if (RedictModule_CreateCommand(ctx, "keyspace.incr_dels", cmdIncrDels,
                                  "write", 0, 0, 0) == REDICTMODULE_ERR){
        return REDICTMODULE_ERR;
    }

    if (RedictModule_CreateCommand(ctx, "keyspace.get_dels", cmdGetDels,
                                  "readonly", 0, 0, 0) == REDICTMODULE_ERR){
        return REDICTMODULE_ERR;
    }

    if (argc == 1) {
        const char *ptr = RedictModule_StringPtrLen(argv[0], NULL);
        if (!strcasecmp(ptr, "noload")) {
            /* This is a hint that we return ERR at the last moment of OnLoad. */
            RedictModule_FreeDict(ctx, loaded_event_log);
            RedictModule_FreeDict(ctx, module_event_log);
            return REDICTMODULE_ERR;
        }
    }

    return REDICTMODULE_OK;
}

int RedictModule_OnUnload(RedictModuleCtx *ctx) {
    RedictModuleDictIter *iter = RedictModule_DictIteratorStartC(loaded_event_log, "^", NULL, 0);
    char* key;
    size_t keyLen;
    RedictModuleString* val;
    while((key = RedictModule_DictNextC(iter, &keyLen, (void**)&val))){
        RedictModule_FreeString(ctx, val);
    }
    RedictModule_FreeDict(ctx, loaded_event_log);
    RedictModule_DictIteratorStop(iter);
    loaded_event_log = NULL;

    iter = RedictModule_DictIteratorStartC(module_event_log, "^", NULL, 0);
    while((key = RedictModule_DictNextC(iter, &keyLen, (void**)&val))){
        RedictModule_FreeString(ctx, val);
    }
    RedictModule_FreeDict(ctx, module_event_log);
    RedictModule_DictIteratorStop(iter);
    module_event_log = NULL;

    return REDICTMODULE_OK;
}
