smith

//! Built binary attached to active writer via a native PTY (`SMH-SPEC-SPEC0001`, Observer mode).
use assert_cmd::cargo::cargo_bin;
use portable_pty::{Child as PtyChild, CommandBuilder, MasterPty, PtySize, native_pty_system};
use serde_json::Value;
use std::io::{BufRead, BufReader, Read, Write};
use std::path::Path;
use std::process::{Child, Command, Stdio};
use std::sync::mpsc::{self, Receiver};
use std::time::{Duration, Instant};

const DEADLINE: Duration = Duration::from_secs(15);
const PTY_SIZE: PtySize = PtySize {
    rows: 24,
    cols: 120,
    pixel_width: 0,
    pixel_height: 0,
};

const OUTPUT_LIMIT: usize = 65_536;

fn completion_markers(lines: &Receiver<String>, output: &mut String, markers: &[&str]) {
    assert!(!markers.is_empty());
    let mut found = vec![false; markers.len()];
    let end = Instant::now() + DEADLINE;
    loop {
        let mut last_completion = None;
        for (index, marker) in markers.iter().enumerate() {
            if !found[index]
                && let Some(at) = output.find(marker)
            {
                found[index] = true;
                last_completion = Some(last_completion.unwrap_or(0).max(at + marker.len()));
            }
        }
        if found.iter().all(|found| *found) {
            output.drain(..last_completion.unwrap());
            return;
        }
        let chunk = lines
            .recv_timeout(end.saturating_duration_since(Instant::now()))
            .unwrap_or_else(|e| {
                panic!(
                    "waiting for {markers:?}: {e}; tail={}",
                    output
                        .chars()
                        .rev()
                        .take(1400)
                        .collect::<String>()
                        .chars()
                        .rev()
                        .collect::<String>()
                )
            });
        output.push_str(&chunk);
        if output.len() > OUTPUT_LIMIT {
            let suffix = markers.iter().map(|marker| marker.len()).max().unwrap() - 1;
            let keep_from = output.len().saturating_sub(suffix);
            let at = output
                .char_indices()
                .find_map(|(at, _)| (at >= keep_from).then_some(at))
                .unwrap_or(output.len());
            output.drain(..at);
        }
    }
}

fn marker(lines: &Receiver<String>, output: &mut String, marker: &str) {
    completion_markers(lines, output, &[marker]);
}

fn plain(bytes: &[u8], csi: &mut u8) -> String {
    let mut output = Vec::with_capacity(bytes.len());
    for &byte in bytes {
        match *csi {
            0 if byte == 0x1b => *csi = 1,
            1 if byte == b'[' => *csi = 2,
            1 => *csi = 0,
            2 if (0x40..=0x7e).contains(&byte) => *csi = 0,
            2 => {}
            _ => output.push(byte),
        }
    }
    String::from_utf8_lossy(&output).into_owned()
}

fn pty_output(mut output: Box<dyn Read + Send>) -> Receiver<String> {
    let (tx, rx) = mpsc::sync_channel(64);
    std::thread::spawn(move || {
        let mut buffer = [0; 4096];
        let mut csi = 0;
        while let Ok(read) = output.read(&mut buffer) {
            if read == 0 || tx.send(plain(&buffer[..read], &mut csi)).is_err() {
                break;
            }
        }
    });
    rx
}

struct Process(Child);

impl Drop for Process {
    fn drop(&mut self) {
        if self.0.try_wait().ok().flatten().is_none() {
            let _ = self.0.kill();
            let _ = self.0.wait();
        }
    }
}

struct Pty {
    master: Box<dyn MasterPty + Send>,
    child: Option<Box<dyn PtyChild + Send + Sync>>,
    keys: Option<Box<dyn Write + Send>>,
}

impl Pty {
    fn write_all(&mut self, bytes: &[u8]) {
        self.keys.as_mut().unwrap().write_all(bytes).unwrap();
    }

    fn close_input(&mut self) {
        self.keys.take();
    }

    fn wait(&mut self) {
        let mut child = self.child.take().unwrap();
        let mut killer = child.clone_killer();
        let (tx, rx) = mpsc::sync_channel(1);
        std::thread::spawn(move || {
            let _ = tx.send(child.wait());
        });
        let status = rx.recv_timeout(DEADLINE).unwrap_or_else(|error| {
            let _ = killer.kill();
            panic!("waiting for observer exit: {error}")
        });
        let status = status.unwrap();
        assert!(status.success(), "observer exited: {status}");
    }
}

impl Drop for Pty {
    fn drop(&mut self) {
        if let Some(child) = &mut self.child
            && child.try_wait().ok().flatten().is_none()
        {
            let _ = child.kill();
            let _ = child.wait();
        }
    }
}

