#include "Rts.h"
#include "Prelude.h"
#include "RaiseAsync.h"
#include "Schedule.h"
#include "Threads.h"
#include "sm/Sanity.h"

#include <sysexits.h>

#if defined(__wasm_reference_types__)

extern HsBool rts_JSFFI_flag;
extern HsStablePtr rts_threadDelay_impl;

__attribute__((__weak__))
int __main_argc_argv(int argc, char *argv[]);

#if !defined(__PIC__)
void init_ghc_hs_iface(void);

extern StgClosure ghczminternal_GHCziInternalziWasmziPrimziImports_raiseJSException_closure;
extern const StgInfoTable ghczminternal_GHCziInternalziWasmziPrimziTypes_JSVal_con_info;
extern StgClosure ghczminternal_GHCziInternalziWasmziPrimziConcziInternal_threadDelay_closure;

__attribute__((constructor(100)))
static void __init_ghc_hs_iface_jsffi(void) {
  init_ghc_hs_iface();
  ghc_hs_iface->raiseJSException_closure = &ghczminternal_GHCziInternalziWasmziPrimziImports_raiseJSException_closure;
  ghc_hs_iface->JSVal_con_info = &ghczminternal_GHCziInternalziWasmziPrimziTypes_JSVal_con_info;
  ghc_hs_iface->threadDelay_closure = &ghczminternal_GHCziInternalziWasmziPrimziConcziInternal_threadDelay_closure;
}
#endif

// Note [JSFFI initialization]
// ~~~~~~~~~~~~~~~~~~~~~~~~~~~
//
// In the wasi preview1 spec, there are two kinds of wasm32-wasi
// modules: commands and reactors. A command module exports _start()
// that's intended to be invoked exactly once, after which all
// instance state is invalidated. A reactor module supports
// user-defined entrypoints that can be called multiple times after
// _initialize() has been called. In Haskell, when JSFFI is used, we
// expect user-defined JSFFI exports to be entrypoints that can be
// called from JavaScript multiple times, in which case the user needs
// to target wasm32-wasi reactor modules by passing the correct
// link-time options, see user manual for details.
//
// What about RTS initialization? We can tell users to export
// functions like hs_init() as well and call that before calling JSFFI
// exports, but this is very inconvenient. So in this case, we choose
// to do it via a ctor defined here. When JSFFI is not used, JSFFI.o
// will not be included by wasm-ld so this ctor will not get in the
// way. When JSFFI is indeed used, this ctor will be called by the
// linker generated __wasm_call_ctors() function, which is called by
// _initialize(), so the user only needs to call _initialize() once
// and then they can call user exported functions directly.
//
// When it comes to ctors, we must pay close attention to ctor
// priorities to guarantee they are invoked in the correct order. Both
// wasi-libc and emscripten libc makes use of ctors to initialize some
// libc state, and we want them to be invoked first, so priority range
// [0..100] has been fully booked. If priority is ommited, it defaults
// the lowest value 65535, and now there's a problem: there are other
// ctors generated by the GHC codegen (e.g. registering foreign export
// closures as GC roots), and we must ensure those ctors are invoked
// before our RTS initialization logic kicks in!
//
// Therefore, on wasm32, we designate priority 101 to ctors generated
// by the GHC codegen, and priority 102 to the initialization logic
// here to ensure hs_init_ghc() sees everything it needs to see.
//
// It's simpler when it comes to shared libraries: we'll need to load
// them via dyld written in JS anyway, and hence we always compile the
// rts shared library with JSVal functionality. The constructors here
// are only for non-shared objects, so that the user would only need
// to invoke _initialize for all the RTS initialization logic to be
// done before invoking user-specified exports; for the shared library
// case, we explicitly export __ghc_wasm_jsffi_init to be invoked by
// dyld, after at least rts & ghc-internal has been loaded.

#if defined(__PIC__)
__attribute__((export_name("__ghc_wasm_jsffi_init"))) void __ghc_wasm_jsffi_init(void);
#else
__attribute__((constructor(102))) static void __ghc_wasm_jsffi_init(void);
#endif

