//! A turn's speaker and its number (M5c spec, sections 5 and 11). `loopd` logs the speaker on the
//! turn's `User` record as it came, and never puts it in a request: the model sees the same bytes
//! with a speaker and without one. After a finished turn it tells `memoryd` the turn's number: the
//! count of `User` records in the session's log, so the first turn is 1, a failed turn that was
//! logged counts too, and the count goes on after a reopen or a restart. Do not edit.

mod support;
#[path = "support/tags.rs"]
mod tags;

use std::os::unix::net::{UnixListener, UnixStream};
use std::path::{Path, PathBuf};
use std::sync::mpsc;
use std::sync::{Arc, Mutex};
use std::thread;
use std::time::Duration;

use bxcore::{
    Empty, Envelope, LogRecord, Message, PROTOCOL_VERSION, SessionId, SessionUpdated, Timestamp,
    Turn, read_frame, write_frame,
};
use loopd::baseline::Baseline;
use loopd::channel::{Context, serve};
use loopd::llama::Client;
use loopd::memory::{MemoryPort, MemorySource};
use loopd::session::Session;
use loopd::tools::Registry;
use support::{FakeServer, Home, Reply, ScriptedPort};
use tags::{holds_a_tag, same};

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

/// A `memoryd` stand-in that records each `(session, turn)` it is told about.
#[derive(Default)]
struct Recorder(Mutex<Vec<(String, u64)>>);

impl Recorder {
    fn told(&self) -> Vec<(String, u64)> {
        self.0.lock().unwrap().clone()
    }
}

/// Lets the test keep its `Arc<Recorder>` while the context owns a `Box<dyn MemorySource>`.
struct Shared(Arc<Recorder>);

impl MemorySource for Shared {
    fn block(&self) -> Result<String, String> {
        Ok(String::new())
    }
    fn updated(&self, session: &SessionId, turn: u64) -> Result<(), String> {
        let entry = (session.as_str().to_string(), turn);
        self.0.0.lock().unwrap().push(entry);
        Ok(())
    }
}

/// A fake inference server and a home that one `loopd` or several, one after the other, serve
/// from.
struct Setup {
    server: FakeServer,
    // Declared last so it is dropped last: fields drop in order, and the home must outlive
    // everything that may still write into it (tip T44).
    home: Home,
}

impl Setup {
    /// The server answers the completions in `chat` in order; the last one repeats.
    fn new(chat: Vec<Reply>) -> Setup {
        let home = Home::new();
        let server = FakeServer::start();
        server.route(CHAT, chat);
        Setup { server, home }
    }

    /// Starts a channel server on the home, as a new `loopd` process would: its own context,
    /// nothing carried over but the files. `name` names its socket, under the home; the returned
    /// recorder hears what this `loopd` tells `memoryd`.
    fn loopd(&self, name: &str) -> (PathBuf, Arc<Recorder>) {
        let mut cfg = self.home.config(&self.server.socket);
        cfg.channel.socket = self.home.dir.join(name);
        let socket = cfg.channel_socket();
        let listener = UnixListener::bind(&socket).unwrap();
        let memory = Arc::new(Recorder::default());
        let ctx = Arc::new(Context::new(
            cfg.clone(),
            Client::new(cfg),
            Box::new(ScriptedPort::new(vec![])),
            Registry::m2b(),
            Box::new(Shared(Arc::clone(&memory))),
        ));
        thread::spawn(move || serve(listener, ctx));
        (socket, memory)
    }

    /// The `User` records of a session's log, in order.
    fn users(&self, session: &str) -> Vec<LogRecord> {
        let records = self.home.records(session);
        let is_user = |r: &LogRecord| matches!(r, LogRecord::User { .. });
        records.into_iter().filter(is_user).collect()
    }

    /// The bodies of the completion requests so far, as sent, in order.
    fn bodies(&self) -> Vec<String> {
        let requests = self.server.requests_to(CHAT);
        let body = |r: &support::Recorded| String::from_utf8(r.body.clone()).unwrap();
        requests.iter().map(body).collect()
    }
}

