smith

//! Golden session corpus (`SMH-SPEC-SPEC0001`, Sessions; Replay and VCS).
//!
//! `tests/golden/` holds session files written through the real store
//! (`create_session`, then `SessionWriter::append_frame`) with fixed ids and
//! timestamps, and each one's `build_trace` as pretty JSON:
//!
//! - `linear_chat`: three user and assistant turns,
//! - `tool_compaction`: a tool call, its result, and a compaction boundary,
//! - `sibling_branches`: two answers to one root, the second one continued.
//!
//! For every `<name>.smh` the test asserts that
//!
//! - loading reports no damage and re-encoding every frame reproduces the file
//!   byte for byte,
//! - `build_trace` of the loaded session equals `<name>.trace.json` byte for
//!   byte,
//! - forking at the selected leaf keeps every record and decodes to the same
//!   content,
//! - the session has the shape its name promises,
//! - writing the corpus again from code produces the same bytes, so writer
//!   determinism is pinned too.
//!
//! Regeneration is a deliberate, versioned act and has one command:
//!
//! ```text
//! SMITH_GOLDEN=regenerate cargo test -p smith-core --test golden
//! ```
//!
//! It refuses unless `frame::VERSION_ENTRY` differs from the number recorded
//! in `golden/VERSION` (or that file is missing), then rewrites every file
//! and `VERSION`. Without regeneration, a `VERSION` that differs from
//! `frame::VERSION_ENTRY` fails the test: an entry frame version bump must
//! come with a regenerated corpus.

use smith::config::Config;
use smith::id::{BranchId, EntryId, MessageId, SessionId, ToolCallId};
use smith::message::{ContentBlock, Message, Role};
use smith_core::frame::{Frame, VERSION_ENTRY};
use smith_core::session::{COMPACTION_KIND, EntryContent, EntryFrame, Session, SessionHeader};
use smith_core::store::{create_session, load_session};
use smith_core::trace::build_trace;
use std::path::{Path, PathBuf};

/// Timestamp of the first entry of every corpus; entries tick by one second.
const EPOCH_MS: i64 = 1_767_225_600_000;

/// Environment switch and the one value that regenerates the corpus.
const SWITCH: &str = "SMITH_GOLDEN";
const REGENERATE: &str = "regenerate";

fn golden_dir() -> PathBuf {
    Path::new(env!("CARGO_MANIFEST_DIR")).join("tests/golden")
}

/// One corpus session: fixed identity and entry frames in file order.
struct Corpus {
    name: &'static str,
    /// Distinguishes the ids of one corpus from every other one.
    tag: u128,
    header: SessionHeader,
    entries: Vec<EntryFrame>,
    /// Asserts the loaded session has the shape the name promises.
    shape: fn(&Session),
}

impl Corpus {
    fn new(name: &'static str, tag: u128, shape: fn(&Session)) -> Self {
        Self {
            name,
            tag,
            header: SessionHeader {
                session_id: SessionId::from_u128(tag << 64 | 0x5E55),
                branch_id: BranchId::from_u128(tag << 64 | 0xB4A0),
            },
            entries: Vec::new(),
            shape,
        }
    }

    /// The id with serial `n` in namespace `kind` of this corpus.
    const fn bits(&self, kind: u128, n: usize) -> u128 {
        self.tag << 64 | kind << 32 | n as u128
    }

    fn last(&self) -> Option<EntryId> {
        self.entries.last().map(|entry| entry.id)
    }

    fn next_ms(&self) -> i64 {
        EPOCH_MS + 1_000 * i64::try_from(self.entries.len()).unwrap()
    }

    /// Append `content` under `parent` with the next fixed id and timestamp.
    fn push_at(&mut self, parent: Option<EntryId>, content: &EntryContent) -> EntryId {
        let id = EntryId::from_u128(self.bits(1, self.entries.len() + 1));
        self.entries.push(EntryFrame {
            id,
            timestamp_ms: self.next_ms(),
            ..EntryFrame::new(parent, content).unwrap()
        });
        id
    }

    /// Append `content` under the last entry.
    fn push(&mut self, content: &EntryContent) -> EntryId {
        self.push_at(self.last(), content)
    }

