#include "quickjs.h"

#define KIM_IMPLEMENTATION
#define NOTA_IMPLEMENTATION
#include "nota.h"

typedef struct NotaEncodeContext {
    JSContext *ctx;
    JSValue visitedStack;
    NotaBuffer nb;  // use the dynamic NotaBuffer
    int cycle;
} NotaEncodeContext;

static void nota_stack_push(NotaEncodeContext *enc, JSValueConst val)
{
    JSContext *ctx = enc->ctx;
    int len = JS_ArrayLength(ctx, enc->visitedStack);
    JS_SetPropertyInt64(ctx, enc->visitedStack, len, JS_DupValue(ctx, val));
}

static void nota_stack_pop(NotaEncodeContext *enc)
{
    JSContext *ctx = enc->ctx;
    int len = JS_ArrayLength(ctx, enc->visitedStack);
    JS_SetPropertyStr(ctx, enc->visitedStack, "length", JS_NewUint32(ctx, len - 1));
}

static int nota_stack_has(NotaEncodeContext *enc, JSValueConst val)
{
    JSContext *ctx = enc->ctx;

    int len = JS_ArrayLength(ctx, enc->visitedStack);
    for (int i = 0; i < len; i++) {
        JSValue elem = JS_GetPropertyUint32(ctx, enc->visitedStack, i);
        if (JS_IsObject(elem) && JS_IsObject(val)) {
            if (JS_VALUE_GET_OBJ(elem) == JS_VALUE_GET_OBJ(val)) {
                JS_FreeValue(ctx, elem);
                return 1;
            }
        }
        JS_FreeValue(ctx, elem);
    }
    return 0;
}

JSValue number;

char *js_do_nota_decode(JSContext *js, JSValue *tmp, char *nota)
{
    int type = nota_type(nota);
    JSValue ret2;
    long long n;
    double d;
    int b;
    char *str;
    uint8_t *blob;

    switch(type) {
      case NOTA_BLOB:
        nota = nota_read_blob(&n, (char**)&blob, nota);
        *tmp = JS_NewArrayBufferCopy(js, blob, n);
        free(blob);
        break;
      case NOTA_TEXT:
        nota = nota_read_text(&str, nota);
        *tmp = JS_NewString(js, str);
        free(str);
        break;
      case NOTA_ARR:
        nota = nota_read_array(&n, nota);
        *tmp = JS_NewArray(js);
        for (int i = 0; i < n; i++) {
          nota = js_do_nota_decode(js, &ret2, nota);
          JS_SetPropertyInt64(js, *tmp, i, ret2);
        }
        break;
      case NOTA_REC:
        nota = nota_read_record(&n, nota);
        *tmp = JS_NewObject(js);
        for (int i = 0; i < n; i++) {
          nota = nota_read_text(&str, nota);
          nota = js_do_nota_decode(js, &ret2, nota);
          JS_SetPropertyStr(js, *tmp, str, ret2);
          free(str);
        }
        break;
      case NOTA_INT:
        nota = nota_read_int(&n, nota);
        *tmp = JS_NewInt64(js,n);
        break;
      case NOTA_SYM:
        nota = nota_read_sym(&b, nota);
        switch(b) {
          case NOTA_NULL:
            *tmp = JS_UNDEFINED;
            break;
          case NOTA_FALSE:
            *tmp = JS_NewBool(js,0);
            break;
          case NOTA_TRUE:
            *tmp = JS_NewBool(js,1);
            break;
        }
        break;
      default:
      case NOTA_FLOAT:
        nota = nota_read_float(&d, nota);
        *tmp = JS_NewFloat64(js,d);
        break;
    }

    return nota;
}

static void nota_encode_value(NotaEncodeContext *enc, JSValueConst val);

static void encode_object_properties(NotaEncodeContext *enc, JSValueConst val)
{
    JSContext *ctx = enc->ctx;

    JSPropertyEnum *ptab;
    uint32_t plen;
    if (JS_GetOwnPropertyNames(ctx, &ptab, &plen, val, JS_GPN_ENUM_ONLY | JS_GPN_STRING_MASK) < 0) {
        nota_write_sym(&enc->nb, NOTA_NULL);
        return;
    }

    nota_write_record(&enc->nb, plen);

    for (uint32_t i = 0; i < plen; i++) {
        // property name
        const char *propName = JS_AtomToCString(ctx, ptab[i].atom);
        nota_write_text(&enc->nb, propName);
        JS_FreeCString(ctx, propName);

        // property value
        JSValue propVal = JS_GetProperty(ctx, val, ptab[i].atom);
        nota_encode_value(enc, propVal);
        JS_FreeValue(ctx, propVal);

        // free the atom
        JS_FreeAtom(ctx, ptab[i].atom);
    }
    js_free(ctx, ptab);
}

