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

typedef struct ObjectRef {
  void *ptr;
  struct ObjectRef *next;
} ObjectRef;

typedef struct WotaEncodeContext {
  JSContext *ctx;
  ObjectRef *visited_stack;
  WotaBuffer wb;
  int cycle;
  JSValue replacer;
} WotaEncodeContext;

static void wota_stack_push (WotaEncodeContext *enc, JSValueConst val) {
  (void)enc; (void)val;
  /* Cycle detection disabled for performance */
}

static void wota_stack_pop (WotaEncodeContext *enc) {
  if (!enc->visited_stack) return;

  ObjectRef *top = enc->visited_stack;
  enc->visited_stack = top->next;
  sys_free (top);
}

static int wota_stack_has (WotaEncodeContext *enc, JSValueConst val) {
  (void)enc; (void)val;
  return 0;
  /* Cycle detection disabled for performance */
}

static void wota_stack_free (WotaEncodeContext *enc) {
  while (enc->visited_stack) {
    wota_stack_pop (enc);
  }
}

static JSValue wota_apply_replacer (WotaEncodeContext *enc, JSValueConst holder, JSValue key, JSValueConst val) {
  if (JS_IsNull (enc->replacer)) return JS_DupValue (enc->ctx, val);
  JSValue key_val = JS_IsNull (key) ? JS_NULL : JS_DupValue (enc->ctx, key);
  JSValue args[2] = { key_val, 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;
}

static void wota_encode_value (WotaEncodeContext *enc, JSValueConst val, JSValueConst holder, JSValue key);

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

  /* Root the input value to protect it during property enumeration */
  JSGCRef val_ref, keys_ref;
  JS_PushGCRef (ctx, &val_ref);
  JS_PushGCRef (ctx, &keys_ref);
  val_ref.val = JS_DupValue (ctx, val);

  keys_ref.val = JS_GetOwnPropertyNames (ctx, val_ref.val);
  if (JS_IsException (keys_ref.val)) {
    wota_write_sym (&enc->wb, WOTA_NULL);
    JS_FreeValue (ctx, val_ref.val);
    JS_PopGCRef (ctx, &keys_ref);
    JS_PopGCRef (ctx, &val_ref);
    return;
  }
  int64_t plen64;
  if (JS_GetLength (ctx, keys_ref.val, &plen64) < 0) {
    JS_FreeValue (ctx, keys_ref.val);
    JS_FreeValue (ctx, val_ref.val);
    wota_write_sym (&enc->wb, WOTA_NULL);
    JS_PopGCRef (ctx, &keys_ref);
    JS_PopGCRef (ctx, &val_ref);
    return;
  }
  uint32_t plen = (uint32_t)plen64;
  uint32_t non_function_count = 0;

  /* Allocate GC-rooted arrays for props and keys */
  JSGCRef *prop_refs = sys_malloc (sizeof (JSGCRef) * plen);
  JSGCRef *key_refs = sys_malloc (sizeof (JSGCRef) * plen);
  for (uint32_t i = 0; i < plen; i++) {
    JS_PushGCRef (ctx, &prop_refs[i]);
    JS_PushGCRef (ctx, &key_refs[i]);
    prop_refs[i].val = JS_NULL;
    key_refs[i].val = JS_NULL;
  }

  for (uint32_t i = 0; i < plen; i++) {
    /* Store key into its GCRef slot immediately so it's rooted before
       JS_GetProperty can trigger GC and relocate the string. */
    key_refs[i].val = JS_GetPropertyNumber (ctx, keys_ref.val, i);
    JSValue prop_val = JS_GetProperty (ctx, val_ref.val, key_refs[i].val);
    if (!JS_IsFunction (prop_val)) {
      if (i != non_function_count) {
        key_refs[non_function_count].val = key_refs[i].val;
        key_refs[i].val = JS_NULL;
      }
      prop_refs[non_function_count].val = prop_val;
      non_function_count++;
    } else {
      JS_FreeValue (ctx, prop_val);
      JS_FreeValue (ctx, key_refs[i].val);
      key_refs[i].val = JS_NULL;
    }
  }
  JS_FreeValue (ctx, keys_ref.val);
  wota_write_record (&enc->wb, non_function_count);
  for (uint32_t i = 0; i < non_function_count; i++) {
    size_t klen;
    const char *prop_name = JS_ToCStringLen (ctx, &klen, key_refs[i].val);
    wota_write_text_len (&enc->wb, prop_name ? prop_name : "", prop_name ? klen : 0);
    wota_encode_value (enc, prop_refs[i].val, val_ref.val, key_refs[i].val);
    JS_FreeCString (ctx, prop_name);
    JS_FreeValue (ctx, prop_refs[i].val);
    JS_FreeValue (ctx, key_refs[i].val);
  }
  /* Pop all GC refs in reverse order */
  for (int i = plen - 1; i >= 0; i--) {
    JS_PopGCRef (ctx, &key_refs[i]);
    JS_PopGCRef (ctx, &prop_refs[i]);
  }
  sys_free (prop_refs);
  sys_free (key_refs);
  JS_FreeValue (ctx, val_ref.val);
  JS_PopGCRef (ctx, &keys_ref);
  JS_PopGCRef (ctx, &val_ref);
}

