///|
/// 生成循环中可以替换的最小 native 操作集合。
///
/// 正常公开路径使用真实 raw API；白盒测试可以替换 decode 和 sample，以覆盖不易由
/// 小型模型稳定触发的状态边界，而不改变公开接口。
priv struct GenerationDriver {
  decode : (@raw.LlamaContextRef, FixedArray[Int]) -> Int
  sample : (@raw.LlamaSamplerRef, @raw.LlamaContextRef, Int) -> Int
}

///|
let native_generation_driver : GenerationDriver = {
  decode: (context_ref, tokens) => @raw.llama_decode_tokens(context_ref, tokens),
  sample: (sampler_ref, context_ref, index) => {
    @raw.llama_sampler_sample(sampler_ref, context_ref, index)
  },
}

///|
/// 在回调期间保持 SessionOwner 存活，防止 borrowed vocabulary 先于 owner 失效。
///
/// 机制说明见 [F-001. SessionOwner::with_vocab_ref](docs/notes/F001-with_vocav_ref.md)。
fn[A] SessionOwner::with_vocab_ref(
  self : SessionOwner,
  work : (@raw.LlamaVocabRef) -> A raise GenerationError,
) -> A raise GenerationError {
  let vocab_ref = self.get_vocab_ref()
  try work(vocab_ref) catch {
    error => {
      self.get_raw_ref() |> ignore
      raise error
    }
  } noraise {
    result => {
      self.get_raw_ref() |> ignore
      result
    }
  }
}

///|
fn tokenize_prompt(
  vocab_ref : @raw.LlamaVocabRef,
  prompt : StringView,
  add_special~ : Bool,
) -> FixedArray[Int] raise GenerationError {
  let text = @utf8.encode(prompt)
  let empty : FixedArray[Int] = []
  let query = @raw.llama_tokenize(vocab_ref, text, empty, add_special, true)
  guard query != -2147483648 else { raise TokenizationFailed(query) }
  guard query <= 0 else { raise TokenizationFailed(query) }
  guard query < 0 else { return empty }
  let required = -query
  let tokens = FixedArray::make(required, 0)
  let written = @raw.llama_tokenize(vocab_ref, text, tokens, add_special, true)
  guard written == required else { raise TokenizationFailed(written) }
  tokens
}

///|
fn token_to_piece(
  vocab_ref : @raw.LlamaVocabRef,
  token : Int,
) -> Bytes raise GenerationError {
  let empty : Bytes = Bytes::new(0)
  let query = @raw.llama_token_to_piece(vocab_ref, token, empty, 0, true)
  guard query != -2147483648 else { raise TokenToPieceFailed(token, query) }
  guard query <= 0 else { raise TokenToPieceFailed(token, query) }
  guard query < 0 else { return b"" }
  let required = -query
  let piece = Bytes::new(required)
  let written = @raw.llama_token_to_piece(vocab_ref, token, piece, 0, true)
  guard written == required else { raise TokenToPieceFailed(token, written) }
  piece
}

///|
fn finish_completion(
  output : Array[Byte],
  prompt_tokens : Int,
  generated_tokens : Int,
  finish_reason : FinishReason,
) -> Completion {
  {
    text: @utf8.decode_lossy(Bytes::from_array(output)),
    prompt_tokens,
    generated_tokens,
    finish_reason,
  }
}

///|
/// 为一次 `complete` 创建独立的 owned sampler。
///
/// random strategy 固定组装为 `top-k → top-p → temperature → dist`，其中
/// top-p 的 `min_keep` 固定为 1。构造中途失败时，已经转移给 chain 的 child
/// sampler 会随 chain 一起释放；成功返回的 sampler 由本次调用负责释放。
fn SamplingStrategy::make_sampler(
  self : SamplingStrategy,
) -> @raw.LlamaSamplerRef raise GenerationError {
  match self {
    Greedy => {
      let sampler = @raw.llama_sampler_init_greedy()
      guard !sampler.is_null() else { raise SamplerCreateFailed }
      sampler
    }
    Random(seed~, temperature~, top_k~, top_p~) => {
      let chain = @raw.llama_sampler_chain_init()
      guard !chain.is_null() else { raise SamplerCreateFailed }
      let top_k_sampler = @raw.llama_sampler_init_top_k(top_k)
      guard !top_k_sampler.is_null() else {
        @raw.llama_sampler_free(chain)
        raise SamplerCreateFailed
      }
      @raw.llama_sampler_chain_add(chain, top_k_sampler)
      let top_p_sampler = @raw.llama_sampler_init_top_p(
        Float::from_double(top_p),
        1UL,
      )
      guard !top_p_sampler.is_null() else {
        @raw.llama_sampler_free(chain)
        raise SamplerCreateFailed
      }
      @raw.llama_sampler_chain_add(chain, top_p_sampler)
      let temperature_sampler = @raw.llama_sampler_init_temp(
        Float::from_double(temperature),
      )
      guard !temperature_sampler.is_null() else {
        @raw.llama_sampler_free(chain)
        raise SamplerCreateFailed
      }
      @raw.llama_sampler_chain_add(chain, temperature_sampler)
      let seed = match seed {
        Some(seed) => seed
        None => @uint.MAX_VALUE
      }
      let distribution = @raw.llama_sampler_init_dist(seed)
      guard !distribution.is_null() else {
        @raw.llama_sampler_free(chain)
        raise SamplerCreateFailed
      }
      @raw.llama_sampler_chain_add(chain, distribution)
      chain
    }
  }
}