static void nota_encode_value(NotaEncodeContext *enc, JSValueConst val)
{
    JSContext *ctx = enc->ctx;
    int tag = JS_VALUE_GET_TAG(val);

    switch (tag) {
        case JS_TAG_INT: {
          double d;
          JS_ToFloat64(ctx, &d, val);
          nota_write_number(&enc->nb, d);
          return;
        }
        case JS_TAG_BIG_INT:
        case JS_TAG_FLOAT64:
        case JS_TAG_BIG_DECIMAL:
        case JS_TAG_BIG_FLOAT: {
            const char *str = JS_ToCString(ctx, val);
            nota_write_number_str(&enc->nb, str);
            JS_FreeCString(ctx, str);
            return;
        }

        case JS_TAG_STRING: {
            const char *str = JS_ToCString(ctx, val);
            nota_write_text(&enc->nb, str);
            JS_FreeCString(ctx, str);
            return;
        }

        case JS_TAG_BOOL: {
            if (JS_VALUE_GET_BOOL(val))
                nota_write_sym(&enc->nb, NOTA_TRUE);
            else
                nota_write_sym(&enc->nb, NOTA_FALSE);
            return;
        }

        case JS_TAG_NULL:
        case JS_TAG_UNDEFINED:
            nota_write_sym(&enc->nb, NOTA_NULL);
            return;

        case JS_TAG_OBJECT: {
            if (JS_IsArrayBuffer(ctx, val)) {
            size_t bufLen;
            void *bufData = JS_GetArrayBuffer(ctx, &bufLen, val);
                /* Write as a blob of bits (bufLen * 8). */
                nota_write_blob(&enc->nb, (unsigned long long)bufLen * 8, (const char*)bufData);
                return;
            }

            if (JS_IsArray(ctx, val)) {
                if (nota_stack_has(enc, val)) {
                    enc->cycle = 1;
                    return;  // bail out
                }
                nota_stack_push(enc, val);

                int arrLen = JS_ArrayLength(ctx, val);
                nota_write_array(&enc->nb, arrLen);
                for (int i = 0; i < arrLen; i++) {
                    JSValue elemVal = JS_GetPropertyUint32(ctx, val, i);
                    nota_encode_value(enc, elemVal);
                    JS_FreeValue(ctx, elemVal);
                }

                nota_stack_pop(enc);
                return;
            }

            if (nota_stack_has(enc, val)) {
                enc->cycle = 1;
                return;  // bail out
            }

            nota_stack_push(enc, val);
            encode_object_properties(enc, val);
            nota_stack_pop(enc);
            return;
        }

        default:
            nota_write_sym(&enc->nb, NOTA_NULL);
            return;
    }
}

static JSValue js_nota_encode(JSContext *ctx, JSValueConst this_val, int argc, JSValueConst *argv)
{
    if (argc < 1)
        return JS_ThrowTypeError(ctx, "nota.encode requires 1 argument");

    NotaEncodeContext enc_s, *enc = &enc_s;
    enc->ctx = ctx;
    enc->visitedStack = JS_NewArray(ctx);  // empty array initially
    enc->cycle = 0;

    nota_buffer_init(&enc->nb, 128);

    nota_encode_value(enc, argv[0]);

    if (enc->cycle) {
        JS_FreeValue(ctx, enc->visitedStack);
        nota_buffer_free(&enc->nb);
        return JS_ThrowReferenceError(ctx, "Tried to encode something to nota with a cycle.");
    }

    JS_FreeValue(ctx, enc->visitedStack);

    size_t totalLen = enc->nb.size;     // how many bytes used
    void* dataPtr   = enc->nb.data;     // pointer to the raw data
    JSValue ret = JS_NewArrayBufferCopy(ctx, (uint8_t*)dataPtr, totalLen);

    nota_buffer_free(&enc->nb);

    return ret;
}

JSValue js_nota_decode(JSContext *js, JSValue self, int argc, JSValue *argv)
{
    if (argc < 1) return JS_UNDEFINED;

    size_t len;
    unsigned char *nota = JS_GetArrayBuffer(js, &len, argv[0]);
    if (!nota) return JS_UNDEFINED;

    JSValue ret;
    js_do_nota_decode(js, &ret, (char*)nota);
    return ret;
}

static const JSCFunctionListEntry js_nota_funcs[] = {
    JS_CFUNC_DEF("encode", 1, js_nota_encode),
    JS_CFUNC_DEF("decode", 1, js_nota_decode),
};

static int js_nota_init(JSContext *ctx, JSModuleDef *m) {
    JS_SetModuleExportList(ctx, m, js_nota_funcs,
                           sizeof(js_nota_funcs)/sizeof(JSCFunctionListEntry));
    return 0;
}

JSValue js_nota_use(JSContext *js)
{
    JSValue export = JS_NewObject(js);
    JS_SetPropertyFunctionList(js, export,
                               js_nota_funcs,
                               sizeof(js_nota_funcs)/sizeof(JSCFunctionListEntry));
    number = JS_GetPropertyStr(js, JS_GetGlobalObject(js), "Number");
    return export;
}

#ifdef JS_SHARED_LIBRARY
#define JS_INIT_MODULE js_init_module
#else
#define JS_INIT_MODULE js_init_module_nota
#endif

JSModuleDef *JS_INIT_MODULE(JSContext *ctx, const char *module_name) {
    JSModuleDef *m = JS_NewCModule(ctx, module_name, js_nota_init);
    if (!m) return NULL;
    JS_AddModuleExportList(ctx, m, js_nota_funcs,
                           sizeof(js_nota_funcs)/sizeof(JSCFunctionListEntry));
    return m;
}