procyon 0.0.1

use crate::agent::{ContentPart, Message};

// 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.
pub const CONTEXT_WINDOW: usize = 128_000;

/// The context window to budget against for a model id.
///
/// Matched on substrings because provider ids carry suffixes (dates, `-latest`, an OpenRouter
/// `vendor/` prefix) that a table of exact names would miss.
pub fn context_window(model: &str) -> usize {
    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
        ("gpt-4.1", 1_047_576),
        ("gpt-4o", 128_000),
        ("o3", 200_000),
        ("o4", 200_000),
        // DeepSeek
        ("deepseek", 128_000),
        // Google
        ("gemini-2.5", 1_048_576),
        ("gemini", 1_048_576),
        // Open weights commonly served locally
        ("qwen", 128_000),
        ("llama-3.3", 128_000),
        ("mistral", 128_000),
    ];

    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;

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()
}

// 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,
}

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

    /// `reported` is input + cache read + cache write + output for the request that just
    /// completed. 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 = price_history(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 = price_history(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),
        }
    }

    #[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));
    }
}