///|
struct Backend(Int) derive(Eq, Debug)

///|
pub fn Backend::metal() -> Backend {
  Backend(0)
}

///|
pub fn Backend::cuda() -> Backend {
  Backend(1)
}

///|
pub fn Backend::cpu() -> Backend {
  Backend(2)
}

///|
extern "c" fn ds4_backend_default() -> Int = "moonbit_ds4_backend_default"

///|
pub fn Backend::default() -> Backend {
  Backend(ds4_backend_default())
}

///|
fn Backend::to_int(self : Backend) -> Int {
  self.0
}

///|
extern "c" fn ds4_backend_name(backend : Int) -> Bytes = "moonbit_ds4_backend_name"

///|
pub fn Backend::name(self : Backend) -> String {
  @utf8.decode_lossy(ds4_backend_name(self.0))
}

///|
struct ThinkMode(Int) derive(Eq, Debug)

///|
pub fn ThinkMode::none() -> ThinkMode {
  ThinkMode(0)
}

///|
pub fn ThinkMode::high() -> ThinkMode {
  ThinkMode(1)
}

///|
pub fn ThinkMode::max() -> ThinkMode {
  ThinkMode(2)
}

///|
fn ThinkMode::to_int(self : ThinkMode) -> Int {
  self.0
}

///|
extern "c" fn ds4_think_mode_name(mode : Int) -> Bytes = "moonbit_ds4_think_mode_name"

///|
pub fn ThinkMode::name(self : ThinkMode) -> String {
  @utf8.decode_lossy(ds4_think_mode_name(self.0))
}

///|
extern "c" fn ds4_think_mode_enabled(mode : Int) -> Bool = "moonbit_ds4_think_mode_enabled"

///|
pub fn ThinkMode::enabled(self : ThinkMode) -> Bool {
  ds4_think_mode_enabled(self.0)
}

///|
extern "c" fn ds4_think_max_prefix() -> Bytes = "moonbit_ds4_think_max_prefix"

///|
pub fn ThinkMode::max_prefix() -> String {
  @utf8.decode_lossy(ds4_think_max_prefix())
}

///|
extern "c" fn ds4_think_max_min_context() -> UInt = "moonbit_ds4_think_max_min_context"

///|
pub fn ThinkMode::max_min_context() -> UInt {
  ds4_think_max_min_context()
}

///|
extern "c" fn ds4_think_mode_for_context(mode : Int, ctx_size : Int) -> Int = "moonbit_ds4_think_mode_for_context"

///|
pub fn ThinkMode::for_context(self : ThinkMode, ctx_size : Int) -> ThinkMode {
  ThinkMode(ds4_think_mode_for_context(self.0, ctx_size))
}

///|
pub(all) struct ContextMemory {
  total_bytes : UInt64
  raw_bytes : UInt64
  compressed_bytes : UInt64
  scratch_bytes : UInt64
  prefill_cap : UInt
  raw_cap : UInt
  comp_cap : UInt
} derive(Eq, Debug)

///|
extern "c" fn ds4_context_memory_total(backend : Int, ctx_size : Int) -> UInt64 = "moonbit_ds4_context_memory_total"

///|
extern "c" fn ds4_context_memory_raw(backend : Int, ctx_size : Int) -> UInt64 = "moonbit_ds4_context_memory_raw"

///|
extern "c" fn ds4_context_memory_compressed(
  backend : Int,
  ctx_size : Int,
) -> UInt64 = "moonbit_ds4_context_memory_compressed"

///|
extern "c" fn ds4_context_memory_scratch(
  backend : Int,
  ctx_size : Int,
) -> UInt64 = "moonbit_ds4_context_memory_scratch"

///|
extern "c" fn ds4_context_memory_prefill_cap(
  backend : Int,
  ctx_size : Int,
) -> UInt = "moonbit_ds4_context_memory_prefill_cap"

///|
extern "c" fn ds4_context_memory_raw_cap(backend : Int, ctx_size : Int) -> UInt = "moonbit_ds4_context_memory_raw_cap"

