//! Duration and response-token budgets.
//!
//! Why these live together: they are not content matchers, they are budgets.
//! A budget answers "did we stay under N units?" for some unit (milliseconds,
//! tokens). They share a failure shape: actual value, budget, overage. Putting
//! them in one module keeps that shape consistent and avoids two near-identical
//! modules.
//!
//! ## Tokenizers
//!
//! Real OpenAI-family BPE tokenizers ship via `tiktoken-rs`: `cl100k_base`
//! (GPT-3.5/4), `o200k_base` (GPT-4o family), and `gpt2` (mapped to the
//! r50k_base encoding tiktoken-rs uses for GPT-2). A `claude` tokenizer is
//! aliased to `cl100k_base` for now; a Claude-native BPE will land in a
//! follow-up. The whitespace tokenizer stays available for tests and for
//! callers who explicitly want a transport-free counter.

use serde_json::Value;

use super::{AssertionDiff, AssertionResult, MatcherKind};
use crate::{Error, Result};

/// Evaluate the `max_duration_ms` budget.
///
/// Returns a passing [`AssertionResult`] when `actual_ms <= budget_ms`. When
/// the call was served from cache (`cache_hit = true`) the check is skipped
/// and the result is a pass with an explanatory note. Failures carry a
/// diff that reports the actual duration, the budget, and the overage in
/// milliseconds.
pub fn evaluate_duration(actual_ms: u64, budget_ms: u64, cache_hit: bool) -> AssertionResult {
    if cache_hit {
        return AssertionResult::passed_with_note(
            MatcherKind::Duration,
            "duration check skipped (cache hit)",
        );
    }
    if actual_ms <= budget_ms {
        return AssertionResult::passed(MatcherKind::Duration);
    }
    let over = actual_ms - budget_ms;
    AssertionResult::failed(
        MatcherKind::Duration,
        AssertionDiff {
            path: String::new(),
            expected: Some(Value::from(budget_ms)),
            actual: Some(Value::from(actual_ms)),
            note: Some(format!("{actual_ms} ms ({over} ms over budget)")),
        },
    )
}

/// Mode for the response-token budget.
///
/// `Fail` is the default and produces a failing [`AssertionResult`] when the
/// budget is exceeded. `Warn` returns a passing result with a note so the
/// reporter can surface "you are close to the cap" without turning the suite
/// red.
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum BudgetMode {
    /// Exceeding the budget produces a failing assertion.
    Fail,
    /// Exceeding the budget produces a passing assertion with a warning note.
    Warn,
}

/// Token budget configuration for the `max_response_tokens` matcher.
///
/// Captures the long-form fields: budget, mode, and the per-image cost. The
/// short form `max_response_tokens: N` lowers into this struct with
/// `mode: Fail` and `image_cost: 1000`.
#[derive(Debug, Clone)]
pub struct TokenBudget {
    /// Maximum allowed token count for the response.
    pub budget: usize,
    /// How the matcher behaves when the budget is exceeded.
    pub mode: BudgetMode,
    /// Token cost charged for each image content block. Defaults to 1000 in
    /// the loader; surfaced here so callers can tune for cheap-image models.
    pub image_cost: usize,
}

impl Default for TokenBudget {
    fn default() -> Self {
        Self {
            budget: 0,
            mode: BudgetMode::Fail,
            image_cost: 1000,
        }
    }
}

/// Token counter interface.
///
/// Why a trait and not a closure: tokenizers are loaded by name from YAML
/// and may have heavy initialization (vocabulary files). A trait lets us hold
/// an instance and reuse it across many calls without re-parsing config.
pub trait Tokenizer {
    /// Return an approximate token count for `text`.
    fn count(&self, text: &str) -> usize;

    /// Return the tokenizer's stable identifier (for example `"whitespace"`).
    ///
    /// Reports surface this name verbatim so callers do not have to maintain a
    /// parallel mapping from impl type to label.
    fn name(&self) -> &'static str;
}

