#define NOTA_IMPLEMENTATION
#include "pit_internal.h"
#include "cell.h"

static int nota_get_arr_len (JSContext *ctx, JSValue arr) {
  int64_t len;
  JS_GetLength (ctx, arr, &len);
  return (int)len;
}

typedef struct NotaVisitedNode {
  JSGCRef ref;
  struct NotaVisitedNode *next;
} NotaVisitedNode;

typedef struct NotaEncodeContext {
  JSContext *ctx;
  NotaVisitedNode *visited_list;
  NotaBuffer nb;
  int cycle;
  JSGCRef *replacer_ref; /* pointer to GC-rooted ref */
} NotaEncodeContext;

static void nota_stack_push (NotaEncodeContext *enc, JSValueConst val) {
  NotaVisitedNode *node = (NotaVisitedNode *)sys_malloc (sizeof (NotaVisitedNode));
  JS_PushGCRef (enc->ctx, &node->ref);
  node->ref.val = val;
  node->next = enc->visited_list;
  enc->visited_list = node;
}

static void nota_stack_pop (NotaEncodeContext *enc) {
  NotaVisitedNode *node = enc->visited_list;
  enc->visited_list = node->next;
  JS_PopGCRef (enc->ctx, &node->ref);
  sys_free (node);
}

static int nota_stack_has (NotaEncodeContext *enc, JSValueConst val) {
  NotaVisitedNode *node = enc->visited_list;
  while (node) {
    if (JS_StrictEq (enc->ctx, node->ref.val, val))
      return 1;
    node = node->next;
  }
  return 0;
}

static JSValue nota_apply_replacer (NotaEncodeContext *enc, JSValueConst holder, JSValueConst key, JSValueConst val) {
  if (!enc->replacer_ref || JS_IsNull (enc->replacer_ref->val)) return val;

  JSValue args[2] = { key, val };
  JSValue result = JS_Call (enc->ctx, enc->replacer_ref->val, holder, 2, args);

  if (JS_IsException (result)) return val;
  return result;
}

static 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);
      sys_free (blob);
      break;
    case NOTA_TEXT:
      nota = nota_read_text (&str, nota);
      *tmp = JS_NewString (js, str);
      sys_free (str);
      break;
    case NOTA_ARR:
      nota = nota_read_array (&n, nota);
      *tmp = JS_NewArrayLen (js, n);
      for (int i = 0; i < n; i++) {
        nota = js_do_nota_decode (js, &ret2, nota, *tmp, JS_NewInt32 (js, i), reviver);
        JS_SetPropertyNumber (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++) {
        JSGCRef prop_key_ref, sub_val_ref;
        JS_PushGCRef (js, &prop_key_ref);
        JS_PushGCRef (js, &sub_val_ref);
        nota = nota_read_text (&str, nota);
        prop_key_ref.val = JS_NewString (js, str);
        sub_val_ref.val = JS_NULL;
        nota = js_do_nota_decode (js, &sub_val_ref.val, nota, *tmp, prop_key_ref.val, reviver);
        JS_SetPropertyStr (js, *tmp, str, sub_val_ref.val);
        JS_PopGCRef (js, &sub_val_ref);
        JS_PopGCRef (js, &prop_key_ref);
        sys_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) {
        JSGCRef inner_ref, obj_ref2;
        JS_PushGCRef (js, &inner_ref);
        JS_PushGCRef (js, &obj_ref2);
        inner_ref.val = JS_NULL;
        nota = js_do_nota_decode (js, &inner_ref.val, nota, holder, JS_NULL, reviver);
        obj_ref2.val = JS_NewObject (js);
        if (!JS_IsNull (js->actor_sym))
          JS_SetPropertyKey (js, obj_ref2.val, js->actor_sym, inner_ref.val);
        JS_CellStone (js, obj_ref2.val);
        *tmp = obj_ref2.val;
        JS_PopGCRef (js, &obj_ref2);
        JS_PopGCRef (js, &inner_ref);
      } 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] = { key, *tmp };
    JSValue revived = JS_Call (js, reviver, holder, 2, args);
    if (!JS_IsException (revived)) {
      *tmp = revived;
    } else {
    }
  }

  return nota;
}

