#include "qjs_wota.h"
#include "qjs_blob.h"

#include "wota.h"
#include <stdlib.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)
{
  if (!JS_IsObject(val)) return;
  
  ObjectRef *ref = malloc(sizeof(ObjectRef));
  if (!ref) return;
  
  ref->ptr = JS_VALUE_GET_PTR(val);
  ref->next = enc->visited_stack;
  enc->visited_stack = ref;
}

static void wota_stack_pop(WotaEncodeContext *enc)
{
  if (!enc->visited_stack) return;
  
  ObjectRef *top = enc->visited_stack;
  enc->visited_stack = top->next;
  free(top);
}

static int wota_stack_has(WotaEncodeContext *enc, JSValueConst val)
{
  if (!JS_IsObject(val)) return 0;
  
  void *ptr = JS_VALUE_GET_PTR(val);
  ObjectRef *current = enc->visited_stack;
  
  while (current) {
    if (current->ptr == ptr) return 1;
    current = current->next;
  }
  return 0;
}

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

static JSValue apply_replacer(WotaEncodeContext *enc, JSValueConst holder, JSAtom key, JSValueConst val)
{
  if (JS_IsNull(enc->replacer)) return JS_DupValue(enc->ctx, val);
  JSValue key_val = (key == JS_ATOM_NULL) ? JS_NULL : JS_AtomToValue(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, JSAtom key);

static void encode_object_properties(WotaEncodeContext *enc, JSValueConst val, JSValueConst holder)
{
  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) {
    wota_write_sym(&enc->wb, WOTA_NULL);
    return;
  }
  uint32_t non_function_count = 0;
  JSValue props[plen];
  JSAtom atoms[plen];
  
  for (uint32_t i = 0; i < plen; i++) {
    JSValue prop_val = JS_GetProperty(ctx, val, ptab[i].atom);
    if (!JS_IsFunction(ctx, prop_val)) {
      atoms[non_function_count] = ptab[i].atom;
      props[non_function_count++] = prop_val;
    } else
      JS_FreeValue(ctx, prop_val);
  }
  wota_write_record(&enc->wb, non_function_count);
  for (uint32_t i = 0; i < non_function_count; i++) {
    size_t plen;
    const char *prop_name = JS_AtomToCStringLen(ctx, &plen, atoms[i]);
    JSValue prop_val = props[i];
    wota_write_text_len(&enc->wb, prop_name, plen);
    wota_encode_value(enc, prop_val, val, atoms[i]);
    JS_FreeCString(ctx, prop_name);
    JS_FreeValue(ctx, prop_val);
  }

  for (int i = 0; i < plen; i++)
    JS_FreeAtom(ctx, ptab[i].atom);

  js_free(ctx, ptab);
}