///|
extern "c" fn ds4_context_memory_comp_cap(
  backend : Int,
  ctx_size : Int,
) -> UInt = "moonbit_ds4_context_memory_comp_cap"

///|
pub fn Backend::context_memory_estimate(
  self : Backend,
  ctx_size : Int,
) -> ContextMemory {
  let backend = self.0
  {
    total_bytes: ds4_context_memory_total(backend, ctx_size),
    raw_bytes: ds4_context_memory_raw(backend, ctx_size),
    compressed_bytes: ds4_context_memory_compressed(backend, ctx_size),
    scratch_bytes: ds4_context_memory_scratch(backend, ctx_size),
    prefill_cap: ds4_context_memory_prefill_cap(backend, ctx_size),
    raw_cap: ds4_context_memory_raw_cap(backend, ctx_size),
    comp_cap: ds4_context_memory_comp_cap(backend, ctx_size),
  }
}

///|
pub(all) suberror Ds4Error {
  OpenFailed(Int)
  SessionCreateFailed(Int, String)
  OperationFailed(op~ : String, status~ : Int, message~ : String)
} derive(Eq, Debug)

///|
pub(all) struct EngineConfig {
  model_path : String
  mtp_path : String
  backend : Backend
  n_threads : Int
  mtp_draft_tokens : Int
  mtp_margin : Float
  directional_steering_file : String
  directional_steering_attn : Float
  directional_steering_ffn : Float
  warm_weights : Bool
  quality : Bool
  inspect_only : Bool
} derive(Eq, Debug)

///|
pub fn EngineConfig::new(
  model_path : String,
  backend? : Backend = Backend::default(),
  mtp_path? : String = "",
  n_threads? : Int = 0,
  mtp_draft_tokens? : Int = 1,
  mtp_margin? : Float = -1.0,
  directional_steering_file? : String = "",
  directional_steering_attn? : Float = 0.0,
  directional_steering_ffn? : Float = 0.0,
  warm_weights? : Bool = false,
  quality? : Bool = false,
  inspect_only? : Bool = false,
) -> EngineConfig {
  {
    model_path,
    mtp_path,
    backend,
    n_threads,
    mtp_draft_tokens,
    mtp_margin,
    directional_steering_file,
    directional_steering_attn,
    directional_steering_ffn,
    warm_weights,
    quality,
    inspect_only,
  }
}

///|
type Engine

///|
type Session

///|
type Tokens

///|
type Sampler

///|
extern "c" fn ds4_engine_new() -> Engine = "moonbit_ds4_engine_new"

///|
#borrow(engine, model_path, mtp_path, directional_steering_file)
extern "c" fn ds4_engine_open(
  engine : Engine,
  model_path : Bytes,
  mtp_path : Bytes,
  backend : Int,
  n_threads : Int,
  mtp_draft_tokens : Int,
  mtp_margin : Float,
  directional_steering_file : Bytes,
  directional_steering_attn : Float,
  directional_steering_ffn : Float,
  warm_weights : Bool,
  quality : Bool,
  inspect_only : Bool,
) -> Int = "moonbit_ds4_engine_open"

///|
pub fn Engine::open(config : EngineConfig) -> Engine raise Ds4Error {
  let engine = ds4_engine_new()
  let status = ds4_engine_open(
    engine,
    @utf8.encode(config.model_path),
    @utf8.encode(config.mtp_path),
    config.backend.to_int(),
    config.n_threads,
    config.mtp_draft_tokens,
    config.mtp_margin,
    @utf8.encode(config.directional_steering_file),
    config.directional_steering_attn,
    config.directional_steering_ffn,
    config.warm_weights,
    config.quality,
    config.inspect_only,
  )
  if status != 0 {
    raise Ds4Error::OpenFailed(status)
  }
  engine
}

///|
#borrow(engine)
extern "c" fn ds4_engine_summary(engine : Engine) -> Unit = "moonbit_ds4_engine_summary"

///|
pub fn Engine::summary(self : Engine) -> Unit {
  ds4_engine_summary(self)
}

///|
#borrow(engine)
extern "c" fn ds4_engine_token_eos(engine : Engine) -> Int = "moonbit_ds4_engine_token_eos"