static void nota_encode_value (NotaEncodeContext *enc, JSValueConst val, JSValueConst holder, JSValueConst key) {
  JSContext *ctx = enc->ctx;
  JSGCRef replaced_ref, keys_ref, elem_ref, prop_ref;
  JS_PushGCRef (ctx, &replaced_ref);
  replaced_ref.val = nota_apply_replacer (enc, holder, key, val);
  int tag = JS_VALUE_GET_TAG (replaced_ref.val);

  switch (tag) {
    case JS_TAG_INT:
    case JS_TAG_FLOAT64: {
      double d;
      JS_ToFloat64 (ctx, &d, replaced_ref.val);
      nota_write_number (&enc->nb, d);
      break;
    }
    case JS_TAG_STRING: {
      const char *str = JS_ToCString (ctx, replaced_ref.val);
      nota_write_text (&enc->nb, str);
      JS_FreeCString (ctx, str);
      break;
    }
    case JS_TAG_BOOL:
      if (JS_VALUE_GET_BOOL (replaced_ref.val)) 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_PTR: {
      if (JS_IsText (replaced_ref.val)) {
        const char *str = JS_ToCString (ctx, replaced_ref.val);
        nota_write_text (&enc->nb, str);
        JS_FreeCString (ctx, str);
        break;
      }

      if (js_is_blob (ctx, replaced_ref.val)) {
        size_t buf_len;
        void *buf_data = js_get_blob_data (ctx, &buf_len, replaced_ref.val);
        if (buf_data == (void *)-1) {
          JS_PopGCRef (ctx, &replaced_ref);
          return;
        }
        nota_write_blob (&enc->nb, (unsigned long long)buf_len * 8, (const char *)buf_data);
        break;
      }

      if (JS_IsArray (replaced_ref.val)) {
        if (nota_stack_has (enc, replaced_ref.val)) {
          enc->cycle = 1;
          break;
        }
        nota_stack_push (enc, replaced_ref.val);
        int arr_len = nota_get_arr_len (ctx, replaced_ref.val);
        nota_write_array (&enc->nb, arr_len);
        JS_PushGCRef (ctx, &elem_ref);
        for (int i = 0; i < arr_len; i++) {
          elem_ref.val = JS_GetPropertyNumber (ctx, replaced_ref.val, i);
          JSValue elem_key = JS_NewInt32 (ctx, i);
          nota_encode_value (enc, elem_ref.val, replaced_ref.val, elem_key);
        }
        JS_PopGCRef (ctx, &elem_ref);
        nota_stack_pop (enc);
        break;
      }

      JSValue adata = JS_NULL;
      if (!JS_IsNull (ctx->actor_sym)) {
        int has = JS_HasPropertyKey (ctx, replaced_ref.val, ctx->actor_sym);
        if (has > 0) adata = JS_GetPropertyKey (ctx, replaced_ref.val, ctx->actor_sym);
      }
      if (!JS_IsNull (adata)) {
        nota_write_sym (&enc->nb, NOTA_PRIVATE);
        nota_encode_value (enc, adata, replaced_ref.val, JS_NULL);
        break;
      }
      if (nota_stack_has (enc, replaced_ref.val)) {
        enc->cycle = 1;
        break;
      }
      nota_stack_push (enc, replaced_ref.val);

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

      JS_PushGCRef (ctx, &keys_ref);
      keys_ref.val = JS_GetOwnPropertyNames (ctx, replaced_ref.val);
      if (JS_IsException (keys_ref.val)) {
        nota_write_sym (&enc->nb, NOTA_NULL);
        nota_stack_pop (enc);
        JS_PopGCRef (ctx, &keys_ref);
        break;
      }
      int64_t plen64;
      if (JS_GetLength (ctx, keys_ref.val, &plen64) < 0) {
        nota_write_sym (&enc->nb, NOTA_NULL);
        nota_stack_pop (enc);
        JS_PopGCRef (ctx, &keys_ref);
        break;
      }
      uint32_t plen = (uint32_t)plen64;

      JS_PushGCRef (ctx, &prop_ref);
      JS_PushGCRef (ctx, &elem_ref);
      uint32_t non_function_count = 0;
      for (uint32_t i = 0; i < plen; i++) {
        elem_ref.val = JS_GetPropertyNumber (ctx, keys_ref.val, i);
        prop_ref.val = JS_GetProperty (ctx, replaced_ref.val, elem_ref.val);
        if (!JS_IsFunction (prop_ref.val)) non_function_count++;
      }

      nota_write_record (&enc->nb, non_function_count);
      for (uint32_t i = 0; i < plen; i++) {
        elem_ref.val = JS_GetPropertyNumber (ctx, keys_ref.val, i);
        prop_ref.val = JS_GetProperty (ctx, replaced_ref.val, elem_ref.val);
        if (!JS_IsFunction (prop_ref.val)) {
          const char *prop_name = JS_ToCString (ctx, elem_ref.val);
          nota_write_text (&enc->nb, prop_name ? prop_name : "");
          nota_encode_value (enc, prop_ref.val, replaced_ref.val, elem_ref.val);
          JS_FreeCString (ctx, prop_name);
        }
      }
      JS_PopGCRef (ctx, &elem_ref);
      JS_PopGCRef (ctx, &prop_ref);
      JS_PopGCRef (ctx, &keys_ref);
      nota_stack_pop (enc);
      break;
    }
    default:
      nota_write_sym (&enc->nb, NOTA_NULL);
      break;
  }
  JS_PopGCRef (ctx, &replaced_ref);
}

