procyon 0.1.2

use crate::agent::{ContentPart, Message, ToolDefinition};
use crate::config::Provider;

// Deliberately crude: the estimate is only ever used as a *delta* on top of an anchor taken from
// the provider's own usage numbers, so it does not need to be accurate in absolute terms.
const CHARS_PER_TOKEN: usize = 4;
const BLOCK_OVERHEAD: usize = 4;
const ROLE_OVERHEAD: usize = 4;

// Fallback for a model this build does not recognize: small enough that an unknown model
// compacts too eagerly rather than blowing a window the provider will reject. Deliberately below
// the smallest window a hosted model ships with, because the unrecognized case is dominated by
// locally served open weights, and those are routinely started with 8k-32k of context.
pub const CONTEXT_WINDOW: usize = 32_000;

// Confirmed against a live server (Ollama 0.20.4): a request through `/v1/chat/completions` —
// the only endpoint this build speaks to Ollama over — that asked for `options.num_ctx: 16384`
// still loaded a 4,096-cell KV cache and logged `truncating input prompt limit=4096`. The OpenAI
// compatibility shim does not forward that field at all, so no request over it can ever get more
// than Ollama's own built-in default, no matter what the loaded weights actually support or what
// this build asks for. That makes the ceiling a property of the transport, not of the model
// picked — the old table entries for "llama"/"qwen"/"mistral" (32k–128k) described what the
// weights could do, not what a request could ever obtain, so a long system prompt plus this
// build's full tool set (which alone were measured at 4k-7k tokens) was silently cut down with no
// error surfaced anywhere, and the model answered from whatever fragment survived.
const OLLAMA_CONTEXT_WINDOW: usize = 4_096;

/// The context window to budget against for a provider/model pair.
///
/// Takes the provider rather than trusting the model name alone: the same weights served by
/// Ollama, LM Studio, or a hosted API can have completely different *usable* windows, since that
/// depends on how the serving stack handles the request, not on the model id in it.
pub fn context_window(provider: Provider, model: &str) -> usize {
    if provider == Provider::Ollama {
        return OLLAMA_CONTEXT_WINDOW;
    }

    let model = model.to_lowercase();

    const TABLE: &[(&str, usize)] = &[
        // Anthropic
        ("claude-opus-5", 200_000),
        ("claude-sonnet-5", 200_000),
        ("claude-haiku", 200_000),
        ("claude-fable", 200_000),
        ("claude-3", 200_000),
        ("claude", 200_000),
        // OpenAI. Longest ids first: `gpt-4.1` must not be shadowed by a `gpt-4` entry.
        ("gpt-5", 400_000),
        ("gpt-4.1", 1_047_576),
        ("gpt-4o", 128_000),
        ("o1", 200_000),
        ("o3", 200_000),
        ("o4", 200_000),
        // DeepSeek
        ("deepseek", 128_000),
        // Google
        ("gemini-3", 1_048_576),
        ("gemini-2.5", 1_048_576),
        ("gemini", 1_048_576),
        // Open weights, served by something other than Ollama (LM Studio, a hosted endpoint):
        // unlike Ollama's OpenAI shim, these are assumed to honor the context length the operator
        // actually configured, so the weights' own advertised window is used as given.
        ("qwen", 32_768),
        ("llama-3.3", 128_000),
        ("llama", 32_768),
        ("mistral", 32_768),
    ];

    TABLE
        .iter()
        .find(|(needle, _)| model.contains(needle))
        .map(|(_, window)| *window)
        .unwrap_or(CONTEXT_WINDOW)
}
const THRESHOLD_RATIO: f64 = 0.8;
const RETAIN_RATIO: f64 = 0.16;

/// The part of the window a prompt may actually occupy.
///
/// A provider counts the reply against the same window as the prompt, and refuses a request whose
/// prompt plus `max_tokens` exceeds it — so budgeting against the raw window means a conversation
/// sitting inside the threshold is still rejected. The reserve is floored well under the window so
/// an absurd `max_tokens` cannot leave zero room for the conversation.
pub fn usable_window(window: usize, max_output: usize) -> usize {
    window.saturating_sub(max_output).max(window / 4)
}