///|
pub fn Engine::token_eos(self : Engine) -> Int {
  ds4_engine_token_eos(self)
}

///|
#borrow(engine)
extern "c" fn ds4_engine_routed_quant_bits(engine : Engine) -> Int = "moonbit_ds4_engine_routed_quant_bits"

///|
pub fn Engine::routed_quant_bits(self : Engine) -> Int {
  ds4_engine_routed_quant_bits(self)
}

///|
#borrow(engine)
extern "c" fn ds4_engine_has_mtp(engine : Engine) -> Bool = "moonbit_ds4_engine_has_mtp"

///|
pub fn Engine::has_mtp(self : Engine) -> Bool {
  ds4_engine_has_mtp(self)
}

///|
#borrow(engine)
extern "c" fn ds4_engine_mtp_draft_tokens(engine : Engine) -> Int = "moonbit_ds4_engine_mtp_draft_tokens"

///|
pub fn Engine::mtp_draft_tokens(self : Engine) -> Int {
  ds4_engine_mtp_draft_tokens(self)
}

///|
#borrow(engine)
extern "c" fn ds4_engine_token_text(engine : Engine, token : Int) -> Bytes = "moonbit_ds4_engine_token_text"

///|
pub fn Engine::token_text(self : Engine, token : Int) -> String {
  @utf8.decode_lossy(ds4_engine_token_text(self, token))
}

///|
extern "c" fn ds4_tokens_new() -> Tokens = "moonbit_ds4_tokens_new"

///|
pub fn Tokens::new() -> Tokens {
  ds4_tokens_new()
}

///|
extern "c" fn ds4_sampler_new(seed : UInt64) -> Sampler = "moonbit_ds4_sampler_new"

///|
pub fn Sampler::new(seed : UInt64) -> Sampler {
  ds4_sampler_new(seed)
}

///|
#borrow(tokens)
extern "c" fn ds4_tokens_copy(tokens : Tokens) -> Tokens = "moonbit_ds4_tokens_copy"

///|
pub fn Tokens::copy(self : Tokens) -> Tokens {
  ds4_tokens_copy(self)
}

///|
#borrow(tokens)
extern "c" fn ds4_tokens_push(tokens : Tokens, token : Int) -> Unit = "moonbit_ds4_tokens_push"

///|
pub fn Tokens::push(self : Tokens, token : Int) -> Unit {
  ds4_tokens_push(self, token)
}

///|
#borrow(tokens)
extern "c" fn ds4_tokens_length(tokens : Tokens) -> Int = "moonbit_ds4_tokens_length"

///|
pub fn Tokens::length(self : Tokens) -> Int {
  ds4_tokens_length(self)
}

///|
#borrow(tokens)
extern "c" fn ds4_tokens_get(tokens : Tokens, index : Int) -> Int = "moonbit_ds4_tokens_get"

///|
pub fn Tokens::get(self : Tokens, index : Int) -> Int? {
  if index < 0 || index >= self.length() {
    None
  } else {
    Some(ds4_tokens_get(self, index))
  }
}

///|
pub fn Tokens::op_get(self : Tokens, index : Int) -> Int {
  if index < 0 || index >= self.length() {
    abort("token index out of bounds")
  }
  ds4_tokens_get(self, index)
}

///|
#borrow(tokens)
extern "c" fn ds4_tokens_to_fixed_array(tokens : Tokens) -> FixedArray[Int] = "moonbit_ds4_tokens_to_fixed_array"

///|
pub fn Tokens::to_fixed_array(self : Tokens) -> FixedArray[Int] {
  ds4_tokens_to_fixed_array(self)
}

///|
#borrow(tokens, prefix)
extern "c" fn ds4_tokens_starts_with(tokens : Tokens, prefix : Tokens) -> Bool = "moonbit_ds4_tokens_starts_with"

///|
pub fn Tokens::starts_with(self : Tokens, prefix : Tokens) -> Bool {
  ds4_tokens_starts_with(self, prefix)
}