static void wota_encode_value (WotaEncodeContext *enc, JSValueConst val, JSValueConst holder, JSValue key) {
  JSContext *ctx = enc->ctx;
  JSValue replaced;
  if (!JS_IsNull (enc->replacer) && !JS_IsNull (key))
    replaced = wota_apply_replacer (enc, holder, key, val);
  else
    replaced = JS_DupValue (enc->ctx, val);

  int tag = JS_VALUE_GET_TAG (replaced);
  switch (tag) {
    case JS_TAG_INT: {
      int32_t d;
      JS_ToInt32 (ctx, &d, replaced);
      wota_write_int_word (&enc->wb, d);
      break;
    }
    case JS_TAG_FLOAT64: {
      double d;
      if (JS_ToFloat64 (ctx, &d, replaced) < 0) {
        wota_write_sym (&enc->wb, WOTA_NULL);
        break;
      }
      wota_write_float_word (&enc->wb, d);
      break;
    }
    case JS_TAG_STRING: {
      size_t plen;
      const char *str = JS_ToCStringLen (ctx, &plen, replaced);
      wota_write_text_len (&enc->wb, str ? str : "", str ? plen : 0);
      JS_FreeCString (ctx, str);
      break;
    }
    case JS_TAG_BOOL:
      wota_write_sym (&enc->wb, JS_VALUE_GET_BOOL (replaced) ? WOTA_TRUE : WOTA_FALSE);
      break;
    case JS_TAG_NULL:
      wota_write_sym (&enc->wb, WOTA_NULL);
      break;
    case JS_TAG_PTR: {
      if (JS_IsText (replaced)) {
        size_t plen;
        const char *str = JS_ToCStringLen (ctx, &plen, replaced);
        wota_write_text_len (&enc->wb, str ? str : "", str ? plen : 0);
        JS_FreeCString (ctx, str);
        break;
      }
      if (js_is_blob (ctx, replaced)) {
        size_t buf_len;
        void *buf_data = js_get_blob_data (ctx, &buf_len, replaced);
        if (buf_data == (void *)-1) {
          JS_FreeValue (ctx, replaced);
          return;
        }
        if (buf_len == 0) {
          wota_write_blob (&enc->wb, 0, "");
        } else {
          wota_write_blob (&enc->wb, (unsigned long long)buf_len * 8, (const char *)buf_data);
        }
        break;
      }
      if (JS_IsArray (replaced)) {
        if (wota_stack_has (enc, replaced)) {
          enc->cycle = 1;
          break;
        }
        wota_stack_push (enc, replaced);
        int64_t arr_len;
        JS_GetLength (ctx, replaced, &arr_len);
        wota_write_array (&enc->wb, arr_len);
        for (int64_t i = 0; i < arr_len; i++) {
          JSValue elem_val = JS_GetPropertyNumber (ctx, replaced, i);
          wota_encode_value (enc, elem_val, replaced, JS_NewInt32 (ctx, (int32_t)i));
          JS_FreeValue (ctx, elem_val);
        }
        wota_stack_pop (enc);
        break;
      }
      JSValue adata = JS_NULL;
      if (!JS_IsNull (ctx->actor_sym)) {
        int has = JS_HasPropertyKey (ctx, replaced, ctx->actor_sym);
        if (has > 0) adata = JS_GetPropertyKey (ctx, replaced, ctx->actor_sym);
      }
      if (!JS_IsNull (adata)) {
        wota_write_sym (&enc->wb, WOTA_PRIVATE);
        wota_encode_value (enc, adata, replaced, JS_NULL);
        JS_FreeValue (ctx, adata);
        break;
      }
      JS_FreeValue (ctx, adata);
      if (wota_stack_has (enc, replaced)) {
        enc->cycle = 1;
        break;
      }
      wota_stack_push (enc, replaced);
      JSValue to_json = JS_GetPropertyStr (ctx, replaced, "toJSON");
      if (JS_IsFunction (to_json)) {
        JSValue result = JS_Call (ctx, to_json, replaced, 0, NULL);
        JS_FreeValue (ctx, to_json);
        if (!JS_IsException (result)) {
          wota_encode_value (enc, result, holder, key);
          JS_FreeValue (ctx, result);
        } else
          wota_write_sym (&enc->wb, WOTA_NULL);
        wota_stack_pop (enc);
        break;
      }
      JS_FreeValue (ctx, to_json);
      encode_object_properties (enc, replaced, holder);
      wota_stack_pop (enc);
      break;
    }
    default:
      wota_write_sym (&enc->wb, WOTA_NULL);
      break;
  }
  JS_FreeValue (ctx, replaced);
}