void *value2nota (JSContext *ctx, JSValue v) {
  JSGCRef val_ref, key_ref;
  JS_PushGCRef (ctx, &val_ref);
  JS_PushGCRef (ctx, &key_ref);
  val_ref.val = v;

  NotaEncodeContext enc_s, *enc = &enc_s;
  enc->ctx = ctx;
  enc->visited_list = NULL;
  enc->cycle = 0;
  enc->replacer_ref = NULL;

  nota_buffer_init (&enc->nb, 128);
  key_ref.val = JS_NewString (ctx, "");
  nota_encode_value (enc, val_ref.val, JS_NULL, key_ref.val);

  if (enc->cycle) {
    JS_PopGCRef (ctx, &key_ref);
    JS_PopGCRef (ctx, &val_ref);
    nota_buffer_free (&enc->nb);
    return NULL;
  }

  JS_PopGCRef (ctx, &key_ref);
  JS_PopGCRef (ctx, &val_ref);
  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;
  JSGCRef holder_ref, key_ref, ret_ref;
  JS_PushGCRef (js, &holder_ref);
  JS_PushGCRef (js, &key_ref);
  JS_PushGCRef (js, &ret_ref);
  holder_ref.val = JS_NewObject (js);
  key_ref.val = JS_NewString (js, "");
  ret_ref.val = JS_NULL;
  js_do_nota_decode (js, &ret_ref.val, nota, holder_ref.val, key_ref.val, JS_NULL);
  JSValue result = ret_ref.val;
  JS_PopGCRef (js, &ret_ref);
  JS_PopGCRef (js, &key_ref);
  JS_PopGCRef (js, &holder_ref);
  return result;
}

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

  JSGCRef val_ref, replacer_ref, key_ref;
  JS_PushGCRef (ctx, &val_ref);
  JS_PushGCRef (ctx, &replacer_ref);
  JS_PushGCRef (ctx, &key_ref);
  val_ref.val = argv[0];
  replacer_ref.val = (argc > 1 && JS_IsFunction (argv[1])) ? argv[1] : JS_NULL;

  NotaEncodeContext enc_s, *enc = &enc_s;
  enc->ctx = ctx;
  enc->visited_list = NULL;
  enc->cycle = 0;
  enc->replacer_ref = &replacer_ref;

  nota_buffer_init (&enc->nb, 128);
  key_ref.val = JS_NewString (ctx, "");
  nota_encode_value (enc, val_ref.val, JS_NULL, key_ref.val);

  JSValue ret;
  if (enc->cycle) {
    nota_buffer_free (&enc->nb);
    ret = JS_RaiseDisrupt (ctx, "Tried to encode something to nota with a cycle.");
  } else {
    size_t total_len = enc->nb.size;
    void *data_ptr = enc->nb.data;
    ret = js_new_blob_stoned_copy (ctx, (uint8_t *)data_ptr, total_len);
    nota_buffer_free (&enc->nb);
  }

  JS_PopGCRef (ctx, &key_ref);
  JS_PopGCRef (ctx, &replacer_ref);
  JS_PopGCRef (ctx, &val_ref);
  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 == (unsigned char *)-1) return JS_EXCEPTION;
  if (!nota) return JS_NULL;

  JSGCRef holder_ref, key_ref, ret_ref, reviver_ref;
  JS_PushGCRef (js, &holder_ref);
  JS_PushGCRef (js, &key_ref);
  JS_PushGCRef (js, &ret_ref);
  JS_PushGCRef (js, &reviver_ref);

  reviver_ref.val = (argc > 1 && JS_IsFunction (argv[1])) ? argv[1] : JS_NULL;
  holder_ref.val = JS_NewObject (js);
  key_ref.val = JS_NewString (js, "");
  ret_ref.val = JS_NULL;

  js_do_nota_decode (js, &ret_ref.val, (char *)nota, holder_ref.val, key_ref.val, reviver_ref.val);

  JSValue result = ret_ref.val;
  JS_PopGCRef (js, &reviver_ref);
  JS_PopGCRef (js, &ret_ref);
  JS_PopGCRef (js, &key_ref);
  JS_PopGCRef (js, &holder_ref);
  return result;
}

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

JSValue js_core_internal_nota_use (JSContext *js) {
  JSGCRef export_ref;
  JS_PushGCRef (js, &export_ref);
  export_ref.val = JS_NewObject (js);
  JS_SetPropertyFunctionList (js, export_ref.val, js_nota_funcs, sizeof (js_nota_funcs) / sizeof (JSCFunctionListEntry));
  JSValue result = export_ref.val;
  JS_PopGCRef (js, &export_ref);
  return result;
}