#include "qjs_nota.h"
#include "qjs_blob.h"

#define NOTA_IMPLEMENTATION
#include "nota.h"

typedef struct NotaEncodeContext {
  JSContext *ctx;
  JSValue visitedStack;
  NotaBuffer nb;
  int cycle;
  JSValue replacer;
} 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_StrictEq(ctx, elem, val)) {
        JS_FreeValue(ctx, elem);
        return 1;
      }
    }
    JS_FreeValue(ctx, elem);
  }
  return 0;
}

static JSValue apply_replacer(NotaEncodeContext *enc, JSValueConst holder, JSValueConst key, JSValueConst val) {
  if (JS_IsNull(enc->replacer)) return JS_DupValue(enc->ctx, val);

  JSValue args[2] = { JS_DupValue(enc->ctx, key), JS_DupValue(enc->ctx, val) };
  JSValue result = JS_Call(enc->ctx, enc->replacer, holder, 2, args);
  JS_FreeValue(enc->ctx, args[0]);
  JS_FreeValue(enc->ctx, args[1]);
  
  if (JS_IsException(result)) return JS_DupValue(enc->ctx, val);
  return result;
}

char *js_do_nota_decode(JSContext *js, JSValue *tmp, char *nota, JSValue holder, JSValue key, JSValue reviver) {
  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_new_blob_stoned_copy(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, *tmp, JS_NewInt32(js, i), reviver);
        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);
        JSValue prop_key = JS_NewString(js, str);
        nota = js_do_nota_decode(js, &ret2, nota, *tmp, prop_key, reviver);
        JS_SetPropertyStr(js, *tmp, str, ret2);
        JS_FreeValue(js, prop_key);
        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);
      if (b == NOTA_PRIVATE) {
        JSValue inner;
        nota = js_do_nota_decode(js, &inner, nota, holder, JS_NULL, reviver);
        JSValue obj = JS_NewObject(js);
        cell_rt *crt = JS_GetContextOpaque(js);
        JS_SetProperty(js, obj, crt->actor_sym, inner);
        *tmp = obj;
      } else {
        switch(b) {
          case NOTA_NULL: *tmp = JS_NULL; break;
          case NOTA_FALSE: *tmp = JS_NewBool(js, 0); break;
          case NOTA_TRUE: *tmp = JS_NewBool(js, 1); break;
          default: *tmp = JS_NULL; break;
        }
      }
      break;
    default:
    case NOTA_FLOAT:
      nota = nota_read_float(&d, nota);
      *tmp = JS_NewFloat64(js, d);
      break;
  }

  if (!JS_IsNull(reviver)) {
    JSValue args[2] = { JS_DupValue(js, key), JS_DupValue(js, *tmp) };
    JSValue revived = JS_Call(js, reviver, holder, 2, args);
    JS_FreeValue(js, args[0]);
    JS_FreeValue(js, args[1]);
    if (!JS_IsException(revived)) {
      JS_FreeValue(js, *tmp);
      *tmp = revived;
    } else {
      JS_FreeValue(js, revived);
    }
  }

  return nota;
}

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

  switch (tag) {
    case JS_TAG_INT: {
      double d;
      JS_ToFloat64(ctx, &d, replaced);
      nota_write_number(&enc->nb, d);
      break;
    }
    case JS_TAG_FLOAT64: {
      const char *str = JS_ToCString(ctx, replaced);
      nota_write_number_str(&enc->nb, str);
      JS_FreeCString(ctx, str);
      break;
    }
    case JS_TAG_STRING: {
      const char *str = JS_ToCString(ctx, replaced);
      nota_write_text(&enc->nb, str);
      JS_FreeCString(ctx, str);
      break;
    }
    case JS_TAG_BOOL:
      if (JS_VALUE_GET_BOOL(replaced)) nota_write_sym(&enc->nb, NOTA_TRUE);
      else nota_write_sym(&enc->nb, NOTA_FALSE);
      break;
    case JS_TAG_NULL:
      nota_write_sym(&enc->nb, NOTA_NULL);
      break;
    case JS_TAG_OBJECT: {
      if (js_is_blob(ctx, replaced)) {
        size_t buf_len;
        void *buf_data = js_get_blob_data(ctx, &buf_len, replaced);
        nota_write_blob(&enc->nb, (unsigned long long)buf_len * 8, (const char*)buf_data);
        break;
      }

      if (JS_IsArray(ctx, replaced)) {
        if (nota_stack_has(enc, replaced)) {
          enc->cycle = 1;
          break;
        }
        nota_stack_push(enc, replaced);
        int arr_len = JS_ArrayLength(ctx, replaced);
        nota_write_array(&enc->nb, arr_len);
        for (int i = 0; i < arr_len; i++) {
          JSValue elem_val = JS_GetPropertyUint32(ctx, replaced, i);
          JSValue elem_key = JS_NewInt32(ctx, i);
          nota_encode_value(enc, elem_val, replaced, elem_key);
          JS_FreeValue(ctx, elem_val);
          JS_FreeValue(ctx, elem_key);
        }
        nota_stack_pop(enc);
        break;
      }

      cell_rt *crt = JS_GetContextOpaque(ctx);
      JSValue adata = JS_GetProperty(ctx, replaced, crt->actor_sym);
      if (!JS_IsNull(adata)) {
        nota_write_sym(&enc->nb, NOTA_PRIVATE);
        nota_encode_value(enc, adata, replaced, JS_NULL);
        JS_FreeValue(ctx, adata);
        break;
      }
      JS_FreeValue(ctx, adata);
      if (nota_stack_has(enc, replaced)) {
        enc->cycle = 1;
        break;
      }
      nota_stack_push(enc, replaced);

      JSValue to_json = JS_GetPropertyStr(ctx, replaced, "toJSON");
      if (JS_IsFunction(ctx, to_json)) {
        JSValue result = JS_Call(ctx, to_json, replaced, 0, NULL);
        JS_FreeValue(ctx, to_json);
        if (!JS_IsException(result)) {
          nota_encode_value(enc, result, holder, key);
          JS_FreeValue(ctx, result);
        } else {
          nota_write_sym(&enc->nb, NOTA_NULL);
        }
        nota_stack_pop(enc);
        break;
      }
      JS_FreeValue(ctx, to_json);

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

      uint32_t non_function_count = 0;
      for (uint32_t i = 0; i < plen; i++) {
        JSValue prop_val = JS_GetProperty(ctx, replaced, ptab[i].atom);
        if (!JS_IsFunction(ctx, prop_val)) non_function_count++;
        JS_FreeValue(ctx, prop_val);
      }

      nota_write_record(&enc->nb, non_function_count);
      for (uint32_t i = 0; i < plen; i++) {
        JSValue prop_val = JS_GetProperty(ctx, replaced, ptab[i].atom);
        if (!JS_IsFunction(ctx, prop_val)) {
          const char *prop_name = JS_AtomToCString(ctx, ptab[i].atom);
          JSValue prop_key = JS_AtomToValue(ctx, ptab[i].atom);
          nota_write_text(&enc->nb, prop_name);
          nota_encode_value(enc, prop_val, replaced, prop_key);
          JS_FreeCString(ctx, prop_name);
          JS_FreeValue(ctx, prop_key);
        }
        JS_FreeValue(ctx, prop_val);
        JS_FreeAtom(ctx, ptab[i].atom);
      }
      js_free(ctx, ptab);
      nota_stack_pop(enc);
      break;
    }
    default:
      nota_write_sym(&enc->nb, NOTA_NULL);
      break;
  }
  JS_FreeValue(ctx, replaced);
}

