///|
pub(all) enum ContextBudgetStatus {
  ContextFits
  ContextNeedsPruning
  ContextTooLarge
} derive(Debug, Eq, ToJson, FromJson)

///|
pub(all) struct ContextBudgetInput {
  run_id : @domain.RunId
  stage : @routine.ModelStage
  model : String
  gateway : ModelGateway
  packet : @prompt.PromptPacket
  reserved_completion_tokens : Int
  max_prompt_ratio_percent : Int
} derive(Debug, Eq, ToJson, FromJson)

///|
pub(all) struct ContextBudgetPlan {
  run_id : @domain.RunId
  stage : @routine.ModelStage
  model : String
  context_window : Int
  estimated_prompt_tokens : Int
  reserved_completion_tokens : Int
  max_prompt_tokens : Int
  available_prompt_tokens : Int
  total_estimated_tokens : Int
  prune_required : Bool
  status : ContextBudgetStatus
  summary : String
} derive(Debug, Eq, ToJson, FromJson)

///|
pub(all) struct ContextBudgetReport {
  run_id : @domain.RunId
  ok : Bool
  plans : Array[ContextBudgetPlan]
  summary : String
} derive(Debug, Eq, ToJson, FromJson)

///|
pub fn ContextBudgetStatus::label(self : ContextBudgetStatus) -> String {
  match self {
    ContextFits => "fits"
    ContextNeedsPruning => "needs_pruning"
    ContextTooLarge => "too_large"
  }
}

///|
pub fn estimate_prompt_tokens(packet : @prompt.PromptPacket) -> Int {
  let chars = packet.text().length()
  let text_tokens = (chars + 3) / 4
  let tool_tokens = packet.required_tools.length() * 4
  text_tokens + tool_tokens
}

///|
fn context_window_for_model(gateway : ModelGateway, model : String) -> Int {
  let mut window : Int? = None
  for entry in gateway.catalog {
    if entry.id == model {
      window = entry.context_window
    }
  }
  match window {
    Some(value) => value
    None => 8192
  }
}

///|
fn clamp_ratio(percent : Int) -> Int {
  if percent < 1 {
    1
  } else if percent > 100 {
    100
  } else {
    percent
  }
}

///|
fn min_int(left : Int, right : Int) -> Int {
  if left < right {
    left
  } else {
    right
  }
}

///|
pub fn context_budget_input(
  run_id : @domain.RunId,
  gateway : ModelGateway,
  packet : @prompt.PromptPacket,
  model? : String = gateway.default_model,
  reserved_completion_tokens? : Int = 1024,
  max_prompt_ratio_percent? : Int = 80,
) -> ContextBudgetInput {
  {
    run_id,
    stage: packet.stage,
    model,
    gateway,
    packet,
    reserved_completion_tokens,
    max_prompt_ratio_percent,
  }
}

///|
fn context_budget_summary(plan : ContextBudgetPlan) -> String {
  "context budget \{plan.status.label()} for \{plan.model}: estimated \{plan.estimated_prompt_tokens} prompt token(s), reserved \{plan.reserved_completion_tokens}, window \{plan.context_window}"
}

///|
pub fn prepare_context_budget_plan(
  input : ContextBudgetInput,
) -> ContextBudgetPlan {
  let context_window = context_window_for_model(input.gateway, input.model)
  let reserved = if input.reserved_completion_tokens < 0 {
    0
  } else {
    input.reserved_completion_tokens
  }
  let max_prompt_tokens = context_window *
    clamp_ratio(input.max_prompt_ratio_percent) /
    100
  let available_prompt_tokens = min_int(
    context_window - reserved,
    max_prompt_tokens,
  )
  let estimated_prompt_tokens = estimate_prompt_tokens(input.packet)
  let total_estimated_tokens = estimated_prompt_tokens + reserved
  let status = if total_estimated_tokens > context_window ||
    available_prompt_tokens < 0 {
    ContextTooLarge
  } else if estimated_prompt_tokens > available_prompt_tokens {
    ContextNeedsPruning
  } else {
    ContextFits
  }
  let plan : ContextBudgetPlan = {
    run_id: input.run_id,
    stage: input.stage,
    model: input.model,
    context_window,
    estimated_prompt_tokens,
    reserved_completion_tokens: reserved,
    max_prompt_tokens,
    available_prompt_tokens,
    total_estimated_tokens,
    prune_required: !(status is ContextFits),
    status,
    summary: "",
  }
  { ..plan, summary: context_budget_summary(plan) }
}

///|
fn report_summary(report : ContextBudgetReport) -> String {
  let blocked = report.plans.fold(init=0, fn(count, plan) {
    if plan.status is ContextTooLarge {
      count + 1
    } else {
      count
    }
  })
  let pruning = report.plans.fold(init=0, fn(count, plan) {
    if plan.status is ContextNeedsPruning {
      count + 1
    } else {
      count
    }
  })
  if blocked > 0 {
    "context budget blocks \{blocked}/\{report.plans.length()} model stage(s)"
  } else if pruning > 0 {
    "context budget recommends pruning for \{pruning}/\{report.plans.length()} model stage(s)"
  } else {
    "context budget fits \{report.plans.length()} model stage(s)"
  }
}

///|
pub fn context_budget_report(
  run_id : @domain.RunId,
  plans : Array[ContextBudgetPlan],
) -> ContextBudgetReport {
  let report : ContextBudgetReport = {
    run_id,
    ok: plans.all(fn(plan) { !(plan.status is ContextTooLarge) }),
    plans,
    summary: "",
  }
  { ..report, summary: report_summary(report) }
}

///|
pub fn context_budget_validation(
  plan : ContextBudgetPlan,
) -> @domain.ValidationReport {
  if plan.status is ContextTooLarge {
    @domain.validation_error(
      "model.context_window",
      "\{plan.summary}; shrink prompt evidence or choose a larger model",
    )
  } else {
    @domain.validation_ok()
  }
}