pub fn threshold_tokens(window: usize) -> usize {
    (window as f64 * THRESHOLD_RATIO) as usize
}

pub fn retain_tokens(window: usize) -> usize {
    (window as f64 * RETAIN_RATIO) as usize
}

fn text_price(text: &str) -> usize {
    text.chars().count() / CHARS_PER_TOKEN + BLOCK_OVERHEAD
}

pub fn price_part(part: &ContentPart) -> usize {
    match part {
        ContentPart::Text { text } => text_price(text),
        ContentPart::ToolUse { id, name, input } => {
            text_price(id) + text_price(name) + text_price(&input.to_string())
        }
        ContentPart::ToolResult {
            tool_use_id,
            content,
        } => text_price(tool_use_id) + text_price(content),
    }
}

pub fn price_message(message: &Message) -> usize {
    message.content.iter().map(price_part).sum::<usize>() + ROLE_OVERHEAD
}

pub fn price_history(history: &[Message]) -> usize {
    history.iter().map(price_message).sum()
}

/// Prices the part of the request that is not the conversation: the system prompt and the tool
/// definitions.
///
/// Left out of the estimate for a long time, and it is not a rounding error — the system prompt is
/// rebuilt per turn with the project, contracts and accounts, and the tool block carries every
/// connected MCP server's schemas. That is tens of thousands of tokens that the threshold could
/// not see, and the anchor only hid it until the next `invalidate` — which compaction itself
/// performs, precisely when the pressure is highest.
pub fn price_envelope(system: Option<&str>, tools: &[ToolDefinition]) -> usize {
    let system = system.map(text_price).unwrap_or(0);
    let tools: usize = tools
        .iter()
        .map(|tool| {
            text_price(&tool.name)
                + text_price(&tool.description)
                + text_price(&tool.input_schema.to_string())
        })
        .sum();
    system + tools
}

// A `tool_use` block must be answered by a `tool_result` in the following user message. Cutting
// between the two produces a request the API rejects, so cuts are only allowed where every
// outstanding call has been answered.
fn pairing_delta(message: &Message) -> isize {
    message
        .content
        .iter()
        .map(|part| match part {
            ContentPart::ToolUse { .. } => 1,
            ContentPart::ToolResult { .. } => -1,
            ContentPart::Text { .. } => 0,
        })
        .sum()
}

pub fn is_balanced_cut(history: &[Message], index: usize) -> bool {
    history[..index.min(history.len())]
        .iter()
        .map(pairing_delta)
        .sum::<isize>()
        == 0
}

#[derive(Debug, PartialEq)]
pub enum CutChoice {
    /// Compact `history[..index]`, keeping the rest verbatim.
    Compact(usize),
    /// The retained window already covers the whole history, so there is nothing older to fold.
    NothingToCompact,
    /// A cut exists but none of them is balanced; truncating anyway would corrupt the transcript.
    NoSafeCut,
}

// Walks back from the newest message until `retain` tokens are held verbatim, then keeps walking
// back until the cut is balanced.
pub fn select_cut(history: &[Message], retain: usize) -> CutChoice {
    if history.is_empty() {
        return CutChoice::NothingToCompact;
    }

    let mut kept = 0usize;
    let mut cut = history.len();
    while cut > 0 && kept < retain {
        cut -= 1;
        kept += price_message(&history[cut]);
    }

    if cut == 0 {
        return CutChoice::NothingToCompact;
    }

    while cut > 0 && !is_balanced_cut(history, cut) {
        cut -= 1;
    }

    if cut == 0 {
        CutChoice::NoSafeCut
    } else {
        CutChoice::Compact(cut)
    }
}

/// Tracks how many tokens the next request will cost, anchored on real provider usage so the
/// crude estimator is only ever responsible for the tail appended since the last response.
#[derive(Debug, Default)]
pub struct Budget {
    anchor_tokens: Option<usize>,
    anchor_estimate: usize,
    envelope: usize,
}

impl Budget {
    pub fn new() -> Self {
        Self::default()
    }