void *value2nota(JSContext *ctx, JSValue v) {
  NotaEncodeContext enc_s, *enc = &enc_s;
  enc->ctx = ctx;
  enc->visitedStack = JS_NewArray(ctx);
  enc->cycle = 0;
  enc->replacer = JS_NULL;

  nota_buffer_init(&enc->nb, 128);
  nota_encode_value(enc, v, JS_NULL, JS_NewString(ctx, ""));

  if (enc->cycle) {
    JS_FreeValue(ctx, enc->visitedStack);
    nota_buffer_free(&enc->nb);
    return NULL;
  }

  JS_FreeValue(ctx, enc->visitedStack);
  void *data_ptr = enc->nb.data;
  enc->nb.data = NULL;
  nota_buffer_free(&enc->nb);
  return data_ptr;
}

JSValue nota2value(JSContext *js, void *nota) {
  if (!nota) return JS_NULL;
  JSValue ret;
  JSValue holder = JS_NewObject(js);
  js_do_nota_decode(js, &ret, nota, holder, JS_NewString(js, ""), JS_NULL);
  JS_FreeValue(js, holder);
  return ret;
}

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

  NotaEncodeContext enc_s, *enc = &enc_s;
  enc->ctx = ctx;
  enc->visitedStack = JS_NewArray(ctx);
  enc->cycle = 0;
  enc->replacer = (argc > 1 && JS_IsFunction(ctx, argv[1])) ? argv[1] : JS_NULL;

  nota_buffer_init(&enc->nb, 128);
  nota_encode_value(enc, argv[0], JS_NULL, JS_NewString(ctx, ""));

  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 total_len = enc->nb.size;
  void *data_ptr = enc->nb.data;
  JSValue ret = js_new_blob_stoned_copy(ctx, (uint8_t*)data_ptr, total_len);

  nota_buffer_free(&enc->nb);
  return ret;
}

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

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

  JSValue reviver = (argc > 1 && JS_IsFunction(js, argv[1])) ? argv[1] : JS_NULL;
  JSValue ret;
  JSValue holder = JS_NewObject(js);
  js_do_nota_decode(js, &ret, (char*)nota, holder, JS_NewString(js, ""), reviver);
  JS_FreeValue(js, holder);
  return ret;
}

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

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