///|
fn Session::run_completion(
  self : Session,
  prompt_tokens : FixedArray[Int],
  options : GenerationOptions,
  vocab_ref : @raw.LlamaVocabRef,
  sampler_ref : @raw.LlamaSamplerRef,
  driver : GenerationDriver,
) -> Completion raise GenerationError {
  let context_ref = self.owner.get_raw_ref()
  let context_size = self.get_context_size()
  let prompt_token_count = prompt_tokens.length()
  let output : Array[Byte] = []
  for pending = prompt_tokens, generated_tokens = 0 {
    if pending.length() > 0 {
      self.owner.mark_needs_reset()
      let status = (driver.decode)(context_ref, pending)
      guard status == 0 else { raise DecodeFailed(status) }
      self.owner.advance_used_tokens(pending.length())
    }
    if options.max_tokens >= 0 && generated_tokens >= options.max_tokens {
      break finish_completion(
        output,
        prompt_token_count,
        generated_tokens,
        MaxTokens,
      )
    }
    if self.owner.get_used_tokens() >= context_size {
      break finish_completion(
        output,
        prompt_token_count,
        generated_tokens,
        ContextFull,
      )
    }
    let token = (driver.sample)(sampler_ref, context_ref, -1)
    if @raw.llama_vocab_is_eog(vocab_ref, token) {
      break finish_completion(
        output,
        prompt_token_count,
        generated_tokens,
        EndOfGeneration,
      )
    }
    let piece = token_to_piece(vocab_ref, token)
    for byte in piece {
      output.push(byte)
    }
    continue [token], generated_tokens + 1
  }
}

///|
fn Session::complete_with_driver(
  self : Session,
  prompt : StringView,
  options~ : GenerationOptions,
  driver~ : GenerationDriver,
) -> Completion raise GenerationError {
  guard self.owner.is_ready() else { raise SessionNeedsReset }
  self.owner.with_vocab_ref(vocab_ref => {
    guard !vocab_ref.is_null() else { raise VocabularyUnavailable }
    let used_tokens = self.owner.get_used_tokens()
    let prompt_tokens = tokenize_prompt(
      vocab_ref,
      prompt,
      add_special=used_tokens == 0U,
    )
    guard used_tokens > 0U || prompt_tokens.length() > 0 else {
      raise EmptyPrompt
    }
    let context_size = self.get_context_size()
    let available = if used_tokens >= context_size {
      0U
    } else {
      context_size - used_tokens
    }
    guard prompt_tokens.length().reinterpret_as_uint() <= available else {
      raise PromptExceedsContext(prompt_tokens.length(), available)
    }
    let sampler_ref = options.strategy.make_sampler()
    try
      self.run_completion(
        prompt_tokens, options, vocab_ref, sampler_ref, driver,
      )
    catch {
      error => {
        @raw.llama_sampler_free(sampler_ref)
        raise error
      }
    } noraise {
      completion => {
        self.owner.mark_ready()
        @raw.llama_sampler_free(sampler_ref)
        completion
      }
    }
  })
}

///|
/// 将 prompt 追加到当前 Session，并按本次 `SamplingStrategy` 同步生成文本。
///
/// prompt 只在本次调用期间读取；返回文本只包含本次新生成的内容。成功返回时，
/// prompt 和所有非 EOG 生成 token 都保留在 Session 中，供下一次调用继续使用。
/// sampler 只属于本次调用；固定 seed 会在每次调用开始时重新初始化 RNG。
/// decode 已经开始后若发生错误，后续调用会抛出 `SessionNeedsReset`，直到用户调用
/// `Session::reset` 明确丢弃无法确认的部分状态。
pub fn Session::complete(
  self : Session,
  prompt : StringView,
  options~ : GenerationOptions,
) -> Completion raise GenerationError {
  self.complete_with_driver(prompt, options~, driver=native_generation_driver)
}