/// Sends one turn and returns true if it ended with `turn_done`, false if with an error.
fn turn(socket: &Path, session: &str, resume: bool, speaker: Option<&str>) -> bool {
    let mut stream = UnixStream::connect(socket).unwrap();
    stream
        .set_read_timeout(Some(Duration::from_secs(10)))
        .unwrap();
    let msg = Message::Turn(Turn {
        session: SessionId::new(session).unwrap(),
        content: "hello".to_string(),
        resume,
        speaker: speaker.map(str::to_string),
    });
    let env = Envelope {
        v: PROTOCOL_VERSION,
        id: 7,
        r#final: true,
        msg,
    };
    write_frame(&mut stream, &env).unwrap();
    loop {
        let frame = read_frame(&mut stream).unwrap();
        if !frame.r#final {
            continue;
        }
        match frame.msg {
            Message::TurnDone(_) => return true,
            Message::Error(_) => return false,
            other => panic!("the turn ended with neither turn_done nor error: {other:?}"),
        }
    }
}

fn told(session: &str, turn: u64) -> (String, u64) {
    (session.to_string(), turn)
}

/// The speaker of a `User` record, or a panic for any other record.
fn speaker(record: &LogRecord) -> Option<String> {
    match record {
        LogRecord::User { speaker, .. } => speaker.clone(),
        other => panic!("not a user record: {other:?}"),
    }
}

#[test]
fn the_speaker_is_logged_on_the_user_record_as_it_came() {
    let s = Setup::new(vec![Reply::fixture("plain")]);
    let (socket, _) = s.loopd("loop.sock");
    // Not a valid Honcho peer name, and one that holds one of the model's tags: `loopd` neither
    // checks nor neutralises a speaker, since it never reaches the model.
    let tagged = format!("x{}y", bxcore::sanitize::TAGS[0].text);
    assert!(turn(&socket, "a", false, Some("first.last")));
    assert!(turn(&socket, "a", true, None));
    assert!(turn(&socket, "a", true, Some(&tagged)));
    let users = s.users("a");
    assert_eq!(users.len(), 3, "one user record per turn");
    assert_eq!(speaker(&users[0]), Some("first.last".to_string()));
    assert_eq!(speaker(&users[1]), None, "a turn without a speaker");
    same(&speaker(&users[2]), &Some(tagged), "the third speaker");
    let log = s.home.read("sessions/a/0.jsonl");
    assert_eq!(
        log.matches("\"speaker\":").count(),
        2,
        "no speaker key for the turn without one"
    );
}

#[test]
fn the_requests_are_the_same_with_and_without_a_speaker() {
    let s = Setup::new(vec![Reply::fixture("plain")]);
    let (socket, _) = s.loopd("loop.sock");
    let tagged = format!("x{}y", bxcore::sanitize::TAGS[0].text);
    assert!(turn(&socket, "with", false, Some("Zelda_Q")));
    assert!(turn(&socket, "with", true, Some(&tagged)));
    assert!(turn(&socket, "without", false, None));
    assert!(turn(&socket, "without", true, None));
    let bodies = s.bodies();
    assert_eq!(bodies.len(), 4, "one request per turn");
    // Each session's second request replays its first turn from the log: a speaker there would
    // change the replayed bytes too.
    assert!(bodies[0] == bodies[2], "the first requests differ");
    assert!(bodies[1] == bodies[3], "the second requests differ");
    for (n, body) in bodies.iter().enumerate() {
        assert!(!body.contains("Zelda_Q"), "request {n} holds the speaker");
        assert!(!holds_a_tag(body), "request {n} holds the tagged speaker");
    }
}

#[test]
fn the_turn_number_starts_at_one_in_each_session() {
    let s = Setup::new(vec![Reply::fixture("plain")]);
    let (socket, memory) = s.loopd("loop.sock");
    assert!(turn(&socket, "a", false, None));
    assert!(turn(&socket, "a", true, Some("kyle")));
    assert!(turn(&socket, "b", false, None));
    assert!(turn(&socket, "a", true, None));
    assert_eq!(
        memory.told(),
        [told("a", 1), told("a", 2), told("b", 1), told("a", 3)]
    );
}