    /// A message with the next fixed message id, stamped like its entry.
    fn message(&self, role: Role, blocks: Vec<ContentBlock>) -> EntryContent {
        EntryContent::Message(Message {
            id: MessageId::from_u128(self.bits(2, self.entries.len() + 1)),
            role,
            blocks,
            timestamp_ms: self.next_ms(),
        })
    }

    fn say(&mut self, role: Role, text: &str) -> EntryId {
        let content = self.message(role, vec![ContentBlock::text(text)]);
        self.push(&content)
    }

    /// Write the corpus through the store into a scratch directory and
    /// return the file bytes: the first entry is published by
    /// `create_session`, every later one appended by its writer.
    fn write(&self) -> Vec<u8> {
        let dir =
            std::env::temp_dir().join(format!("smith_golden_{}_{}", std::process::id(), self.name));
        let _ = std::fs::remove_dir_all(&dir);
        std::fs::create_dir_all(&dir).unwrap();
        let path = dir.join(format!("{}.smh", self.name));
        let frames = |entries: &[EntryFrame]| {
            std::iter::once(Frame::Header {
                header: self.header,
            })
            .chain(entries.iter().cloned().map(|entry| Frame::Known { entry }))
            .collect::<Vec<_>>()
        };
        let (created, appended) = self.entries.split_at(1);
        {
            let seed = Session::from_frames(&frames(created)).unwrap();
            let mut opened = create_session(&path, &seed, Config::default_valid()).unwrap();
            for entry in appended {
                opened
                    .writer
                    .append_frame(&Frame::Known {
                        entry: entry.clone(),
                    })
                    .unwrap();
            }
        }
        let bytes = std::fs::read(&path).unwrap();
        std::fs::remove_dir_all(&dir).unwrap();
        bytes
    }
}

fn linear_chat() -> Corpus {
    let mut corpus = Corpus::new("linear_chat", 0xA, |session| {
        assert_eq!(session.branches().len(), 1);
        assert_eq!(session.active_branch().entries().len(), 6);
    });
    for (question, answer) in [
        ("hello", "hi, what are we building?"),
        ("list the files", "a.txt and b.txt"),
        ("thanks", "you are welcome"),
    ] {
        corpus.say(Role::User, question);
        corpus.say(Role::Assistant, answer);
    }
    corpus
}

fn tool_compaction() -> Corpus {
    let mut corpus = Corpus::new("tool_compaction", 0xB, |session| {
        let boundary = session.compaction_boundary().unwrap();
        let entries = session.active_branch().entries();
        assert!(
            entries
                .iter()
                .any(|entry| entry.id == boundary && entry.is_compaction())
        );
    });
    let call_id = ToolCallId::from_u128(corpus.bits(3, 1));
    let input = serde_json::json!({"path": "notes.txt"});
    corpus.say(Role::User, "read notes.txt");
    let request = corpus.message(
        Role::Assistant,
        vec![ContentBlock::tool_use("read", input.clone(), call_id)],
    );
    corpus.push(&request);
    corpus.push(&EntryContent::ToolCall {
        name: "read".to_string(),
        input,
        call_id,
    });
    corpus.push(&EntryContent::ToolResult {
        ok: true,
        output: "buy milk".to_string(),
        call_id,
    });
    corpus.say(Role::Assistant, "notes.txt says: buy milk");
    corpus.push(&EntryContent::Meta {
        kind: COMPACTION_KIND.to_string(),
        detail: serde_json::json!({"summary": "user read notes.txt", "entries": 5}),
    });
    corpus.say(Role::User, "and after compaction?");
    corpus.say(Role::Assistant, "still: buy milk");
    corpus
}

fn sibling_branches() -> Corpus {
    let mut corpus = Corpus::new("sibling_branches", 0xC, |session| {
        assert_eq!(session.branches().len(), 2);
        assert_eq!(session.active_branch().entries().len(), 3);
    });
    let root = corpus.say(Role::User, "name the project");
    corpus.say(Role::Assistant, "Forge");
    let retry = corpus.message(Role::Assistant, vec![ContentBlock::text("Smith")]);
    corpus.push_at(Some(root), &retry);
    corpus.say(Role::User, "go with Smith");
    corpus
}