static char *decode_wota_value (JSContext *ctx, char *data_ptr, JSValue *out_val, JSValue holder, JSValue key, JSValue reviver) {
  uint64_t first_word = *(uint64_t *)data_ptr;
  int type = (int)(first_word & 0xffU);
  switch (type) {
    case WOTA_INT: {
      long long val;
      data_ptr = wota_read_int (&val, data_ptr);
      *out_val = JS_NewInt64 (ctx, val);
      break;
    }
    case WOTA_FLOAT: {
      double d;
      data_ptr = wota_read_float (&d, data_ptr);
      *out_val = JS_NewFloat64 (ctx, d);
      break;
    }
    case WOTA_SYM: {
      int scode;
      data_ptr = wota_read_sym (&scode, data_ptr);
      if (scode == WOTA_PRIVATE) {
        JSGCRef inner_ref, obj_ref2;
        JS_PushGCRef (ctx, &inner_ref);
        JS_PushGCRef (ctx, &obj_ref2);
        inner_ref.val = JS_NULL;
        data_ptr = decode_wota_value (ctx, data_ptr, &inner_ref.val, holder, JS_NULL, reviver);
        obj_ref2.val = JS_NewObject (ctx);
        if (!JS_IsNull (ctx->actor_sym))
          JS_SetPropertyKey (ctx, obj_ref2.val, ctx->actor_sym, inner_ref.val);
        JS_CellStone (ctx, obj_ref2.val);
        *out_val = obj_ref2.val;
        JS_PopGCRef (ctx, &obj_ref2);
        JS_PopGCRef (ctx, &inner_ref);
      } else if (scode == WOTA_NULL) *out_val = JS_NULL;
      else if (scode == WOTA_FALSE) *out_val = JS_NewBool (ctx, 0);
      else if (scode == WOTA_TRUE) *out_val = JS_NewBool (ctx, 1);
      else *out_val = JS_NULL;
      break;
    }
    case WOTA_BLOB: {
      long long blen;
      char *bdata = NULL;
      data_ptr = wota_read_blob (&blen, &bdata, data_ptr);
      *out_val = bdata ? js_new_blob_stoned_copy (ctx, (uint8_t *)bdata, (size_t)blen) : js_new_blob_stoned_copy (ctx, NULL, 0);
      if (bdata) sys_free (bdata);
      break;
    }
    case WOTA_TEXT: {
      char *utf8 = NULL;
      data_ptr = wota_read_text (&utf8, data_ptr);
      *out_val = JS_NewString (ctx, utf8 ? utf8 : "");
      if (utf8) sys_free (utf8);
      break;
    }
    case WOTA_ARR: {
      long long c;
      data_ptr = wota_read_array (&c, data_ptr);
      JSGCRef arr_ref;
      JS_PushGCRef (ctx, &arr_ref);
      arr_ref.val = JS_NewArrayLen (ctx, c);
      for (long long i = 0; i < c; i++) {
        JSGCRef elem_ref;
        JS_PushGCRef (ctx, &elem_ref);
        elem_ref.val = JS_NULL;
        JSValue idx_key = JS_NewInt32 (ctx, (int32_t)i);
        data_ptr = decode_wota_value (ctx, data_ptr, &elem_ref.val, arr_ref.val, idx_key, reviver);
        JS_SetPropertyNumber (ctx, arr_ref.val, i, elem_ref.val);
        JS_PopGCRef (ctx, &elem_ref);
      }
      *out_val = arr_ref.val;
      JS_PopGCRef (ctx, &arr_ref);
      break;
    }
    case WOTA_REC: {
      long long c;
      data_ptr = wota_read_record (&c, data_ptr);
      JSGCRef obj_ref;
      JS_PushGCRef (ctx, &obj_ref);
      obj_ref.val = JS_NewObject (ctx);
      for (long long i = 0; i < c; i++) {
        char *tkey = NULL;
        size_t key_len;
        data_ptr = wota_read_text_len (&key_len, &tkey, data_ptr);
        if (!tkey) continue;
        JSGCRef prop_key_ref, sub_val_ref;
        JS_PushGCRef (ctx, &prop_key_ref);
        JS_PushGCRef (ctx, &sub_val_ref);
        prop_key_ref.val = JS_NewStringLen (ctx, tkey, key_len);
        sub_val_ref.val = JS_NULL;
        data_ptr = decode_wota_value (ctx, data_ptr, &sub_val_ref.val, obj_ref.val, prop_key_ref.val, reviver);
        JS_SetPropertyStr (ctx, obj_ref.val, tkey, sub_val_ref.val);
        JS_PopGCRef (ctx, &sub_val_ref);
        JS_PopGCRef (ctx, &prop_key_ref);
        sys_free (tkey);
      }
      *out_val = obj_ref.val;
      JS_PopGCRef (ctx, &obj_ref);
      break;
    }
    default:
      data_ptr += 8;
      *out_val = JS_NULL;
      break;
  }
  if (!JS_IsNull (reviver)) {
    JSValue key_val = JS_IsNull (key) ? JS_NULL : JS_DupValue (ctx, key);
    JSValue args[2] = { key_val, JS_DupValue (ctx, *out_val) };
    JSValue revived = JS_Call (ctx, reviver, holder, 2, args);
    JS_FreeValue (ctx, args[0]);
    JS_FreeValue (ctx, args[1]);
    if (!JS_IsException (revived)) {
      JS_FreeValue (ctx, *out_val);
      *out_val = revived;
    } else
      JS_FreeValue (ctx, revived);
  }
  return data_ptr;
}