#[test]
fn a_turn_with_tool_calls_is_one_turn() {
    // The first turn asks for a tool, then answers: two assistant records, one user record.
    let s = Setup::new(vec![Reply::fixture("tool_call"), Reply::fixture("plain")]);
    let (socket, memory) = s.loopd("loop.sock");
    assert!(turn(&socket, "a", false, None));
    assert!(turn(&socket, "a", true, None));
    assert_eq!(memory.told(), [told("a", 1), told("a", 2)]);
}

#[test]
fn a_failed_turn_is_not_told_but_is_counted() {
    let s = Setup::new(vec![
        Reply::fixture("plain"),
        Reply::fixture("bad_request"),
        Reply::fixture("plain"),
    ]);
    let (socket, memory) = s.loopd("loop.sock");
    assert!(turn(&socket, "a", false, None));
    assert!(!turn(&socket, "a", true, None), "the second turn fails");
    assert!(!turn(&socket, "nobody", true, None), "no such session");
    assert!(turn(&socket, "a", true, None));
    assert_eq!(s.users("a").len(), 3, "the failed turn's message is logged");
    assert_eq!(memory.told(), [told("a", 1), told("a", 3)]);
}

#[test]
fn a_restarted_loopd_goes_on_counting_from_the_log() {
    let s = Setup::new(vec![Reply::fixture("plain")]);
    let (first, before) = s.loopd("first.sock");
    assert!(turn(&first, "a", false, None));
    assert!(turn(&first, "a", true, None));
    let (second, after) = s.loopd("second.sock");
    assert!(turn(&second, "a", true, None));
    assert!(turn(&second, "b", false, None));
    assert_eq!(before.told(), [told("a", 1), told("a", 2)]);
    assert_eq!(after.told(), [told("a", 3), told("b", 1)]);
}

#[test]
fn session_turns_counts_the_user_records_of_the_log() {
    let home = Home::new();
    let cfg = home.config(Path::new("/tmp/unused.sock"));
    let baseline = Baseline::assemble(&cfg, &Registry::m2b()).unwrap();
    let id = SessionId::new("a").unwrap();
    let mut session = Session::create(&home.dir, id.clone(), baseline, 0).unwrap();
    assert_eq!(session.turns(), 0, "a new session has had no turn");
    let user = |speaker: Option<&str>| LogRecord::User {
        time: Timestamp::now(),
        content: "hello".to_string(),
        speaker: speaker.map(str::to_string),
    };
    let assistant = || LogRecord::Assistant {
        time: Timestamp::now(),
        content: Some("hi".to_string()),
        reasoning_content: None,
        tool_calls: vec![],
    };
    session.append(user(None)).unwrap();
    session.append(assistant()).unwrap();
    assert_eq!(session.turns(), 1);
    session.append(user(Some("kyle"))).unwrap();
    session.append(assistant()).unwrap();
    session.append(assistant()).unwrap();
    assert_eq!(session.turns(), 2);
    drop(session);
    let session = Session::open(&home.dir, id).unwrap();
    assert_eq!(session.turns(), 2, "counted again from the log");
}

#[test]
fn the_memory_port_sends_the_turn_number() {
    let home = Home::new();
    let socket = home.dir.join("memory.sock");
    let listener = UnixListener::bind(&socket).unwrap();
    let (tx, rx) = mpsc::channel();
    thread::spawn(move || {
        let (mut stream, _) = listener.accept().unwrap();
        let env = read_frame(&mut stream).unwrap();
        let reply = Envelope {
            v: PROTOCOL_VERSION,
            id: env.id,
            r#final: true,
            msg: Message::Ok(Empty {}),
        };
        write_frame(&mut stream, &reply).unwrap();
        let _ = tx.send(env.msg);
    });
    let port = MemoryPort::new(socket, Duration::from_secs(5));
    let id = SessionId::new("a").unwrap();
    assert_eq!(port.updated(&id, 7), Ok(()));
    let sent = rx.recv_timeout(Duration::from_secs(10)).unwrap();
    let want = Message::SessionUpdated(SessionUpdated {
        session: id,
        turn: 7,
    });
    assert_eq!(sent, want);
}