fn corpora() -> [Corpus; 3] {
    [linear_chat(), tool_compaction(), sibling_branches()]
}

fn trace_json(session: &Session) -> String {
    let mut json = serde_json::to_string_pretty(&build_trace(session)).unwrap();
    json.push('\n');
    json
}

fn recorded_version(dir: &Path) -> Option<u16> {
    let text = std::fs::read_to_string(dir.join("VERSION")).ok()?;
    Some(text.trim().parse().unwrap())
}

/// Rewrite the corpus, refusing unless the entry frame version moved.
fn regenerate(dir: &Path) {
    let recorded = recorded_version(dir);
    assert_ne!(
        recorded,
        Some(VERSION_ENTRY),
        "refusing to regenerate: frame::VERSION_ENTRY is still {VERSION_ENTRY}; \
         the golden corpus changes only with an entry frame version bump"
    );
    std::fs::create_dir_all(dir).unwrap();
    for corpus in corpora() {
        let path = dir.join(format!("{}.smh", corpus.name));
        std::fs::write(&path, corpus.write()).unwrap();
        let session = session_of(&path, corpus.name);
        std::fs::write(
            dir.join(format!("{}.trace.json", corpus.name)),
            trace_json(&session),
        )
        .unwrap();
    }
    std::fs::write(dir.join("VERSION"), format!("{VERSION_ENTRY}\n")).unwrap();
}

/// Load the session file at `path`, asserting that it is undamaged and that
/// re-encoding its frames reproduces it exactly.
fn session_of(path: &Path, name: &str) -> Session {
    let bytes = std::fs::read(path).unwrap();
    let max = Config::default_valid().max_frame_bytes;
    let recovery = load_session(path, max).unwrap();
    assert!(recovery.damage.is_empty(), "{name}: {:?}", recovery.damage);
    let reencoded: Vec<u8> = recovery
        .frames
        .iter()
        .flat_map(|frame| frame.encode(max).unwrap())
        .collect();
    assert!(reencoded == bytes, "{name}: re-encoded frames differ");
    Session::from_frames(&recovery.frames).unwrap()
}

/// Forking at the selected leaf keeps every record and its decoded content.
fn assert_fork_at_leaf_preserves(session: &Session, name: &str) {
    let leaf = session.selected_id().unwrap();
    let fork = session.fork_at(leaf).unwrap();
    let (original, forked) = (session.active_branch(), fork.active_branch());
    assert_eq!(original.entries(), forked.entries(), "{name}: records");
    assert_eq!(fork.selected_id(), Some(leaf), "{name}: selected leaf");
    assert_eq!(
        fork.compaction_boundary(),
        session.compaction_boundary(),
        "{name}: compaction boundary"
    );
    for (left, right) in original.entries().iter().zip(forked.entries()) {
        assert_eq!(
            session.content(left).unwrap(),
            fork.content(right).unwrap(),
            "{name}: entry {}",
            left.id
        );
    }
}

#[test]
fn corpus_replays_byte_identically() {
    let dir = golden_dir();
    match std::env::var(SWITCH).as_deref() {
        Ok(REGENERATE) => regenerate(&dir),
        Ok(other) => panic!("{SWITCH}={other}: the only value is {REGENERATE}"),
        Err(_) => {}
    }
    assert_eq!(
        recorded_version(&dir),
        Some(VERSION_ENTRY),
        "golden/VERSION must equal frame::VERSION_ENTRY; after a bump run \
         `{SWITCH}={REGENERATE} cargo test -p smith-core --test golden`"
    );
    for corpus in corpora() {
        let name = corpus.name;
        let path = dir.join(format!("{name}.smh"));
        let session = session_of(&path, name);
        (corpus.shape)(&session);
        let trace = std::fs::read_to_string(dir.join(format!("{name}.trace.json"))).unwrap();
        assert!(trace_json(&session) == trace, "{name}: trace differs");
        assert_fork_at_leaf_preserves(&session, name);
        let stored = std::fs::read(&path).unwrap();
        assert!(corpus.write() == stored, "{name}: writer output differs");
    }
}