///| Stateful orchestration around proposal, verification, metrics, cache-slot

///| lifecycle, and request budget. This is still model-agnostic: callers pass

///|
/// a `ReplayRoundResult` or adapt their own inference backend to that shape.
pub enum SessionError {
  BudgetFailure
  CacheFailure
  AlreadyFinished
  InvalidRound
  RoundPending
} derive(Eq, Debug)

///| Terminal reason is explicit so a server can distinguish a clean EOS-style

///|
/// finish requested by its caller from an enforced resource budget stop.
pub enum StopReason {
  Running
  Requested
  TargetCallBudget
  OutputTokenBudget
  DraftTokenBudget
  CacheCapacity
} derive(Eq, Debug)

///|
pub fn StopReason::label(self : StopReason) -> String {
  match self {
    Running => "running"
    Requested => "requested"
    TargetCallBudget => "target-call-budget"
    OutputTokenBudget => "output-token-budget"
    DraftTokenBudget => "draft-token-budget"
    CacheCapacity => "cache-capacity"
  }
}

///| Snapshot returned after each successful round. It avoids exposing mutable

///|
/// session internals to logging or UI code.
pub struct SessionSnapshot {
  generated : Array[Int]
  metrics : DecodeMetrics
  stop_reason : StopReason
}

///|
pub fn SessionSnapshot::generated(self : SessionSnapshot) -> Array[Int] {
  self.generated.copy()
}

///|
pub fn SessionSnapshot::stop_reason(self : SessionSnapshot) -> StopReason {
  self.stop_reason
}

///| A session owns cache transactions for retained tokens. Every committed

///| draft prefix gets a transaction id, which lets an embedding layer map the

///|
/// lightweight slot model to a real KV cache implementation.
pub struct DecodeSession {
  generated : Array[Int]
  metrics : DecodeMetrics
  budget : DecodeBudget
  cache : KvSlots
  mut stop_reason : StopReason
  mut pending : (Int, Int)?
}

///|
pub fn DecodeSession::new(
  prefix : Array[Int],
  budget : DecodeBudget,
  cache_capacity : Int,
) -> DecodeSession {
  let generated : Array[Int] = []
  for token in prefix {
    generated.push(token)
  }
  {
    generated,
    metrics: DecodeMetrics::empty(),
    budget,
    cache: KvSlots::new(cache_capacity),
    stop_reason: Running,
    pending: None,
  }
}

///|
pub fn DecodeSession::generated(self : DecodeSession) -> Array[Int] {
  self.generated.copy()
}

///|
pub fn DecodeSession::metrics(self : DecodeSession) -> DecodeMetrics {
  self.metrics.copy()
}

///|
pub fn DecodeSession::is_running(self : DecodeSession) -> Bool {
  self.stop_reason == Running
}

///|
pub fn DecodeSession::stop_reason(self : DecodeSession) -> StopReason {
  self.stop_reason
}

///|
fn DecodeSession::set_budget_stop(
  self : DecodeSession,
  error : BudgetError,
) -> Unit {
  self.stop_reason = match error {
    TargetCallLimit => TargetCallBudget
    OutputTokenLimit => OutputTokenBudget
    DraftTokenLimit => DraftTokenBudget
    InvalidLimit => Requested
  }
}

///| Reserve cache capacity and check policy limits before executing a model

///|
/// round. Call `apply_round` with the result only if this function succeeds.
pub fn DecodeSession::prepare_round(
  self : DecodeSession,
  depth : Int,
) -> Result[(Int, Array[Int]), SessionError] {
  if !self.is_running() {
    return Err(AlreadyFinished)
  }
  if self.pending is Some(_) {
    return Err(RoundPending)
  }
  match self.budget.allow_round(self.metrics, depth) {
    Ok(_) => ()
    Err(error) => {
      self.set_budget_stop(error)
      return Err(BudgetFailure)
    }
  }
  match self.cache.reserve(depth) {
    Ok(reservation) => {
      self.pending = Some((reservation.0, depth))
      Ok(reservation)
    }
    Err(_) => {
      self.stop_reason = CacheCapacity
      Err(CacheFailure)
    }
  }
}