void __ghc_wasm_jsffi_init(void) {
  // If linking static code without -no-hs-main, then the driver
  // emitted main() is in charge of its own RTS initialization, so
  // skip.
#if !defined(__PIC__)
  if (__main_argc_argv) {
    return;
  }
#endif

  // Code below is mirrored from
  // https://gitlab.haskell.org/haskell-wasm/wasi-libc/-/blob/master/libc-bottom-half/sources/__main_void.c,
  // fetches argc/argv using wasi api
  __wasi_errno_t err;

  // Get the sizes of the arrays we'll have to create to copy in the args.
  size_t argv_buf_size;
  size_t argc;
  err = __wasi_args_sizes_get(&argc, &argv_buf_size);
  if (err != __WASI_ERRNO_SUCCESS) {
    _Exit(EX_OSERR);
  }

  // Add 1 for the NULL pointer to mark the end, and check for overflow.
  size_t num_ptrs = argc + 1;
  if (num_ptrs == 0) {
    _Exit(EX_SOFTWARE);
  }

  // Allocate memory for storing the argument chars.
  char *argv_buf = malloc(argv_buf_size);
  if (argv_buf == NULL) {
    _Exit(EX_SOFTWARE);
  }

  // Allocate memory for the array of pointers. This uses `calloc` both to
  // handle overflow and to initialize the NULL pointer at the end.
  char **argv = calloc(num_ptrs, sizeof(char *));
  if (argv == NULL) {
    free(argv_buf);
    _Exit(EX_SOFTWARE);
  }

  // Fill the argument chars, and the argv array with pointers into those chars.
  // TODO: Remove the casts on `argv_ptrs` and `argv_buf` once the witx is
  // updated with char8 support.
  err = __wasi_args_get((uint8_t **)argv, (uint8_t *)argv_buf);
  if (err != __WASI_ERRNO_SUCCESS) {
    free(argv_buf);
    free(argv);
    _Exit(EX_OSERR);
  }

  // Now that we have argc/argv, proceed to initialize the GHC RTS
  RtsConfig __conf = defaultRtsConfig;
  __conf.rts_opts_enabled = RtsOptsAll;
  __conf.rts_hs_main = false;
#if defined(__PIC__)
  __conf.keep_cafs = 1;
#endif
  hs_init_ghc((int *)&argc, &argv, __conf);
  // See Note [threadDelay on wasm] for details.
  rts_JSFFI_flag = HS_BOOL_TRUE;
  getStablePtr((
      StgPtr)ghc_hs_iface->raiseJSException_closure);
  rts_threadDelay_impl = getStablePtr((
      StgPtr)ghc_hs_iface->threadDelay_closure);
}

typedef __externref_t HsJSVal;
typedef StgInt JSValKey;

extern const StgInfoTable stg_JSVAL_info;

// See Note [JSVal representation for wasm] for detailed explanation.

__attribute__((import_module("ghc_wasm_jsffi"), import_name("newJSVal")))
JSValKey __imported_newJSVal(HsJSVal);

__attribute__((import_module("ghc_wasm_jsffi"),
               import_name("freeJSVal"))) void __imported_freeJSVal(JSValKey);

static void __wrapped_freeJSVal(JSValKey k) {
  __imported_freeJSVal(k);
}

HaskellObj rts_mkJSVal(Capability*, HsJSVal);
HaskellObj rts_mkJSVal(Capability *cap, HsJSVal v) {
  JSValKey k = __imported_newJSVal(v);

  HaskellObj p = (HaskellObj)allocate(cap, CONSTR_sizeW(1, 2));
  SET_HDR(p, &stg_JSVAL_info, CCS_SYSTEM);
  p->payload[1] = (HaskellObj)k;
  p->payload[2] = NULL;

  StgCFinalizerList *cfin =
      (StgCFinalizerList *)allocate(cap, sizeofW(StgCFinalizerList));
  SET_HDR(cfin, &stg_C_FINALIZER_LIST_info, CCS_SYSTEM);
  cfin->link = &stg_NO_FINALIZER_closure;
  cfin->fptr = (void (*)(void))__wrapped_freeJSVal;
  cfin->ptr = (void *)k;
  cfin->flag = 0;

  StgWeak *w = (StgWeak *)allocate(cap, sizeofW(StgWeak));
  SET_HDR(w, &stg_WEAK_info, CCS_SYSTEM);
  w->cfinalizers = (StgClosure *)cfin;
  w->key = p;
  w->value = Unit_closure;
  w->finalizer = &stg_NO_FINALIZER_closure;
  w->link = cap->weak_ptr_list_hd;
  cap->weak_ptr_list_hd = w;
  if (cap->weak_ptr_list_tl == NULL) {
    cap->weak_ptr_list_tl = w;
  }

  p->payload[0] = (HaskellObj)w;

  HaskellObj box = (HaskellObj)allocate(cap, CONSTR_sizeW(1, 0));
  SET_HDR(box, ghc_hs_iface->JSVal_con_info, CCS_SYSTEM);
  box->payload[0] = p;

  return TAG_CLOSURE(1, box);
}

__attribute__((import_module("ghc_wasm_jsffi"), import_name("getJSVal")))
HsJSVal __imported_getJSVal(JSValKey);

STATIC_INLINE HsJSVal rts_getJSValzh(HaskellObj p) {
  ASSERT(p->header.info == &stg_JSVAL_info);
  return __imported_getJSVal((JSValKey)p->payload[1]);
}