/// Whitespace tokenizer: splits on Unicode whitespace.
///
/// Kept as a fallback for callers who explicitly want a transport-free
/// counter. Production usage should prefer one of the BPE tokenizers
/// below because they match what models actually bill.
pub struct WhitespaceTokenizer;

impl Tokenizer for WhitespaceTokenizer {
    fn count(&self, text: &str) -> usize {
        text.split_whitespace().count()
    }

    fn name(&self) -> &'static str {
        "whitespace"
    }
}

/// BPE tokenizer backed by `tiktoken-rs`.
///
/// Why one struct for every BPE variant: the OpenAI-family encodings differ
/// only in which BPE table they load and the user-facing name they report.
/// The counting logic (`encode_with_special_tokens`) is identical, so a single
/// struct holding the loaded `CoreBPE` plus its name avoids four byte-identical
/// `impl Tokenizer` bodies. Construct with the `cl100k_base`, `o200k_base`,
/// `gpt2`, or `claude` helpers below.
pub struct BpeTokenizer {
    bpe: tiktoken_rs::CoreBPE,
    name: &'static str,
}

impl BpeTokenizer {
    /// Build a BPE tokenizer from a loader result and a stable name.
    ///
    /// The loader is the `tiktoken-rs` constructor's return value (already
    /// invoked by the caller) so the single error-mapping site here applies to
    /// every variant. The error bound stays generic on `Display` so we do not
    /// pull `anyhow` into this crate just to name the loader's error type.
    /// Returns [`Error::Tokenizer`] when the table cannot be loaded
    /// (effectively impossible at runtime; the tables ship in the crate, but
    /// the constructor returns `Result` so we propagate).
    fn from_loader<E: std::fmt::Display>(
        loaded: std::result::Result<tiktoken_rs::CoreBPE, E>,
        name: &'static str,
    ) -> Result<Self> {
        let bpe = loaded.map_err(|e| Error::Tokenizer(format!("{name}: {e}")))?;
        Ok(Self { bpe, name })
    }

    /// CL100K tokenizer (GPT-3.5 and GPT-4 family).
    pub fn cl100k_base() -> Result<Self> {
        Self::from_loader(tiktoken_rs::cl100k_base(), "cl100k_base")
    }

    /// O200K tokenizer (GPT-4o family).
    pub fn o200k_base() -> Result<Self> {
        Self::from_loader(tiktoken_rs::o200k_base(), "o200k_base")
    }

    /// GPT-2 tokenizer.
    ///
    /// `tiktoken-rs` exposes the GPT-2 encoding under the `r50k_base()`
    /// constructor (GPT-2 and GPT-3 share that BPE), so we load it and surface
    /// the user-facing `gpt2` name to match common pricing tables.
    pub fn gpt2() -> Result<Self> {
        Self::from_loader(tiktoken_rs::r50k_base(), "gpt2")
    }

    /// Claude approximate tokenizer.
    ///
    /// `tiktoken-rs` does not ship a Claude-native BPE, so the `claude`
    /// tokenizer is aliased to the `cl100k_base` table for now; a Claude-native
    /// BPE will land in a follow-up. The over/undercount is typically within
    /// 5%.
    // TODO: swap in a Claude-native BPE when one is available.
    pub fn claude() -> Result<Self> {
        Self::from_loader(tiktoken_rs::cl100k_base(), "claude")
    }
}

impl Tokenizer for BpeTokenizer {
    fn count(&self, text: &str) -> usize {
        self.bpe.encode_with_special_tokens(text).len()
    }

    fn name(&self) -> &'static str {
        self.name
    }
}

/// Factory: build a tokenizer from its YAML name.
///
/// Returns `None` for unknown names; the caller turns that into a
/// user-facing error with surrounding context. Any constructor
/// failure (extremely rare; the BPE tables are bundled) bubbles up as
/// the inner `Result`.
pub fn tokenizer_by_name(name: &str) -> Result<Option<Box<dyn Tokenizer>>> {
    let tok: Box<dyn Tokenizer> = match name {
        "whitespace" => Box::new(WhitespaceTokenizer),
        "cl100k_base" => Box::new(BpeTokenizer::cl100k_base()?),
        "o200k_base" => Box::new(BpeTokenizer::o200k_base()?),
        "gpt2" => Box::new(BpeTokenizer::gpt2()?),
        "claude" => Box::new(BpeTokenizer::claude()?),
        _ => return Ok(None),
    };
    Ok(Some(tok))
}