#[cfg(unix)]
fn terminal_state(master: &dyn MasterPty) -> String {
    format!(
        "{:?}",
        master.get_termios().expect("PTY termios unavailable")
    )
}

fn observer(home: &Path, id: &str, export: Option<&Path>) -> (Pty, Receiver<String>) {
    let system = native_pty_system();
    let pair = system.openpty(PTY_SIZE).unwrap();
    let output = pair.master.try_clone_reader().unwrap();
    let keys = pair.master.take_writer().unwrap();
    let mut command = CommandBuilder::new(cargo_bin("smith"));
    command.args(["observe", id]);
    if let Some(export) = export {
        command.arg("--export");
        command.arg(export);
    }
    command.env("XDG_DATA_HOME", home);
    command.env("TERM", "xterm-256color");
    command.env_remove("SMITH_PROVIDER");
    command.env_remove("OPENAI_API_KEY");
    command.env_remove("ANTHROPIC_API_KEY");
    let child = pair.slave.spawn_command(command).unwrap();
    (
        Pty {
            master: pair.master,
            child: Some(child),
            keys: Some(keys),
        },
        pty_output(output),
    )
}

fn session(home: &Path, prompt: &str) -> String {
    let eval = Command::new(cargo_bin("smith"))
        .args(["eval", prompt, "--mock"])
        .env("XDG_DATA_HOME", home)
        .output()
        .unwrap();
    assert!(
        eval.status.success(),
        "{}",
        String::from_utf8_lossy(&eval.stderr)
    );
    let listed = Command::new(cargo_bin("smith"))
        .args(["sessions", "list", "--json"])
        .env("XDG_DATA_HOME", home)
        .output()
        .unwrap();
    serde_json::from_slice::<Value>(&listed.stdout).unwrap()["session"]
        .as_str()
        .unwrap()
        .to_owned()
}

fn rpc_writer(home: &Path, id: &str) -> (Process, std::process::ChildStdin, Receiver<String>) {
    let mut writer = Process(
        Command::new(cargo_bin("smith"))
            .args(["rpc", "--session", id])
            .env("XDG_DATA_HOME", home)
            .env_remove("SMITH_PROVIDER")
            .stdin(Stdio::piped())
            .stdout(Stdio::piped())
            .stderr(Stdio::piped())
            .spawn()
            .unwrap(),
    );
    let input = writer.0.stdin.take().unwrap();
    let output = writer.0.stdout.take().unwrap();
    let (tx, rx) = mpsc::sync_channel(64);
    std::thread::spawn(move || {
        for line in BufReader::new(output).lines() {
            if tx.send(line.unwrap()).is_err() {
                break;
            }
        }
    });
    (writer, input, rx)
}

fn rpc_line(lines: &Receiver<String>) -> Value {
    serde_json::from_str(&lines.recv_timeout(DEADLINE).unwrap()).unwrap()
}

