///|
/// Recoverable failure while preparing per-instance native state.
pub(all) suberror NativeSetupError {
  SegmentAllocationFailed
  RootFrameAllocationFailed
  ExceptionPayloadAllocationFailed
  DebugRegistrationFailed
  MemoryBindingAllocationFailed
  TableBindingAllocationFailed
  InvalidBinding(String)
} derive(Debug)

///|
pub extend NativeSetupError with Debug::{to_repr}

///|
pub extend NativeSetupError with Show::{output, to_string}

///|
pub impl Show for NativeSetupError with fn output(self, logger) {
  logger.write_string(
    match self {
      SegmentAllocationFailed => "native segment allocation failed"
      RootFrameAllocationFailed => "native root frame allocation failed"
      ExceptionPayloadAllocationFailed =>
        "native exception payload allocation failed"
      DebugRegistrationFailed => "native debug metadata registration failed"
      MemoryBindingAllocationFailed => "native memory binding allocation failed"
      TableBindingAllocationFailed => "native table binding allocation failed"
      InvalidBinding(message) => "invalid native binding: \{message}"
    },
  )
}

/// Execution, callback, exception, fiber-stack, and GC operations for managed JIT contexts.

///|
fn check_exception_status(status : Int32) -> Unit raise NativeSetupError {
  match status {
    1 => ()
    0 => raise ExceptionPayloadAllocationFailed
    _ =>
      raise InvalidBinding(
        "exception delivery requires the active host callback and a valid tag",
      )
  }
}

///|
fn setup_segments(
  context : JITContext,
  datas : Array[@types.Data],
  elem_segments : Array[Array[Int64]],
  data_dropped~ : Array[Bool] = [],
  elem_dropped~ : Array[Bool] = [],
) -> Unit raise NativeSetupError {
  let data = Array::makei(datas.length(), i => {
    if i < data_dropped.length() && data_dropped[i] {
      b""
    } else {
      datas[i].init
    }
  })
  let elements = Array::makei(elem_segments.length(), i => {
    if i < elem_dropped.length() && elem_dropped[i] {
      []
    } else {
      elem_segments[i]
    }
  })
  guard context.set_segments(data, elements) else {
    raise SegmentAllocationFailed
  }
}

///|
fn clear_segments(context : JITContext) -> Unit {
  c_jit_ctx_clear_segments_managed(context)
}

///|
pub fn NativeJITContext::func_count(self : NativeJITContext) -> Int32 {
  context_func_count(self.handle)
}

///|
/// Stage a module-local exception. Guest delivery occurs after the callback returns.
pub fn NativeJITContext::throw_exception_tag(
  self : NativeJITContext,
  tag_addr : Int32,
) -> Unit raise NativeSetupError {
  check_exception_status(
    c_jit_exception_throw_tag_managed(self.handle, tag_addr),
  )
}

///|
/// Stage a module-local exception. Guest delivery occurs after the callback returns.
pub fn NativeJITContext::throw_exception_values(
  self : NativeJITContext,
  tag_addr : Int32,
  values : Array[Int64],
) -> Unit raise NativeSetupError {
  let payload = FixedArray::makei(values.length(), i => values[i])
  check_exception_status(
    c_jit_exception_throw_values_managed(
      self.handle,
      tag_addr,
      payload,
      values.length(),
    ),
  )
}

///|
/// Stage an imported exception using its Store tag identity.
pub fn NativeJITContext::throw_store_exception_values(
  self : NativeJITContext,
  tag_addr : Int32,
  values : Array[Int64],
) -> Unit raise NativeSetupError {
  let payload = FixedArray::makei(values.length(), i => values[i])
  check_exception_status(
    c_jit_exception_throw_store_values_managed(
      self.handle,
      tag_addr,
      payload,
      values.length(),
    ),
  )
}

///|
pub fn NativeJITContext::setup_segments(
  self : NativeJITContext,
  datas : Array[@types.Data],
  elem_segments : Array[Array[Int64]],
  data_dropped~ : Array[Bool] = [],
  elem_dropped~ : Array[Bool] = [],
) -> Unit raise NativeSetupError {
  setup_segments(
    self.handle,
    datas,
    elem_segments,
    data_dropped~,
    elem_dropped~,
  )
}

///|
pub fn NativeJITContext::clear_segments(self : NativeJITContext) -> Unit {
  clear_segments(self.handle)
}

///|
pub fn NativeJITContext::clear_gc_heap(self : NativeJITContext) -> Unit {
  c_jit_gc_clear_heap_managed(self.handle)
}

///|
pub fn NativeJITContext::set_func(
  self : NativeJITContext,
  idx : Int32,
  func_ptr : Int64,
) -> Unit {
  c_jit_ctx_set_func_managed(self.handle, idx, func_ptr)
}

///|
pub fn NativeJITContext::alloc_indirect_table(
  self : NativeJITContext,
  count : Int32,
) -> Bool {
  c_jit_ctx_alloc_indirect_table_managed(self.handle, count) != 0
}

///|
pub fn NativeJITContext::set_indirect(
  self : NativeJITContext,
  table_idx : Int32,
  func_idx : Int32,
  type_idx : Int32,
) -> Unit {
  c_jit_ctx_set_indirect_managed(self.handle, table_idx, func_idx, type_idx)
}