    /// Records what the system prompt and tool block cost, so the threshold accounts for them.
    ///
    /// Called before every request rather than once: the system prompt is rebuilt per turn, and
    /// the tool block changes when a plugin or MCP server is registered. A change here does not
    /// invalidate the anchor — the anchor is a measurement of the whole request, and the estimate
    /// it is compared against moves with the envelope, so the delta stays honest.
    pub fn set_envelope(&mut self, tokens: usize) {
        self.envelope = tokens;
    }

    fn estimate(&self, history: &[Message]) -> usize {
        self.envelope + price_history(history)
    }

    /// `reported` is input + cache read + cache write + output for the request that just
    /// completed, and `history` must already include that response — the reported output tokens
    /// are part of the next request's prompt, so anchoring before the assistant message is pushed
    /// charged for it twice.
    ///
    /// Ignored when it is below what the estimator would have charged for the same prefix, so a
    /// bogus low number cannot mask real pressure.
    pub fn anchor(&mut self, reported: usize, history: &[Message]) {
        let estimate = self.estimate(history);
        if reported >= estimate {
            self.anchor_tokens = Some(reported);
            self.anchor_estimate = estimate;
        } else {
            self.anchor_tokens = Some(estimate);
            self.anchor_estimate = estimate;
        }
    }

    /// Dropped when the request envelope changes, since the anchor priced a different header.
    pub fn invalidate(&mut self) {
        self.anchor_tokens = None;
        self.anchor_estimate = 0;
    }

    pub fn total(&self, history: &[Message]) -> usize {
        let estimate = self.estimate(history);
        match self.anchor_tokens {
            Some(anchor) => anchor.saturating_add(estimate.saturating_sub(self.anchor_estimate)),
            None => estimate,
        }
    }