#[test]
fn active_writer_observer_disables_submission_follows_exports_and_forks() {
    let home = std::env::temp_dir().join(format!("smith-observer-tui-{}", std::process::id()));
    let _ = std::fs::remove_dir_all(&home);
    let id = session(&home, "initial prompt");
    let path = home.join("smith/sessions/run").join(format!("{id}.smh"));
    let (mut writer, mut writer_in, ready_rx) = rpc_writer(&home, &id);
    writer_in
        .write_all(b"{\"id\":1,\"method\":\"ping\"}\n")
        .unwrap();
    assert_eq!(rpc_line(&ready_rx)["pong"], true);
    let before_denied = std::fs::read(&path).unwrap();
    let denied = Command::new(cargo_bin("smith"))
        .args(["eval", "must not run", "--session", &id, "--mock"])
        .env("XDG_DATA_HOME", &home)
        .output()
        .unwrap();
    assert!(!denied.status.success());
    assert!(String::from_utf8_lossy(&denied.stderr).contains("SESSION_READ_ONLY"));
    let rpc_denied = Command::new(cargo_bin("smith"))
        .args(["rpc", "--session", &id])
        .env("XDG_DATA_HOME", &home)
        .env_remove("SMITH_PROVIDER")
        .output()
        .unwrap();
    assert!(!rpc_denied.status.success());
    assert!(String::from_utf8_lossy(&rpc_denied.stderr).contains("SESSION_READ_ONLY"));
    assert_eq!(std::fs::read(&path).unwrap(), before_denied);

    let export = home.join("transcript.jsonl");
    let (mut ui, rx) = observer(&home, &id, Some(&export));
    #[cfg(unix)]
    let before_terminal = terminal_state(ui.master.as_ref());
    let mut seen = String::new();
    marker(&rx, &mut seen, "READONLY");
    marker(&rx, &mut seen, "ctive");
    let original = std::fs::read(&path).unwrap();
    ui.write_all(b"not a prompt\r");
    marker(&rx, &mut seen, "READ ONLY: prompt");
    assert_eq!(std::fs::read(&path).unwrap(), original);

    ui.write_all(b"f");
    marker(&rx, &mut seen, "FOLLOWLIVE");
    for (request, input) in ["filler"; 12]
        .into_iter()
        .chain(["new live entry"])
        .enumerate()
    {
        let request = request + 2;
        writer_in
            .write_all(
                format!(
                    "{{\"id\":{request},\"method\":\"prompt\",\"params\":{{\"input\":\"{input}\"}}}}\n"
                )
                .as_bytes(),
            )
            .unwrap();
        loop {
            if rpc_line(&ready_rx)["event"] == "end" {
                break;
            }
        }
    }
    marker(&rx, &mut seen, "newliveentry");
    ui.write_all(b"b");
    marker(&rx, &mut seen, "paused");
    // Export supersedes this pending navigation fetch.
    // Wait for independent export and current-snapshot completions in either order.
    ui.write_all(b"e");
    completion_markers(&rx, &mut seen, &["Exported to", "initial"]);
    assert!(
        std::fs::read_to_string(&export)
            .unwrap()
            .contains("new live entry")
    );
    let before_fork = std::fs::read(&path).unwrap();
    ui.write_all(b"k");
    marker(&rx, &mut seen, "Fork crea");
    assert_eq!(std::fs::read(&path).unwrap(), before_fork);
    ui.write_all(b"q");
    ui.close_input();
    ui.wait();
    #[cfg(unix)]
    assert_eq!(terminal_state(ui.master.as_ref()), before_terminal);
    writer_in
        .write_all(b"{\"id\":15,\"method\":\"ping\"}\n")
        .unwrap();
    assert_eq!(rpc_line(&ready_rx)["pong"], true);
    drop(writer_in);
    assert!(writer.0.wait().unwrap().success());
    let files: Vec<_> = std::fs::read_dir(home.join("smith/sessions/run"))
        .unwrap()
        .filter_map(|e| {
            let p = e.unwrap().path();
            (p.extension().is_some_and(|ext| ext == "smh")).then_some(p)
        })
        .collect();
    assert_eq!(files.len(), 2);
    let _ = std::fs::remove_dir_all(home);
}

#[test]
fn completion_markers_accept_export_then_snapshot() {
    let (tx, rx) = mpsc::channel();
    tx.send("Exported to transcript".into()).unwrap();
    tx.send(" initial later".into()).unwrap();
    drop(tx);
    let mut output = String::new();
    completion_markers(&rx, &mut output, &["Exported to", "initial"]);
    assert_eq!(output, " later");
}

#[test]
fn completion_markers_accept_snapshot_then_split_export() {
    let (tx, rx) = mpsc::channel();
    tx.send("ini".into()).unwrap();
    tx.send("tial Exported".into()).unwrap();
    tx.send(" to later".into()).unwrap();
    drop(tx);
    let mut output = String::new();
    completion_markers(&rx, &mut output, &["Exported to", "initial"]);
    assert_eq!(output, " later");
}

#[test]
fn completion_markers_require_fresh_snapshot_after_consumed_stale_marker() {
    let (tx, rx) = mpsc::channel();
    tx.send("initial stale".into()).unwrap();
    tx.send(" Exported to transcript initial fresh later".into())
        .unwrap();
    drop(tx);
    let mut output = String::new();
    marker(&rx, &mut output, "initial");
    completion_markers(&rx, &mut output, &["Exported to", "initial"]);
    assert_eq!(output, " fresh later");
}

#[test]
fn ctrl_c_exits_observer_and_restores_terminal() {
    let home = std::env::temp_dir().join(format!("smith-observer-ctrl-c-{}", std::process::id()));
    let _ = std::fs::remove_dir_all(&home);
    let id = session(&home, "ctrl-c session");
    let (mut ui, rx) = observer(&home, &id, None);
    #[cfg(unix)]
    let before_terminal = terminal_state(ui.master.as_ref());
    let mut seen = String::new();
    marker(&rx, &mut seen, "READONLY");
    ui.write_all(&[3]);
    ui.close_input();
    ui.wait();
    #[cfg(unix)]
    assert_eq!(terminal_state(ui.master.as_ref()), before_terminal);
    let _ = std::fs::remove_dir_all(home);
}