///| Commit accepted cache positions, free unused speculative positions, append

///| emitted tokens, and update metrics. A replacement after rejection is not

///|
/// committed here because a target backend owns that new token's cache entry.
pub fn DecodeSession::apply_round(
  self : DecodeSession,
  transaction : Int,
  result : ReplayRoundResult,
) -> Result[SessionSnapshot, SessionError] {
  if !self.is_running() {
    return Err(AlreadyFinished)
  }
  let verification = result.verification()
  let proposal = result.proposal()
  let reserved_depth = match self.pending {
    Some((id, depth)) if id == transaction => depth
    _ => return Err(InvalidRound)
  }
  if proposal.prefix != self.generated ||
    proposal.length() != reserved_depth ||
    verification.accepted_count < 0 ||
    verification.accepted_count > proposal.length() {
    return Err(InvalidRound)
  }
  match validate_proposal(proposal) {
    Err(_) => return Err(InvalidRound)
    Ok(_) => ()
  }
  match verification.rejected_at {
    None =>
      if verification.accepted_count != proposal.length() ||
        verification.emitted.length() != proposal.length() {
        return Err(InvalidRound)
      }
    Some(index) =>
      if index != verification.accepted_count ||
        index >= proposal.length() ||
        verification.emitted.length() != index + 1 {
        return Err(InvalidRound)
      }
  }
  for i in 0..= proposal.tokens[0].distribution.length() {
      return Err(InvalidRound)
    }
    if i < verification.accepted_count && token != proposal.tokens[i].token {
      return Err(InvalidRound)
    }
  }
  match self.budget.allow_round(self.metrics, reserved_depth) {
    Err(error) => {
      self.set_budget_stop(error)
      let _ = self.cache.rollback(transaction)
      self.pending = None
      return Err(BudgetFailure)
    }
    Ok(_) => ()
  }
  match
    self.budget.allow_emission(self.metrics, verification.emitted.length()) {
    Ok(_) => ()
    Err(error) => {
      self.set_budget_stop(error)
      let _ = self.cache.rollback(transaction)
      self.pending = None
      return Err(BudgetFailure)
    }
  }
  match self.cache.commit(transaction, verification.accepted_count) {
    Ok(_) => ()
    Err(_) => return Err(CacheFailure)
  }
  self.metrics.record(verification, result.proposal())
  self.pending = None
  for token in verification.emitted {
    self.generated.push(token)
  }
  Ok({
    generated: self.generated.copy(),
    metrics: self.metrics.copy(),
    stop_reason: self.stop_reason,
  })
}

///| Release cache slots for a previously retained accepted prefix. This gives

///|
/// serving code a distinct lifecycle hook when it evicts a completed request.
pub fn DecodeSession::release_cache(
  self : DecodeSession,
  transaction : Int,
) -> Result[Unit, SessionError] {
  match self.cache.release_committed(transaction) {
    Ok(_) => Ok(())
    Err(_) => Err(CacheFailure)
  }
}

///| Mark an otherwise healthy session finished, for example when an embedding

///|
/// application observes its EOS token.
pub fn DecodeSession::finish(self : DecodeSession) -> Unit {
  if self.pending is Some((transaction, _)) {
    let _ = self.cache.rollback(transaction)
    self.pending = None
  }
  if self.stop_reason == Running {
    self.stop_reason = Requested
  }
}

///|
/// Cancel a reservation after a model provider fails, without ending the request.
pub fn DecodeSession::cancel_round(
  self : DecodeSession,
  transaction : Int,
) -> Result[Unit, SessionError] {
  match self.pending {
    Some((id, _)) if id == transaction => {
      match self.cache.rollback(id) {
        Err(_) => return Err(CacheFailure)
        Ok(_) => ()
      }
      self.pending = None
      Ok(())
    }
    _ => Err(InvalidRound)
  }
}

///|
pub fn DecodeSession::render(self : DecodeSession) -> String {
  "session=" +
  self.stop_reason.label() +
  " generated=" +
  self.generated.length().to_string() +
  " " +
  self.budget.describe_remaining(self.metrics)
}