static void wota_encode_value(WotaEncodeContext *enc, JSValueConst val, JSValueConst holder, JSAtom key)
{
  JSContext *ctx = enc->ctx;
  JSValue replaced;
  if (!JS_IsNull(enc->replacer) && key != JS_ATOM_NULL)
    replaced = 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:
    case JS_TAG_BIG_INT: {
      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_OBJECT: {
      if (js_is_blob(ctx, replaced)) {
        size_t buf_len;
        void *buf_data = js_get_blob_data(ctx, &buf_len, replaced);
        wota_write_blob(&enc->wb, (unsigned long long)buf_len * 8, (const char *)buf_data);
        break;
      }
      if (JS_IsArray(ctx, 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_GetPropertyUint32(ctx, replaced, i);
          JSAtom idx_atom = JS_NewAtomUInt32(ctx, (uint32_t)i);
          wota_encode_value(enc, elem_val, replaced, idx_atom);
          JS_FreeAtom(ctx, idx_atom);
          JS_FreeValue(ctx, elem_val);
        }
        wota_stack_pop(enc);
        break;
      }
      cell_rt *crt = JS_GetContextOpaque(ctx);
      JSValue adata = JS_GetProperty(ctx, replaced, crt->actor_sym);
      if (!JS_IsNull(adata)) {
        wota_write_sym(&enc->wb, WOTA_PRIVATE);
        wota_encode_value(enc, adata, replaced, JS_ATOM_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(ctx, 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, JSAtom 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) {
        JSValue inner = JS_NULL;
        data_ptr = decode_wota_value(ctx, data_ptr, &inner, holder, JS_ATOM_NULL, reviver);
        JSValue obj = JS_NewObject(ctx);
        cell_rt *crt = JS_GetContextOpaque(ctx);
        JS_SetProperty(ctx, obj, crt->actor_sym, inner);
        *out_val = obj;
      } 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) 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) free(utf8);
      break;
    }
    case WOTA_ARR: {
      long long c;
      data_ptr = wota_read_array(&c, data_ptr);
      JSValue arr = JS_NewArray(ctx);
      JS_SetLength(ctx, arr, c);
      for (long long i = 0; i < c; i++) {
        JSValue elem_val = JS_NULL;
        JSAtom idx_atom = JS_NewAtomUInt32(ctx, (uint32_t)i);
        data_ptr = decode_wota_value(ctx, data_ptr, &elem_val, arr, idx_atom, reviver);
        JS_SetPropertyUint32(ctx, arr, i, elem_val);
        JS_FreeAtom(ctx, idx_atom);
      }
      *out_val = arr;
      break;
    }
    case WOTA_REC: {
      long long c;
      data_ptr = wota_read_record(&c, data_ptr);
      JSValue obj = 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; // invalid key
        JSAtom prop_key = JS_NewAtomLen(ctx, tkey, key_len);
        JSValue sub_val = JS_NULL;
        data_ptr = decode_wota_value(ctx, data_ptr, &sub_val, obj, prop_key, reviver);
        JS_SetProperty(ctx, obj, prop_key, sub_val);
        JS_FreeAtom(ctx, prop_key);
        free(tkey);
      }
      *out_val = obj;
      break;
    }
    default:
      data_ptr += 8;
      *out_val = JS_NULL;
      break;
  }
  if (!JS_IsNull(reviver)) {
    JSValue key_val = (key == JS_ATOM_NULL) ? JS_NULL : JS_AtomToValue(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)
{
  WotaEncodeContext enc_s, *enc = &enc_s;
  
  enc->ctx = ctx;
  enc->visited_stack = NULL;
  enc->cycle = 0;
  enc->replacer = replacer;
  wota_buffer_init(&enc->wb, 16);
  wota_encode_value(enc, v, JS_NULL, JS_ATOM_NULL);
  if (enc->cycle) {
    wota_stack_free(enc);
    wota_buffer_free(&enc->wb);
    return NULL;
  }
  wota_stack_free(enc);
  size_t total_bytes = enc->wb.size * sizeof(uint64_t);
  void *wota = realloc(enc->wb.data, total_bytes);
  if (bytes) *bytes = total_bytes;
  return wota;
}

JSValue wota2value(JSContext *ctx, void *wota)
{
  JSValue result = JS_NULL;
  JSValue holder = JS_NewObject(ctx);
  decode_wota_value(ctx, wota, &result, holder, JS_ATOM_NULL, JS_NULL);
  JS_FreeValue(ctx, holder);
  return result;
}

static JSValue js_wota_encode(JSContext *ctx, JSValueConst this_val, int argc, JSValueConst *argv)
{
  if (argc < 1) return JS_ThrowTypeError(ctx, "wota.encode requires at least 1 argument");
  size_t total_bytes;
  void *wota = value2wota(ctx, argv[0], JS_IsFunction(ctx,argv[1]) ? argv[1] : JS_NULL, &total_bytes);
  JSValue ret = js_new_blob_stoned_copy(ctx, wota, total_bytes);
  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) return JS_NULL;
  JSValue reviver = (argc > 1 && JS_IsFunction(ctx, argv[1])) ? argv[1] : JS_NULL;
  char *data_ptr = (char *)buf;
  JSValue result = JS_NULL;
  JSValue holder = JS_NewObject(ctx);
  JSAtom empty_atom = JS_NewAtom(ctx, "");
  decode_wota_value(ctx, data_ptr, &result, holder, empty_atom, reviver);
  JS_FreeAtom(ctx, empty_atom);
  JS_FreeValue(ctx, holder);
  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_wota_use(JSContext *ctx)
{
  JSValue exports = JS_NewObject(ctx);
  JS_SetPropertyFunctionList(ctx, exports, js_wota_funcs, sizeof(js_wota_funcs)/sizeof(js_wota_funcs[0]));
  return exports;
}