HsJSVal rts_getJSVal(HaskellObj);
HsJSVal rts_getJSVal(HaskellObj box) {
  ASSERT(UNTAG_CLOSURE(box)->header.info == ghc_hs_iface->JSVal_con_info);
  return rts_getJSValzh(UNTAG_CLOSURE(box)->payload[0]);
}

INLINE_HEADER void pushClosure   (StgTSO *tso, StgWord c) {
  tso->stackobj->sp--;
  tso->stackobj->sp[0] = (W_) c;
}

extern const StgInfoTable stg_scheduler_loop_info;

// schedule a future round of RTS scheduler loop via setImmediate(),
// to avoid jamming the JavaScript main thread

__attribute__((import_module("ghc_wasm_jsffi"), import_name("scheduleWork")))
void __imported_scheduleWork(void);

void rts_scheduleWork(void);
void rts_scheduleWork(void) {
  __imported_scheduleWork();
}

__attribute__((export_name("rts_schedulerLoop")))
void rts_schedulerLoop(void);
void rts_schedulerLoop(void) {
  Capability *cap = rts_lock();
  StgTSO *tso = createThread(cap, RESERVED_STACK_WORDS);
  pushClosure(tso, (StgWord)&stg_scheduler_loop_info);
  scheduleWaitThread(tso, NULL, &cap);
  rts_checkSchedStatus("rts_schedulerLoop", cap);
  rts_unlock(cap);
}

#define mk_rtsPromiseCallback(obj)                         \
  {                                                        \
  Capability *cap = &MainCapability;                       \
  hs_try_putmvar_with_value(cap->no, sp, obj);             \
  rts_schedulerLoop();                                     \
  }

#define mk_rtsPromiseResolve(T)                            \
  __attribute__((export_name("rts_promiseResolve"#T)))     \
  void rts_promiseResolve##T(HsStablePtr, Hs##T);          \
  void rts_promiseResolve##T(HsStablePtr sp, Hs##T js_res) \
  mk_rtsPromiseCallback(rts_mk##T(cap, js_res))

__attribute__((export_name("rts_promiseResolveUnit")))
void rts_promiseResolveUnit(HsStablePtr);
void rts_promiseResolveUnit(HsStablePtr sp)
  mk_rtsPromiseCallback(TAG_CLOSURE(1, Unit_closure))

mk_rtsPromiseResolve(JSVal)
mk_rtsPromiseResolve(Char)
mk_rtsPromiseResolve(Int)
mk_rtsPromiseResolve(Int8)
mk_rtsPromiseResolve(Int16)
mk_rtsPromiseResolve(Int32)
mk_rtsPromiseResolve(Int64)
mk_rtsPromiseResolve(Word)
mk_rtsPromiseResolve(Word8)
mk_rtsPromiseResolve(Word16)
mk_rtsPromiseResolve(Word32)
mk_rtsPromiseResolve(Word64)
mk_rtsPromiseResolve(Ptr)
mk_rtsPromiseResolve(FunPtr)
mk_rtsPromiseResolve(Float)
mk_rtsPromiseResolve(Double)
mk_rtsPromiseResolve(StablePtr)
mk_rtsPromiseResolve(Bool)

__attribute__((export_name("rts_promiseReject")))
void rts_promiseReject(HsStablePtr, HsJSVal);
void rts_promiseReject(HsStablePtr sp, HsJSVal js_err)
  mk_rtsPromiseCallback(rts_apply(cap, ghc_hs_iface->raiseJSException_closure, rts_mkJSVal(cap, js_err)))

__attribute__((export_name("rts_promiseThrowTo")))
void rts_promiseThrowTo(HsStablePtr, HsJSVal);
void rts_promiseThrowTo(HsStablePtr sp, HsJSVal js_err) {
  Capability *cap = &MainCapability;
  StgWeak *w = (StgWeak *)deRefStablePtr(sp);
  if (w->header.info == &stg_DEAD_WEAK_info) {
    return;
  }
  ASSERT(w->header.info == &stg_WEAK_info);
  StgTSO *tso = (StgTSO *)w->key;
  ASSERT(tso->header.info == &stg_TSO_info);
  throwToSelf(
      cap, tso,
      rts_apply(
          cap,
          ghc_hs_iface->raiseJSException_closure,
          rts_mkJSVal(cap, js_err)));
  tryWakeupThread(cap, tso);
  rts_schedulerLoop();
}

__attribute__((export_name("rts_freeStablePtr")))
void rts_freeStablePtr(HsStablePtr);
void rts_freeStablePtr(HsStablePtr sp) {
  hs_free_stable_ptr(sp);
}

#endif // __wasm_reference_types__
