//! Tests for compacting after a turn and retrying a full context. Do not edit.

mod support;

use bxcore::{Epoch, LogRecord, SessionId, TurnEvent};
use loopd::baseline::{Baseline, CARRY_HEADING};
use loopd::compact;
use loopd::llama::Client;
use loopd::memory::NoMemory;
use loopd::session::Session;
use loopd::tools::Registry;
use loopd::turn::{Runtime, TurnError, run_turn};
use support::{FakeServer, Home, Reply};

const CHAT: &str = "/v1/chat/completions";
const PROMPT: &str = "Summarise this conversation for the next epoch.\n";

fn sse(content: &str, cache_n: u64, prompt_n: u64) -> Reply {
    let content_json = serde_json::to_string(content).unwrap();
    let data1 = format!(
        r#"data: {{"choices":[{{"finish_reason":null,"index":0,"delta":{{"content":{content_json}}}}}],"id":"t","timings":{{"cache_n":{cache_n},"prompt_n":{prompt_n},"predicted_n":1}}}}"#
    );
    let data2 = format!(
        r#"data: {{"choices":[{{"finish_reason":"stop","index":0,"delta":{{}}}}],"id":"t","timings":{{"cache_n":{cache_n},"prompt_n":{prompt_n},"predicted_n":1}}}}"#
    );
    let body = format!("{data1}\n\n{data2}\n\ndata: [DONE]\n\n");
    Reply::raw(format!(
        "HTTP/1.1 200 OK\r\nContent-Type: text/event-stream\r\nContent-Length: {}\r\nConnection: close\r\n\r\n{body}",
        body.len()
    ))
}

struct World {
    home: Home,
    server: FakeServer,
    cfg: loopd::config::Config,
    client: Client,
    port: support::ScriptedPort,
    registry: Registry,
}

impl World {
    fn new() -> World {
        let home = Home::new();
        home.write("compact.md", PROMPT);
        let server = FakeServer::start();
        let mut cfg = home.config(&server.socket);
        cfg.expect.n_ctx = 1_000;
        cfg.compact.threshold = 0.85;
        let client = Client::new(cfg.clone());
        World {
            home,
            server,
            cfg,
            client,
            port: support::ScriptedPort::new(vec![]),
            registry: Registry::m2b(),
        }
    }

    fn session(&self) -> Session {
        let baseline = Baseline::assemble(&self.cfg, &self.registry).unwrap();
        Session::create(
            &self.home.dir,
            SessionId::new("a").unwrap(),
            baseline,
            self.cfg.slots.main,
        )
        .unwrap()
    }

    fn runtime(&self) -> Runtime<'_> {
        Runtime {
            cfg: &self.cfg,
            client: &self.client,
            port: &self.port,
            registry: &self.registry,
        }
    }

    fn finish(
        &self,
        session: &mut Session,
        text: &str,
        outcome: Result<loopd::turn::TurnOutcome, TurnError>,
    ) -> Result<loopd::turn::TurnOutcome, TurnError> {
        let mut events = Vec::new();
        compact::finish_turn(
            session,
            &self.runtime(),
            &NoMemory,
            text,
            None,
            &mut |e: &TurnEvent| events.push(e.clone()),
            outcome,
        )
    }
}

#[test]
fn usage_under_the_threshold_leaves_the_epoch() {
    let w = World::new();
    w.server.route(CHAT, vec![sse("ok", 400, 449)]);
    let mut s = w.session();
    let mut events = Vec::new();
    let outcome = run_turn(&mut s, &w.runtime(), "hi", None, &mut |e| {
        events.push(e.clone())
    });
    let outcome = w.finish(&mut s, "hi", outcome).unwrap();
    assert_eq!(outcome.content, "ok");
    assert_eq!(s.epoch(), Epoch(0));
    assert_eq!(w.server.requests_to(CHAT).len(), 1);
}

#[test]
fn usage_at_the_threshold_compacts_before_the_turn_returns() {
    let w = World::new();
    w.server
        .route(CHAT, vec![sse("ok", 400, 450), sse("the sum", 0, 10)]);
    let mut s = w.session();
    let mut events = Vec::new();
    let outcome = run_turn(&mut s, &w.runtime(), "hi", None, &mut |e| {
        events.push(e.clone())
    });
    w.finish(&mut s, "hi", outcome).unwrap();
    assert_eq!(s.epoch(), Epoch(1));
    assert_eq!(s.carry(), Some("the sum"));
    w.server.route(CHAT, vec![sse("next", 0, 20)]);
    let mut events = Vec::new();
    run_turn(&mut s, &w.runtime(), "again", None, &mut |e| {
        events.push(e.clone())
    })
    .unwrap();
    let sent = w.server.requests_to(CHAT);
    assert_eq!(sent.len(), 3);
    assert_eq!(sent[0].json()["id_slot"], 0);
    assert_eq!(sent[1].json()["id_slot"], 1);
    assert!(sent[1].json().get("tools").is_none());
    let live = sent[2].json();
    assert_eq!(live["id_slot"], 0);
    let contents: Vec<&str> = live["messages"]
        .as_array()
        .unwrap()
        .iter()
        .map(|m| m["content"].as_str().unwrap_or(""))
        .collect();
    assert!(
        contents
            .iter()
            .any(|c| c.contains(CARRY_HEADING) && c.contains("the sum")),
        "{contents:?}"
    );
    let old = w.home.records("a");
    assert!(matches!(
        old.last(),
        Some(LogRecord::EpochEnd { next, .. }) if *next == Epoch(1)
    ));
}