///|
#borrow(engine, text)
extern "c" fn ds4_engine_tokenize_text(engine : Engine, text : Bytes) -> Tokens = "moonbit_ds4_engine_tokenize_text"

///|
pub fn Engine::tokenize_text(self : Engine, text : StringView) -> Tokens {
  ds4_engine_tokenize_text(self, @utf8.encode(text))
}

///|
#borrow(engine, text)
extern "c" fn ds4_engine_tokenize_rendered_chat(
  engine : Engine,
  text : Bytes,
) -> Tokens = "moonbit_ds4_engine_tokenize_rendered_chat"

///|
pub fn Engine::tokenize_rendered_chat(
  self : Engine,
  text : StringView,
) -> Tokens {
  ds4_engine_tokenize_rendered_chat(self, @utf8.encode(text))
}

///|
#borrow(engine, system, prompt)
extern "c" fn ds4_engine_encode_chat_prompt(
  engine : Engine,
  system : Bytes,
  prompt : Bytes,
  think_mode : Int,
) -> Tokens = "moonbit_ds4_engine_encode_chat_prompt"

///|
pub fn Engine::encode_chat_prompt(
  self : Engine,
  system? : StringView = "",
  prompt : StringView,
  think_mode? : ThinkMode = ThinkMode::none(),
) -> Tokens {
  ds4_engine_encode_chat_prompt(
    self,
    @utf8.encode(system),
    @utf8.encode(prompt),
    think_mode.to_int(),
  )
}

///|
#borrow(engine, tokens)
extern "c" fn ds4_engine_chat_begin(engine : Engine, tokens : Tokens) -> Unit = "moonbit_ds4_engine_chat_begin"

///|
pub fn Engine::chat_begin(self : Engine, tokens : Tokens) -> Unit {
  ds4_engine_chat_begin(self, tokens)
}

///|
#borrow(engine, tokens)
extern "c" fn ds4_engine_chat_append_max_effort_prefix(
  engine : Engine,
  tokens : Tokens,
) -> Unit = "moonbit_ds4_engine_chat_append_max_effort_prefix"

///|
pub fn Engine::chat_append_max_effort_prefix(
  self : Engine,
  tokens : Tokens,
) -> Unit {
  ds4_engine_chat_append_max_effort_prefix(self, tokens)
}

///|
#borrow(engine, tokens, role, content)
extern "c" fn ds4_engine_chat_append_message(
  engine : Engine,
  tokens : Tokens,
  role : Bytes,
  content : Bytes,
) -> Unit = "moonbit_ds4_engine_chat_append_message"

///|
pub fn Engine::chat_append_message(
  self : Engine,
  tokens : Tokens,
  role? : StringView = "user",
  content : StringView,
) -> Unit {
  ds4_engine_chat_append_message(
    self,
    tokens,
    @utf8.encode(role),
    @utf8.encode(content),
  )
}

///|
#borrow(engine, tokens)
extern "c" fn ds4_engine_chat_append_assistant_prefix(
  engine : Engine,
  tokens : Tokens,
  think_mode : Int,
) -> Unit = "moonbit_ds4_engine_chat_append_assistant_prefix"

///|
pub fn Engine::chat_append_assistant_prefix(
  self : Engine,
  tokens : Tokens,
  think_mode? : ThinkMode = ThinkMode::none(),
) -> Unit {
  ds4_engine_chat_append_assistant_prefix(self, tokens, think_mode.to_int())
}

///|
extern "c" fn ds4_session_new() -> Session = "moonbit_ds4_session_new"

///|
const DS4_ERROR_BUFFER_SIZE : Int = 1024

///|
fn new_error_buffer() -> FixedArray[Byte] {
  FixedArray::make(DS4_ERROR_BUFFER_SIZE, b'\x00')
}

///|
fn decode_error_buffer(err : FixedArray[Byte]) -> String {
  let bytes = err.unsafe_reinterpret_as_bytes()
  let len = bytes.find(b"\x00").unwrap_or(bytes.length())
  @utf8.decode_lossy(bytes.view(end=len))
}