/// Evaluate `max_response_tokens` against an MCP response.
///
/// `content_blocks` is the `result.content` array from an MCP tool response.
/// Each block is counted per its `type`: text blocks pass through the
/// tokenizer, image blocks cost `budget.image_cost`, resource blocks with a
/// text MIME type are tokenized, and binary resources are skipped. The
/// JSON-RPC envelope is never counted, only response content.
pub fn evaluate_response_tokens(
    content_blocks: &[Value],
    budget: &TokenBudget,
    tokenizer: &dyn Tokenizer,
) -> Result<AssertionResult> {
    if budget.budget == 0 {
        return Err(Error::InvalidMatcherInput(
            "max_response_tokens budget must be greater than 0".into(),
        ));
    }
    let total = count_blocks(content_blocks, budget.image_cost, tokenizer);
    if total <= budget.budget {
        return Ok(AssertionResult::passed(MatcherKind::ResponseTokens));
    }
    let over = total - budget.budget;
    let note = format!(
        "{total} tokens ({over} tokens over budget of {})",
        budget.budget
    );
    match budget.mode {
        BudgetMode::Fail => Ok(AssertionResult::failed(
            MatcherKind::ResponseTokens,
            AssertionDiff {
                path: String::new(),
                expected: Some(Value::from(budget.budget)),
                actual: Some(Value::from(total)),
                note: Some(note),
            },
        )),
        BudgetMode::Warn => Ok(AssertionResult::passed_with_note(
            MatcherKind::ResponseTokens,
            note,
        )),
    }
}

/// Walk a slice of MCP content blocks and sum their token cost.
///
/// Unknown block types contribute 1 token (so a malformed block is not
/// silently free). Resource blocks with no recognizable text fall under the
/// same rule.
fn count_blocks(blocks: &[Value], image_cost: usize, tokenizer: &dyn Tokenizer) -> usize {
    let mut total: usize = 0;
    for block in blocks {
        total = total.saturating_add(count_one_block(block, image_cost, tokenizer));
    }
    total
}

fn count_one_block(block: &Value, image_cost: usize, tokenizer: &dyn Tokenizer) -> usize {
    let Some(obj) = block.as_object() else {
        return 1;
    };
    let kind = obj.get("type").and_then(Value::as_str).unwrap_or("");
    match kind {
        "text" => obj
            .get("text")
            .and_then(Value::as_str)
            .map(|t| tokenizer.count(t))
            .unwrap_or(1),
        "image" => image_cost,
        "resource" => count_resource_block(obj, tokenizer),
        _ => 1,
    }
}