void *value2wota (JSContext *ctx, JSValue v, JSValue replacer, size_t *bytes) {
  JSGCRef val_ref, rep_ref;
  JS_PushGCRef (ctx, &val_ref);
  JS_PushGCRef (ctx, &rep_ref);
  val_ref.val = v;
  rep_ref.val = replacer;

  WotaEncodeContext enc_s, *enc = &enc_s;

  enc->ctx = ctx;
  enc->visited_stack = NULL;
  enc->cycle = 0;
  enc->replacer = rep_ref.val;
  wota_buffer_init (&enc->wb, 16);
  wota_encode_value (enc, val_ref.val, JS_NULL, JS_NULL);
  if (enc->cycle) {
    wota_stack_free (enc);
    wota_buffer_free (&enc->wb);
    JS_PopGCRef (ctx, &rep_ref);
    JS_PopGCRef (ctx, &val_ref);
    return NULL;
  }
  wota_stack_free (enc);
  size_t total_bytes = enc->wb.size * sizeof (uint64_t);
  void *wota = sys_realloc (enc->wb.data, total_bytes);
  if (bytes) *bytes = total_bytes;
  JS_PopGCRef (ctx, &rep_ref);
  JS_PopGCRef (ctx, &val_ref);
  return wota;
}

JSValue wota2value (JSContext *ctx, void *wota) {
  JSGCRef holder_ref, result_ref;
  JS_PushGCRef (ctx, &holder_ref);
  JS_PushGCRef (ctx, &result_ref);
  result_ref.val = JS_NULL;
  holder_ref.val = JS_NewObject (ctx);
  decode_wota_value (ctx, wota, &result_ref.val, holder_ref.val, JS_NULL, JS_NULL);
  JSValue result = result_ref.val;
  JS_PopGCRef (ctx, &result_ref);
  JS_PopGCRef (ctx, &holder_ref);
  return result;
}