///|
pub fn Engine::create_session(
  self : Engine,
  ctx_size : Int,
) -> Session raise Ds4Error {
  let session = ds4_session_new()
  let status = ds4_session_create(session, self, ctx_size)
  if status != 0 {
    raise Ds4Error::SessionCreateFailed(status, "")
  }
  session
}

///|
#borrow(session, engine)
extern "c" fn ds4_session_create(
  session : Session,
  engine : Engine,
  ctx_size : Int,
) -> Int = "moonbit_ds4_session_create"

///|
#borrow(session, prompt, err)
extern "c" fn ds4_session_sync(
  session : Session,
  prompt : Tokens,
  err : FixedArray[Byte],
  errlen : Int,
) -> Int = "moonbit_ds4_session_sync"

///|
pub fn Session::sync(self : Session, prompt : Tokens) -> Unit raise Ds4Error {
  let err = new_error_buffer()
  let status = ds4_session_sync(self, prompt, err, err.length())
  if status != 0 {
    raise Ds4Error::OperationFailed(
      op="sync",
      status~,
      message=decode_error_buffer(err),
    )
  }
}

///|
#borrow(session, err)
extern "c" fn ds4_session_eval(
  session : Session,
  token : Int,
  err : FixedArray[Byte],
  errlen : Int,
) -> Int = "moonbit_ds4_session_eval"

///|
pub fn Session::eval(self : Session, token : Int) -> Unit raise Ds4Error {
  let err = new_error_buffer()
  let status = ds4_session_eval(self, token, err, err.length())
  if status != 0 {
    raise Ds4Error::OperationFailed(
      op="eval",
      status~,
      message=decode_error_buffer(err),
    )
  }
}

///|
#borrow(session)
extern "c" fn ds4_session_argmax(session : Session) -> Int = "moonbit_ds4_session_argmax"

///|
pub fn Session::argmax(self : Session) -> Int {
  ds4_session_argmax(self)
}

///|
#borrow(session)
extern "c" fn ds4_session_argmax_excluding(
  session : Session,
  excluded_id : Int,
) -> Int = "moonbit_ds4_session_argmax_excluding"

///|
pub fn Session::argmax_excluding(self : Session, excluded_id : Int) -> Int {
  ds4_session_argmax_excluding(self, excluded_id)
}

///|
#borrow(session, sampler)
extern "c" fn ds4_session_sample(
  session : Session,
  temperature : Float,
  top_k : Int,
  top_p : Float,
  min_p : Float,
  sampler : Sampler,
) -> Int = "moonbit_ds4_session_sample"

///|
pub fn Session::sample(
  self : Session,
  sampler : Sampler,
  temperature? : Float = 1.0,
  top_k? : Int = 0,
  top_p? : Float = 1.0,
  min_p? : Float = 0.0,
) -> Int {
  ds4_session_sample(self, temperature, top_k, top_p, min_p, sampler)
}

///|
#borrow(session)
extern "c" fn ds4_session_invalidate(session : Session) -> Unit = "moonbit_ds4_session_invalidate"

///|
pub fn Session::invalidate(self : Session) -> Unit {
  ds4_session_invalidate(self)
}

///|
#borrow(session)
extern "c" fn ds4_session_rewind(session : Session, pos : Int) -> Unit = "moonbit_ds4_session_rewind"

///|
pub fn Session::rewind(self : Session, pos : Int) -> Unit {
  ds4_session_rewind(self, pos)
}

///|
#borrow(session)
extern "c" fn ds4_session_pos(session : Session) -> Int = "moonbit_ds4_session_pos"

///|
pub fn Session::pos(self : Session) -> Int {
  ds4_session_pos(self)
}

///|
#borrow(session)
extern "c" fn ds4_session_ctx(session : Session) -> Int = "moonbit_ds4_session_ctx"

///|
pub fn Session::ctx(self : Session) -> Int {
  ds4_session_ctx(self)
}

///|
#borrow(session)
extern "c" fn ds4_session_tokens(session : Session) -> Tokens = "moonbit_ds4_session_tokens"

///|
pub fn Session::tokens(self : Session) -> Tokens {
  ds4_session_tokens(self)
}