fn count_resource_block(obj: &serde_json::Map<String, Value>, tokenizer: &dyn Tokenizer) -> usize {
    let Some(resource) = obj.get("resource").and_then(Value::as_object) else {
        return 1;
    };
    let mime = resource
        .get("mimeType")
        .and_then(Value::as_str)
        .unwrap_or("");
    // Binary resources carry `blob`. Skip them so we do not double-count an
    // image counted as a resource.
    if resource.contains_key("blob") {
        return 0;
    }
    let is_text_mime = mime.starts_with("text/")
        || mime == "application/json"
        || mime.ends_with("+json")
        || mime == "application/xml"
        || mime.ends_with("+xml")
        || mime.is_empty();
    if !is_text_mime {
        return 0;
    }
    resource
        .get("text")
        .and_then(Value::as_str)
        .map(|t| tokenizer.count(t))
        .unwrap_or(1)
}

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

    #[test]
    fn duration_under_budget_passes() {
        let r = evaluate_duration(100, 200, false);
        assert!(r.pass);
        assert!(r.diffs.is_empty());
    }

    #[test]
    fn duration_over_budget_reports_overage_in_note() {
        let r = evaluate_duration(247, 100, false);
        assert!(!r.pass);
        let note = r.diffs[0].note.as_deref().unwrap_or("");
        assert!(note.contains("247 ms"));
        assert!(note.contains("147 ms over budget"));
        assert_eq!(r.diffs[0].expected.as_ref().unwrap(), &json!(100));
        assert_eq!(r.diffs[0].actual.as_ref().unwrap(), &json!(247));
    }

    #[test]
    fn cache_hit_skips_duration_check() {
        let r = evaluate_duration(9999, 1, true);
        assert!(r.pass);
        assert_eq!(
            r.diffs[0].note.as_deref(),
            Some("duration check skipped (cache hit)")
        );
    }

    #[test]
    fn duration_exactly_at_budget_passes() {
        let r = evaluate_duration(100, 100, false);
        assert!(r.pass);
    }

    #[test]
    fn token_budget_under_passes() {
        let blocks = vec![json!({"type": "text", "text": "one two three"})];
        let r = evaluate_response_tokens(
            &blocks,
            &TokenBudget {
                budget: 10,
                ..Default::default()
            },
            &WhitespaceTokenizer,
        )
        .unwrap();
        assert!(r.pass);
    }

    #[test]
    fn token_budget_exceeded_in_fail_mode_fails() {
        let blocks = vec![json!({"type": "text", "text": "one two three four five"})];
        let r = evaluate_response_tokens(
            &blocks,
            &TokenBudget {
                budget: 3,
                ..Default::default()
            },
            &WhitespaceTokenizer,
        )
        .unwrap();
        assert!(!r.pass);
        let note = r.diffs[0].note.as_deref().unwrap_or("");
        assert!(note.contains("5 tokens"));
        assert!(note.contains("over budget"));
    }

    #[test]
    fn image_block_counted_at_default_cost() {
        let blocks = vec![json!({"type": "image", "data": "..."})];
        let budget = TokenBudget {
            budget: 500,
            ..Default::default()
        };
        let r = evaluate_response_tokens(&blocks, &budget, &WhitespaceTokenizer).unwrap();
        // Default image_cost is 1000, well over budget of 500.
        assert!(!r.pass);
        let note = r.diffs[0].note.as_deref().unwrap_or("");
        assert!(note.contains("1000 tokens"));
    }

    #[test]
    fn warn_mode_passes_with_warning_note() {
        let blocks = vec![json!({"type": "text", "text": "one two three four five"})];
        let budget = TokenBudget {
            budget: 2,
            mode: BudgetMode::Warn,
            image_cost: 1000,
        };
        let r = evaluate_response_tokens(&blocks, &budget, &WhitespaceTokenizer).unwrap();
        assert!(r.pass);
        let note = r.diffs[0].note.as_deref().unwrap_or("");
        assert!(note.contains("5 tokens"));
        assert!(note.contains("over budget"));
    }

    #[test]
    fn text_resource_block_is_tokenized() {
        let blocks = vec![json!({
            "type": "resource",
            "resource": {
                "uri": "file:///a.txt",
                "mimeType": "text/plain",
                "text": "alpha beta gamma"
            }
        })];
        let budget = TokenBudget {
            budget: 2,
            ..Default::default()
        };
        let r = evaluate_response_tokens(&blocks, &budget, &WhitespaceTokenizer).unwrap();
        assert!(!r.pass);
    }

    #[test]
    fn binary_resource_block_is_skipped() {
        let blocks = vec![json!({
            "type": "resource",
            "resource": {
                "uri": "file:///a.bin",
                "mimeType": "application/octet-stream",
                "blob": "AAAA"
            }
        })];
        let budget = TokenBudget {
            budget: 1,
            ..Default::default()
        };
        let r = evaluate_response_tokens(&blocks, &budget, &WhitespaceTokenizer).unwrap();
        assert!(r.pass);
    }

    #[test]
    fn zero_budget_is_a_structural_error() {
        let blocks = vec![json!({"type": "text", "text": "x"})];
        let budget = TokenBudget {
            budget: 0,
            ..Default::default()
        };
        let err = evaluate_response_tokens(&blocks, &budget, &WhitespaceTokenizer).unwrap_err();
        assert!(matches!(err, Error::InvalidMatcherInput(_)));
    }

    #[test]
    fn whitespace_tokenizer_counts_split_items() {
        let t = WhitespaceTokenizer;
        assert_eq!(t.count("one two three"), 3);
        assert_eq!(t.count(""), 0);
        assert_eq!(t.count("  spaced  out  "), 2);
    }

    #[test]
    fn whitespace_tokenizer_name() {
        let t = WhitespaceTokenizer;
        assert_eq!(t.name(), "whitespace");
    }

    #[test]
    fn tokenizer_name_is_stable() {
        let t = WhitespaceTokenizer;
        let first = t.name();
        let second = t.name();
        assert_eq!(first, second);
        assert_eq!(first, "whitespace");
    }

    #[test]
    fn cl100k_base_reports_its_name() {
        let t = BpeTokenizer::cl100k_base().unwrap();
        assert_eq!(t.name(), "cl100k_base");
    }

    #[test]
    fn o200k_base_reports_its_name() {
        let t = BpeTokenizer::o200k_base().unwrap();
        assert_eq!(t.name(), "o200k_base");
    }

    #[test]
    fn gpt2_reports_its_name() {
        let t = BpeTokenizer::gpt2().unwrap();
        assert_eq!(t.name(), "gpt2");
    }

    #[test]
    fn claude_reports_its_name() {
        let t = BpeTokenizer::claude().unwrap();
        assert_eq!(t.name(), "claude");
    }

    #[test]
    fn cl100k_base_counts_non_zero_for_hello_world() {
        let t = BpeTokenizer::cl100k_base().unwrap();
        let n = t.count("hello world");
        assert!(n > 0);
        // GPT-3.5/4 tokenizes "hello world" as 2 tokens; keep the
        // assertion bounded but not brittle on whitespace handling.
        assert!((2..=5).contains(&n), "expected 2..=5 tokens, got {n}");
    }

    #[test]
    fn o200k_base_counts_non_zero_for_hello_world() {
        let t = BpeTokenizer::o200k_base().unwrap();
        assert!(t.count("hello world") > 0);
    }

    #[test]
    fn gpt2_counts_non_zero_for_hello_world() {
        let t = BpeTokenizer::gpt2().unwrap();
        assert!(t.count("hello world") > 0);
    }

    #[test]
    fn claude_counts_non_zero_for_hello_world() {
        let t = BpeTokenizer::claude().unwrap();
        assert!(t.count("hello world") > 0);
    }

    #[test]
    fn tokenizer_by_name_resolves_cl100k_base() {
        let t = tokenizer_by_name("cl100k_base").unwrap();
        assert!(t.is_some());
        assert_eq!(t.unwrap().name(), "cl100k_base");
    }

    #[test]
    fn tokenizer_by_name_returns_none_for_bogus() {
        let t = tokenizer_by_name("does-not-exist").unwrap();
        assert!(t.is_none());
    }

    #[test]
    fn tokenizer_by_name_resolves_all_known_names() {
        for name in ["whitespace", "cl100k_base", "o200k_base", "gpt2", "claude"] {
            let t = tokenizer_by_name(name).unwrap();
            assert!(t.is_some(), "expected {name} to resolve");
            assert_eq!(t.unwrap().name(), name);
        }
    }

    // TODO: when a Claude-native BPE lands, this
    // equivalence no longer holds. Update the test to track the
    // expected divergence rather than removing it.
    #[test]
    fn claude_and_cl100k_base_agree_today() {
        let claude = BpeTokenizer::claude().unwrap();
        let cl = BpeTokenizer::cl100k_base().unwrap();
        for sample in [
            "",
            "hello world",
            "The quick brown fox jumps over the lazy dog.",
            "{\"role\":\"user\",\"content\":\"hi\"}",
        ] {
            assert_eq!(claude.count(sample), cl.count(sample), "sample: {sample:?}");
        }
    }
}