    pub fn is_over_threshold(&self, history: &[Message], window: usize) -> bool {
        self.total(history) >= threshold_tokens(window)
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use serde_json::json;

    fn user(text: &str) -> Message {
        Message::user(text)
    }

    fn assistant_text(text: &str) -> Message {
        Message::assistant(vec![ContentPart::Text {
            text: text.to_string(),
        }])
    }

    fn assistant_calls(ids: &[&str]) -> Message {
        Message::assistant(
            ids.iter()
                .map(|id| ContentPart::ToolUse {
                    id: id.to_string(),
                    name: "grep".to_string(),
                    input: json!({"pattern": "x"}),
                })
                .collect(),
        )
    }

    fn results(ids: &[&str]) -> Message {
        Message::tool_results(
            ids.iter()
                .map(|id| (id.to_string(), "ok".to_string()))
                .collect(),
        )
    }

    #[test]
    fn price_grows_with_content() {
        let small = price_message(&user("hi"));
        let big = price_message(&user(&"x".repeat(4000)));
        assert!(big > small + 900, "{} vs {}", big, small);
    }

    #[test]
    fn threshold_and_retain_match_the_reference_ratios() {
        assert_eq!(threshold_tokens(200_000), 160_000);
        assert_eq!(retain_tokens(200_000), 32_000);
        assert!(retain_tokens(200_000) < threshold_tokens(200_000));
    }

    #[test]
    fn a_cut_between_a_call_and_its_result_is_unbalanced() {
        let history = vec![
            user("find it"),
            assistant_calls(&["a"]),
            results(&["a"]),
            assistant_text("done"),
        ];
        assert!(is_balanced_cut(&history, 0));
        assert!(is_balanced_cut(&history, 1));
        assert!(
            !is_balanced_cut(&history, 2),
            "cutting between the call and its result must be rejected"
        );
        assert!(is_balanced_cut(&history, 3));
        assert!(is_balanced_cut(&history, 4));
    }

    #[test]
    fn parallel_calls_answered_in_one_message_balance_out() {
        let history = vec![
            user("do three things"),
            assistant_calls(&["a", "b", "c"]),
            results(&["a", "b", "c"]),
        ];
        assert!(
            !is_balanced_cut(&history, 2),
            "three outstanding calls must not be cuttable"
        );
        assert!(
            is_balanced_cut(&history, 3),
            "one message answering all three closes the balance"
        );
    }

    #[test]
    fn select_cut_moves_back_to_a_balanced_boundary() {
        // Retain 0 so the first candidate is the newest message, which sits mid-pair.
        let history = vec![
            user("q"),
            assistant_calls(&["a"]),
            results(&["a"]),
            assistant_calls(&["b"]),
            results(&["b"]),
        ];
        match select_cut(&history, 1) {
            CutChoice::Compact(index) => {
                assert!(
                    is_balanced_cut(&history, index),
                    "select_cut returned unbalanced index {}",
                    index
                );
            }
            other => panic!("a balanced cut exists, got {:?}", other),
        }
    }

    #[test]
    fn refuses_when_every_candidate_cut_would_orphan_a_result() {
        // Cutting at 1 would summarize the call and leave its result orphaned at the head of the
        // tail, which the API rejects; walking back lands on 0, so there is nothing safe to do.
        let history = vec![assistant_calls(&["a"]), results(&["a"])];
        assert_eq!(select_cut(&history, 1), CutChoice::NoSafeCut);
    }

    #[test]
    fn a_trailing_unanswered_call_may_stay_in_the_retained_tail() {
        // The summarized head is self-contained and the call keeps its future result beside it,
        // so this cut is valid even though a call is outstanding at the end.
        let history = vec![user("q"), assistant_calls(&["a"])];
        assert_eq!(select_cut(&history, 1), CutChoice::Compact(1));
    }

    #[test]
    fn empty_history_has_nothing_to_compact() {
        assert_eq!(select_cut(&[], 100), CutChoice::NothingToCompact);
    }

    #[test]
    fn retaining_everything_reports_nothing_to_compact() {
        let history = vec![user("a"), assistant_text("b")];
        assert_eq!(
            select_cut(&history, usize::MAX),
            CutChoice::NothingToCompact
        );
    }

    #[test]
    fn a_large_retain_keeps_recent_turns_verbatim() {
        let mut history = vec![user("start")];
        for i in 0..40 {
            history.push(assistant_text(&format!("reply {}", i)));
            history.push(user(&format!("follow up {}", i)));
        }
        let retain = price_history(&history) / 4;
        match select_cut(&history, retain) {
            CutChoice::Compact(index) => {
                let kept = price_history(&history[index..]);
                assert!(kept >= retain, "kept {} < retain {}", kept, retain);
                assert!(index > 0 && index < history.len());
            }
            other => panic!("expected a cut, got {:?}", other),
        }
    }

    fn tool_defs(count: usize) -> Vec<ToolDefinition> {
        (0..count)
            .map(|i| ToolDefinition {
                name: format!("tool_{}", i),
                description: "does a thing".repeat(20),
                input_schema: json!({"type": "object", "properties": {"path": {"type": "string"}}}),
            })
            .collect()
    }

    #[test]
    fn the_envelope_counts_the_system_prompt_and_the_tool_block() {
        let with_nothing = price_envelope(None, &[]);
        let with_system = price_envelope(Some(&"x".repeat(4000)), &[]);
        let with_tools = price_envelope(None, &tool_defs(30));

        assert_eq!(with_nothing, 0);
        assert!(with_system > 900, "got {}", with_system);
        assert!(
            with_tools > 1000,
            "30 tool schemas priced at {}",
            with_tools
        );
    }

    // The regression this exists for: a conversation of two short messages next to a large tool
    // block is nowhere near the threshold by message count alone, and was sent anyway.
    #[test]
    fn a_large_tool_block_can_push_the_total_over_the_threshold_on_its_own() {
        let history = vec![user("hi")];
        let mut budget = Budget::new();
        assert!(!budget.is_over_threshold(&history, 32_000));

        budget.set_envelope(threshold_tokens(32_000));
        assert!(
            budget.is_over_threshold(&history, 32_000),
            "an envelope over the threshold must be visible without the anchor"
        );
    }

    #[test]
    fn the_reply_is_reserved_out_of_the_window() {
        assert_eq!(usable_window(200_000, 8192), 191_808);
        assert!(
            threshold_tokens(usable_window(200_000, 64_000)) < threshold_tokens(200_000),
            "reserving output must lower the threshold"
        );
    }

    // An absurd max_tokens must not leave the conversation with no room at all, which would make
    // every turn compact immediately and still fail.
    #[test]
    fn the_reserve_never_swallows_the_whole_window() {
        assert_eq!(usable_window(200_000, 500_000), 50_000);
        assert!(usable_window(32_000, usize::MAX) > 0);
    }

    #[test]
    fn an_unrecognized_model_gets_a_conservative_window() {
        // Stated against a literal rather than against the constant itself, which would be an
        // assertion the compiler can fold away to `true`.
        assert_eq!(
            context_window(Provider::LmStudio, "some-local-thing-v2"),
            32_000
        );
    }

    #[test]
    fn newer_ids_are_recognized_rather_than_falling_back() {
        for (provider, model, window) in [
            (Provider::OpenAi, "gpt-5", 400_000),
            (Provider::OpenAi, "gpt-5-mini-2025-08-07", 400_000),
            (Provider::OpenAi, "o1-preview", 200_000),
            (Provider::OpenAiCompatible, "gemini-3-pro", 1_048_576),
            (Provider::Openrouter, "openrouter/qwen3-32b", 32_768),
        ] {
            assert_eq!(
                context_window(provider, model),
                window,
                "model was {}",
                model
            );
        }
    }

    // `gpt-4.1` has a million-token window and `gpt-4o` has 128k; a shorter needle matching first
    // would silently hand one the other's window.
    #[test]
    fn a_longer_id_is_not_shadowed_by_a_shorter_one() {
        assert_eq!(context_window(Provider::OpenAi, "gpt-4.1-mini"), 1_047_576);
        assert_eq!(context_window(Provider::LmStudio, "llama-3.3-70b"), 128_000);
    }

    // Verified against a live Ollama 0.20.4 server: `options.num_ctx` in the request body is not
    // honored over `/v1/chat/completions`, so every model served through it is capped at Ollama's
    // own default no matter how large the weights' real window is.
    #[test]
    fn ollama_is_capped_regardless_of_the_model_name() {
        assert_eq!(context_window(Provider::Ollama, "llama-3.3-70b"), 4_096);
        assert_eq!(context_window(Provider::Ollama, "qwen2.5-coder:32b"), 4_096);
    }

    #[test]
    fn without_an_anchor_the_total_is_the_estimate() {
        let history = vec![user("hello")];
        let budget = Budget::new();
        assert_eq!(budget.total(&history), price_history(&history));
    }

    #[test]
    fn the_anchor_replaces_the_estimate_and_the_tail_is_added_on_top() {
        let mut history = vec![user("hello")];
        let mut budget = Budget::new();
        budget.anchor(50_000, &history);
        assert_eq!(budget.total(&history), 50_000);

        history.push(assistant_text(&"x".repeat(4000)));
        let grown = budget.total(&history);
        assert!(
            grown > 50_000 && grown < 52_000,
            "expected anchor plus a ~1000 token tail, got {}",
            grown
        );
    }

    #[test]
    fn a_reported_usage_below_the_estimate_does_not_hide_pressure() {
        let history = vec![user(&"x".repeat(40_000))];
        let estimate = price_history(&history);
        let mut budget = Budget::new();
        budget.anchor(1, &history);
        assert_eq!(
            budget.total(&history),
            estimate,
            "an implausibly small usage number must not lower the total"
        );
    }

    #[test]
    fn invalidating_the_anchor_falls_back_to_the_estimate() {
        let history = vec![user("hello")];
        let mut budget = Budget::new();
        budget.anchor(50_000, &history);
        budget.invalidate();
        assert_eq!(budget.total(&history), price_history(&history));
    }

    #[test]
    fn threshold_fires_only_once_the_window_is_mostly_full() {
        let history = vec![user("hi")];
        let mut budget = Budget::new();

        budget.anchor(threshold_tokens(200_000) - 1, &history);
        assert!(!budget.is_over_threshold(&history, 200_000));

        budget.anchor(threshold_tokens(200_000), &history);
        assert!(budget.is_over_threshold(&history, 200_000));
    }
}