///| Resource limits for a decoding request. In a real server these limits are

///| part of admission control: speculative decoding should never turn a short

///|
/// request into an unbounded amount of draft, target, or KV-cache work.
pub enum BudgetError {
  InvalidLimit
  TargetCallLimit
  OutputTokenLimit
  DraftTokenLimit
} derive(Eq, Debug)

///| All limits are inclusive positive caps. A caller can choose a large value

///|
/// for a dimension it does not need to constrain.
pub struct DecodeBudget {
  max_target_batches : Int
  max_output_tokens : Int
  max_draft_tokens : Int
}

///|
pub fn DecodeBudget::new(
  max_target_batches : Int,
  max_output_tokens : Int,
  max_draft_tokens : Int,
) -> Result[DecodeBudget, BudgetError] {
  if max_target_batches <= 0 || max_output_tokens <= 0 || max_draft_tokens <= 0 {
    return Err(InvalidLimit)
  }
  Ok({ max_target_batches, max_output_tokens, max_draft_tokens })
}

///|
/// A modest default for CLI demonstrations, not a hidden global policy.
pub fn DecodeBudget::demo() -> DecodeBudget {
  { max_target_batches: 16, max_output_tokens: 64, max_draft_tokens: 64 }
}

///|
pub fn DecodeBudget::max_target_batches(self : DecodeBudget) -> Int {
  self.max_target_batches
}

///|
pub fn DecodeBudget::max_output_tokens(self : DecodeBudget) -> Int {
  self.max_output_tokens
}

///|
pub fn DecodeBudget::max_draft_tokens(self : DecodeBudget) -> Int {
  self.max_draft_tokens
}

///| Check the cost of a proposed round before dispatching either model. A

///| rejected speculative token can emit at most one token, so output capacity

///|
/// is checked against the proposal depth as a safe upper bound.
pub fn DecodeBudget::allow_round(
  self : DecodeBudget,
  metrics : DecodeMetrics,
  proposal_depth : Int,
) -> Result[Unit, BudgetError] {
  if proposal_depth <= 0 {
    return Err(InvalidLimit)
  }
  if metrics.target_batches + 1 > self.max_target_batches {
    return Err(TargetCallLimit)
  }
  if metrics.draft_tokens + proposal_depth > self.max_draft_tokens {
    return Err(DraftTokenLimit)
  }
  if metrics.emitted_tokens + proposal_depth > self.max_output_tokens {
    return Err(OutputTokenLimit)
  }
  Ok(())
}

///| Check actual emitted work after verification. This is also useful to a

///|
/// caller that supplies a target backend capable of accepting an entire path.
pub fn DecodeBudget::allow_emission(
  self : DecodeBudget,
  metrics : DecodeMetrics,
  emitted_count : Int,
) -> Result[Unit, BudgetError] {
  if emitted_count < 0 {
    return Err(InvalidLimit)
  }
  if metrics.emitted_tokens + emitted_count > self.max_output_tokens {
    Err(OutputTokenLimit)
  } else {
    Ok(())
  }
}

///|
/// Explain the remaining headroom in plain text for a CLI or a request log.
pub fn DecodeBudget::describe_remaining(
  self : DecodeBudget,
  metrics : DecodeMetrics,
) -> String {
  "target_batches=" +
  (self.max_target_batches - metrics.target_batches).to_string() +
  " output_tokens=" +
  (self.max_output_tokens - metrics.emitted_tokens).to_string() +
  " draft_tokens=" +
  (self.max_draft_tokens - metrics.draft_tokens).to_string()
}