static JSValue js_wota_encode (JSContext *ctx, JSValueConst this_val, int argc, JSValueConst *argv) {
  if (argc < 1) return JS_RaiseDisrupt (ctx, "wota.encode requires at least 1 argument");
  size_t total_bytes;
  void *wota = value2wota (ctx, argv[0], JS_IsFunction (argv[1]) ? argv[1] : JS_NULL, &total_bytes);
  JSValue ret = js_new_blob_stoned_copy (ctx, wota, total_bytes);
  sys_free (wota);
  return ret;
}

static JSValue js_wota_decode (JSContext *ctx, JSValueConst this_val, int argc, JSValueConst *argv) {
  if (argc < 1) return JS_NULL;
  size_t len;
  uint8_t *buf = js_get_blob_data (ctx, &len, argv[0]);
  if (buf == (uint8_t *)-1) return JS_EXCEPTION;
  if (!buf || len == 0) return JS_RaiseDisrupt (ctx, "No blob data present");
  JSValue reviver = (argc > 1 && JS_IsFunction (argv[1])) ? argv[1] : JS_NULL;
  char *data_ptr = (char *)buf;
  JSGCRef result_ref, holder_ref, empty_key_ref;
  JS_PushGCRef (ctx, &result_ref);
  JS_PushGCRef (ctx, &holder_ref);
  JS_PushGCRef (ctx, &empty_key_ref);
  result_ref.val = JS_NULL;
  holder_ref.val = JS_NewObject (ctx);
  empty_key_ref.val = JS_NewString (ctx, "");
  decode_wota_value (ctx, data_ptr, &result_ref.val, holder_ref.val, empty_key_ref.val, reviver);
  JSValue result = result_ref.val;
  JS_PopGCRef (ctx, &empty_key_ref);
  JS_PopGCRef (ctx, &holder_ref);
  JS_PopGCRef (ctx, &result_ref);
  return result;
}

static const JSCFunctionListEntry js_wota_funcs[] = {
  JS_CFUNC_DEF ("encode", 2, js_wota_encode),
  JS_CFUNC_DEF ("decode", 2, js_wota_decode),
};

JSValue js_core_internal_wota_use (JSContext *ctx) {
  JSGCRef exports_ref;
  JS_PushGCRef (ctx, &exports_ref);
  exports_ref.val = JS_NewObject (ctx);
  JS_SetPropertyFunctionList (ctx, exports_ref.val, js_wota_funcs, sizeof (js_wota_funcs) / sizeof (js_wota_funcs[0]));
  JSValue result = exports_ref.val;
  JS_PopGCRef (ctx, &exports_ref);
  return result;
}