///|
/// Snapshot borrowed global-cell addresses and retain the table until replacement.
pub fn NativeJITContext::bind_globals(
  self : NativeJITContext,
  globals : FixedArray[Int64],
) -> Int64 {
  self.handle.bind_globals(FixedArray::makei(globals.length(), i => globals[i]))
}

///|
/// Transfer a malloc-owned table of borrowed cell addresses to this context.
/// Prefer bind_globals for MoonBit-owned input.
pub fn NativeJITContext::set_globals(
  self : NativeJITContext,
  globals_ptr : Int64,
) -> Unit {
  c_jit_ctx_set_globals_managed(self.handle, globals_ptr)
}

///|
/// Configure the guarded stack size used by future native JIT continuations.
/// The stack mapping is allocated only when a continuation starts.
#alias(alloc_wasm_stack, deprecated="Use configure_native_fiber_stack_size instead")
pub fn NativeJITContext::configure_native_fiber_stack_size(
  self : NativeJITContext,
  stack_size : Int64,
) -> Bool {
  if stack_size <= 0L {
    false
  } else {
    self.configured_native_fiber_stack_size = stack_size
    true
  }
}

///|
/// Check whether a native-fiber stack size override has been configured.
#alias(has_wasm_stack, deprecated="Use has_configured_native_fiber_stack_size instead")
pub fn NativeJITContext::has_configured_native_fiber_stack_size(
  self : NativeJITContext,
) -> Bool {
  self.configured_native_fiber_stack_size > 0L
}

///|
/// Typed result returned by a JIT hostcall dispatcher.
pub(all) enum NativeHostcallOutcome {
  HostcallCompleted
  HostcallExit(Int32)
  HostcallTrap(Int32)
  HostcallSuspended
} derive(Debug, Eq)

///|
pub fn NativeJITContext::set_hostcall_callback(
  self : NativeJITContext,
  callback : () -> NativeHostcallOutcome,
) -> Unit {
  let encoded = fn() {
    match callback() {
      HostcallCompleted => 0
      HostcallExit(code) => {
        c_jit_set_wasi_exit_code_managed(self.handle, code)
        100
      }
      HostcallTrap(code) => if code > 0 { code } else { 8 }
      HostcallSuspended => -1
    }
  }
  let call_closure : FuncRef[(() -> Int32) -> Int32] = fn(f : () -> Int32) {
    f()
  }
  c_jit_set_hostcall_callback_managed(self.handle, call_closure, encoded)
}

///|
pub fn NativeJITContext::clear_hostcall_callback(
  self : NativeJITContext,
) -> Unit {
  c_jit_clear_hostcall_callback_managed(self.handle)
}

///|
/// Set the cancellation callback for subsequent invocations. Active and parked
/// invocations retain their existing callback; mutate its captured signal to cancel.
pub fn NativeJITContext::set_cancellation_callback(
  self : NativeJITContext,
  callback : () -> Bool,
) -> Unit {
  let call_closure : FuncRef[(() -> Bool) -> Int32] = fn(f : () -> Bool) {
    if f() {
      1
    } else {
      0
    }
  }
  c_jit_set_cancellation_callback_managed(self.handle, call_closure, callback)
}

///|
/// Clear the default callback without changing active or parked invocations.
pub fn NativeJITContext::clear_cancellation_callback(
  self : NativeJITContext,
) -> Unit {
  c_jit_clear_cancellation_callback_managed(self.handle)
}

///|
#borrow(context)
extern "c" fn c_set_cooperative_scheduling(
  context : JITContext,
  enabled : Bool,
) -> Unit = "wasmoon_jit_set_cooperative_scheduling_managed"

///|
/// Configure native scheduling for subsequent invocations at instrumented
/// guest safepoints. Active invocations keep their own budget. The embedding
/// must enter this context through a resumable native invocation.
pub fn NativeJITContext::set_cooperative_scheduling(
  self : NativeJITContext,
  enabled : Bool,
) -> Unit {
  c_set_cooperative_scheduling(self.handle, enabled)
}

///|
pub fn NativeJITContext::call_trampoline(
  self : NativeJITContext,
  trampoline_ptr : Int64,
  func_ptr : Int64,
  values_vec : FixedArray[Int64],
  values_len : Int32,
  use_native_fiber : Bool,
) -> Int32 raise NativeFiberError {
  if use_native_fiber {
    let continuation = self.start_continuation(
      trampoline_ptr, func_ptr, values_vec, values_len,
    )
    let event = {
      errdefer (continuation.cancel() catch { _ => () })
      continuation.continue_with()
    }
    match event {
      NativeJITCompleted => 0
      NativeJITTrapped(code) => code
      NativeSuspended(_) => {
        continuation.cancel()
        raise UnexpectedYield(-1L)
      }
    }
  } else {
    c_jit_call_trampoline_managed(
      self.handle,
      trampoline_ptr,
      func_ptr,
      values_vec,
      values_len,
    )
  }
}

///|
/// Share instance-local dropped bits with the interpreter and Store GC.
pub fn NativeJITContext::bind_segment_drop_flags(
  self : NativeJITContext,
  data : FixedArray[Bool],
  elements : FixedArray[Bool],
) -> Unit {
  self.handle.bind_segment_drop_flags(data, elements)
}