#[test]
fn session_full_compacts_and_retries_the_user_message_once() {
    let w = World::new();
    w.server.route(
        CHAT,
        vec![
            Reply::fixture("context_full"),
            sse("the sum", 0, 10),
            sse("retried", 0, 20),
        ],
    );
    let mut s = w.session();
    let mut events = Vec::new();
    let outcome = run_turn(&mut s, &w.runtime(), "please fit", None, &mut |e| {
        events.push(e.clone())
    });
    assert!(matches!(outcome, Err(TurnError::SessionFull)));
    let done = w.finish(&mut s, "please fit", outcome).unwrap();
    assert_eq!(done.content, "retried");
    assert_eq!(s.epoch(), Epoch(1));
    let users: Vec<_> = s
        .records()
        .iter()
        .filter_map(|r| match r {
            LogRecord::User { content, .. } => Some(content.as_str()),
            _ => None,
        })
        .collect();
    assert_eq!(users, ["please fit"]);
    let old_users = w
        .home
        .records("a")
        .iter()
        .filter(|r| matches!(r, LogRecord::User { .. }))
        .count();
    assert_eq!(old_users, 1);
    assert_eq!(s.turns(), 2);
    let sent = w.server.requests_to(CHAT);
    assert_eq!(sent.len(), 3);
    assert_eq!(sent[1].json()["id_slot"], 1);
    assert_eq!(sent[2].json()["id_slot"], 0);
}

#[test]
fn session_full_stays_session_full_when_compact_fails() {
    let w = World::new();
    w.server.route(
        CHAT,
        vec![
            Reply::fixture("context_full"),
            Reply::fixture("bad_request"),
        ],
    );
    let mut s = w.session();
    let mut events = Vec::new();
    let outcome = run_turn(&mut s, &w.runtime(), "x", None, &mut |e| {
        events.push(e.clone())
    });
    let err = w.finish(&mut s, "x", outcome).unwrap_err();
    assert!(matches!(err, TurnError::SessionFull), "{err:?}");
    assert_eq!(s.epoch(), Epoch(0));
}

#[test]
fn a_successful_turn_stays_successful_when_compact_fails() {
    let home = Home::new();
    let server = FakeServer::start();
    let mut cfg = home.config(&server.socket);
    cfg.expect.n_ctx = 1_000;
    cfg.compact.threshold = 0.85;
    let client = Client::new(cfg.clone());
    server.route(CHAT, vec![sse("ok", 400, 450)]);
    let baseline = Baseline::assemble(&cfg, &Registry::m2b()).unwrap();
    let mut s = Session::create(
        &home.dir,
        SessionId::new("a").unwrap(),
        baseline,
        cfg.slots.main,
    )
    .unwrap();
    let port = support::ScriptedPort::new(vec![]);
    let registry = Registry::m2b();
    let rt = Runtime {
        cfg: &cfg,
        client: &client,
        port: &port,
        registry: &registry,
    };
    let mut events = Vec::new();
    let outcome = run_turn(&mut s, &rt, "hi", None, &mut |e| events.push(e.clone()));
    let done = compact::finish_turn(
        &mut s,
        &rt,
        &NoMemory,
        "hi",
        None,
        &mut |e| events.push(e.clone()),
        outcome,
    )
    .unwrap();
    assert_eq!(done.content, "ok");
    assert_eq!(s.epoch(), Epoch(0), "compact.md is missing, so no epoch");
}

#[test]
fn finish_turn_does_not_compact_an_error_other_than_session_full() {
    let w = World::new();
    w.server.route(CHAT, vec![Reply::fixture("bad_request")]);
    let mut s = w.session();
    let mut events = Vec::new();
    let outcome = run_turn(&mut s, &w.runtime(), "x", None, &mut |e| {
        events.push(e.clone())
    });
    let err = w.finish(&mut s, "x", outcome).unwrap_err();
    assert!(matches!(err, TurnError::Infer(_)), "{err:?}");
    assert_eq!(s.epoch(), Epoch(0));
    assert_eq!(w.server.requests_to(CHAT).len(), 1);
}