//! Tests for the thinking cap. Do not edit.
//!
//! The "capped" recording has 63 reasoning chunks (63 reasoning tokens) and then an answer. The
//! fake server plays it regardless of the control call, so the cap and the overrun allowance
//! decide what the client makes of it.

mod support;

use loopd::llama::{ChatEvent, ChatMessage, ChatRequest, Client, Completion, InferError};
use support::{FakeServer, Reply};

const CHAT: &str = "/v1/chat/completions";
const CONTROL: &str = "/v1/chat/completions/control";

fn run(
    server: &FakeServer,
    cap: u64,
    overrun: u64,
) -> (Result<Completion, InferError>, Vec<ChatEvent>) {
    let mut cfg = support::test_config(&server.socket);
    cfg.limits.thinking_cap = cap;
    cfg.limits.thinking_overrun = overrun;
    let req = ChatRequest {
        slot: 0,
        messages: vec![ChatMessage::User {
            content: "think".to_string(),
        }],
        tools: vec![],
        thinking: true,
    };
    let mut events = Vec::new();
    let result = Client::new(cfg).chat(&req, &mut |e| events.push(e.clone()));
    (result, events)
}

fn capped_events(events: &[ChatEvent]) -> Vec<u64> {
    events
        .iter()
        .filter_map(|e| match e {
            ChatEvent::ThinkingCapped { tokens } => Some(*tokens),
            _ => None,
        })
        .collect()
}

#[test]
fn under_the_cap_nothing_happens() {
    let server = FakeServer::start();
    server.route(CHAT, vec![Reply::fixture("capped")]);
    let (result, events) = run(&server, 64, 10);
    let done = result.unwrap();
    assert!(!done.thinking_capped);
    assert_eq!(done.reasoning_tokens, 63);
    assert!(capped_events(&events).is_empty());
    assert!(
        server.requests_to(CONTROL).is_empty(),
        "63 tokens is under a cap of 64"
    );
}

#[test]
fn at_the_cap_the_control_call_is_sent_once() {
    let server = FakeServer::start();
    server.route(CHAT, vec![Reply::fixture("capped")]);
    server.route(CONTROL, vec![Reply::fixture("control")]);
    let (result, events) = run(&server, 20, 100);
    let done = result.unwrap();
    assert!(done.thinking_capped);
    assert_eq!(
        done.reasoning_tokens, 63,
        "the count keeps running after the cap"
    );
    assert_eq!(
        done.content.as_deref(),
        support::expected("capped")["content"].as_str()
    );
    assert_eq!(
        capped_events(&events),
        vec![20],
        "one event, at the token count that hit the cap"
    );

    let control = server.requests_to(CONTROL);
    assert_eq!(
        control.len(),
        1,
        "the control call is made once, not on every later chunk"
    );
    let id = support::expected("capped")["id"].clone();
    assert_eq!(
        control[0].json(),
        serde_json::json!({"id": id, "action": "reasoning_end", "model": "test-model"})
    );

    // The event sits between the reasoning that hit the cap and what came after.
    let at = events
        .iter()
        .position(|e| matches!(e, ChatEvent::ThinkingCapped { .. }))
        .unwrap();
    let before = events[..at]
        .iter()
        .filter(|e| matches!(e, ChatEvent::Reasoning(_)))
        .count();
    assert_eq!(before, 20);
}

#[test]
fn thinking_on_past_the_allowance_is_an_overrun() {
    let server = FakeServer::start();
    server.route(CHAT, vec![Reply::fixture("capped")]);
    server.route(CONTROL, vec![Reply::fixture("control")]);
    // Capped at 20; 63 reasoning tokens is more than 20 + 30.
    let (result, events) = run(&server, 20, 30);
    assert!(
        matches!(result, Err(InferError::ThinkingOverrun)),
        "{result:?}"
    );
    assert_eq!(capped_events(&events), vec![20]);
    let reasoning = events
        .iter()
        .filter(|e| matches!(e, ChatEvent::Reasoning(_)))
        .count();
    assert_eq!(reasoning, 50, "the client stops reading at 20 + 30 tokens");
}

#[test]
fn the_allowance_is_exact() {
    let server = FakeServer::start();
    server.route(CHAT, vec![Reply::fixture("capped")]);
    server.route(CONTROL, vec![Reply::fixture("control")]);
    // 63 is not 20 + 44, so this is allowed; with 43 it would be an overrun at token 63.
    assert!(run(&server, 20, 44).0.is_ok());
    assert!(matches!(
        run(&server, 20, 43).0,
        Err(InferError::ThinkingOverrun)
    ));
}

#[test]
fn a_control_call_that_fails_is_an_overrun() {
    for reply in [
        Reply::json(200, r#"{"success":false,"message":"no such completion"}"#),
        Reply::json(500, r#"{"error":"boom"}"#),
        Reply::json(200, "not json"),
    ] {
        let server = FakeServer::start();
        server.route(CHAT, vec![Reply::fixture("capped")]);
        server.route(CONTROL, vec![reply]);
        let (result, events) = run(&server, 20, 100);
        assert!(
            matches!(result, Err(InferError::ThinkingOverrun)),
            "{result:?}"
        );
        assert!(
            capped_events(&events).is_empty(),
            "the cap did not take effect, so no event"
        );
    }
}

#[test]
fn a_completion_without_reasoning_is_never_capped() {
    let server = FakeServer::start();
    server.route(CHAT, vec![Reply::fixture("plain")]);
    let (result, _) = run(&server, 1, 0);
    assert!(!result.unwrap().thinking_capped);
    assert!(server.requests_to(CONTROL).is_empty());
}