cli.rs

use std::fs;
use std::fs::File;
use std::io::{self, IsTerminal, Read, Write};
use std::net::Shutdown;
use std::os::unix::net::UnixStream;
use std::path::{Path, PathBuf};
use std::process::{Command, Stdio};
use std::thread;
use std::time::{Duration, Instant};

use anyhow::{Context, Result};
use chrono::{Local, TimeZone};
use clap::{Parser, Subcommand};
#[cfg(target_os = "linux")]
use nix::sys::signal::{Signal, kill};
#[cfg(target_os = "linux")]
use nix::unistd::Pid;
use rand::random;
use serde::Serialize;
use vt100::Parser as VtParser;

use crate::broker;
use crate::protocol::{
    BROKER_API_VERSION, BrokerRequest, BrokerResponse, Event, SessionState, SessionSummary,
};
use crate::store;

const DEFAULT_RUN_YIELD_TIMEOUT_MS: u64 = 10_000;
const DEFAULT_ATTACH_YIELD_TIMEOUT_MS: u64 = 300_000;

#[derive(Parser, Debug)]
#[command(name = "contd")]
#[command(about = "Stateless wrapper for interactive CLI sessions")]
pub struct Cli {
    #[arg(long, global = true, help = "Emit machine-readable JSON responses")]
    json: bool,
    #[arg(
        long,
        global = true,
        help = "Print extra diagnostics (stdin tty, process state, wait channel)"
    )]
    verbose: bool,

    #[command(subcommand)]
    command: Commands,
}

#[derive(Subcommand, Debug)]
enum Commands {
    #[command(about = "Start a new managed interactive session")]
    Run {
        #[arg(
            long = "yield_timeout",
            value_name = "DURATION",
            value_parser = parse_yield_timeout_ms,
            help = "Hard deadline for this call (default 10s). When reached, return with collected output even if the process is still producing output (supports: ms, s, min/m, h/hr, d; max 1d)"
        )]
        yield_timeout: Option<u64>,
        #[arg(
            required = true,
            trailing_var_arg = true,
            help = "Command and arguments to run under contd session management"
        )]
        command: Vec<String>,
        #[arg(
            long,
            help = "Render PTY output through a virtual 80-column screen with frame (safe for TUI control sequences)"
        )]
        raw_pty: bool,
    },
    #[command(about = "Attach to a session and optionally send input")]
    Attach {
        #[arg(help = "Session id or unique prefix to attach to")]
        session_id: String,
        #[arg(
            long,
            help = "Input payload; supports terminal escapes (e.g. \\n newline, \\r Enter, \\t Tab, \\e ESC, \\x1b byte, \\u{2191} Unicode)"
        )]
        input: Option<String>,
        #[arg(
            long = "yield_timeout",
            value_name = "DURATION",
            value_parser = parse_yield_timeout_ms,
            help = "Hard deadline for this call (default 300s). When reached, return with collected output even if the process is still producing output (supports: ms, s, min/m, h/hr, d; max 1d)"
        )]
        yield_timeout: Option<u64>,
        #[arg(
            long,
            help = "Render PTY output through a virtual 80-column screen with frame (safe for TUI control sequences)"
        )]
        raw_pty: bool,
    },
    #[command(about = "Terminate a running session")]
    Kill {
        #[arg(help = "Session id or unique prefix to terminate")]
        session_id: String,
    },
    #[command(about = "List known sessions (or one matched by id/prefix)")]
    List {
        #[arg(help = "Optional session id/prefix filter")]
        session_id: Option<String>,
    },
    #[command(about = "Garbage-collect finished sessions")]
    Gc,
    #[command(about = "Stop brokers and remove all state")]
    Clean,
    #[command(hide = true)]
    Broker {
        #[arg(
            long = "state-root",
            help = "State root directory used by the background broker"
        )]
        state_root: PathBuf,
    },
}

pub fn run() -> i32 {
    let cli = Cli::parse();
    let json = cli.json;

    match run_inner(cli) {
        Ok(code) => code,
        Err(err) => {
            if json {
                let event = Event::Error {
                    code: "runtime_error".into(),
                    message: err.to_string(),
                };
                println!("{}", event.to_json_line());
            } else {
                eprintln!("{}", style("error", COLOR_RED, true));
                eprintln!("{}", err);
            }
            1
        }
    }
}

fn run_inner(cli: Cli) -> Result<i32> {
    let json = cli.json;
    let verbose = cli.verbose;

    match cli.command {
        Commands::Broker { state_root } => {
            broker::run_broker(&state_root)?;
            Ok(0)
        }
        Commands::Run {
            command,
            yield_timeout,
            raw_pty,
        } => {
            let response = send_to_broker(BrokerRequest::Run {
                command,
                yield_timeout_ms: Some(yield_timeout.unwrap_or(DEFAULT_RUN_YIELD_TIMEOUT_MS)),
            })?;
            Ok(handle_response(response, json, raw_pty, verbose, None))
        }
        Commands::Attach {
            session_id,
            input,
            yield_timeout,
            raw_pty,
        } => {
            let payload = match input {
                Some(value) => decode_cli_input_escapes(&value)?,
                None => read_stdin_if_piped()?,
            };

            let response = send_to_broker(BrokerRequest::Attach {
                session_id,
                input: payload,
                yield_timeout_ms: Some(yield_timeout.unwrap_or(DEFAULT_ATTACH_YIELD_TIMEOUT_MS)),
            })?;
            Ok(handle_response(response, json, raw_pty, verbose, None))
        }
        Commands::Kill { session_id } => {
            let response = send_to_broker(BrokerRequest::Kill { session_id })?;
            Ok(handle_response(response, json, false, verbose, None))
        }
        Commands::List { session_id } => {
            let highlight = session_id.clone();
            let response = send_to_broker(BrokerRequest::List { session_id })?;
            Ok(handle_response(
                response,
                json,
                false,
                verbose,
                highlight.as_deref(),
            ))
        }
        Commands::Gc => {
            let response = send_to_broker(BrokerRequest::Gc)?;
            Ok(handle_response(response, json, false, verbose, None))
        }
        Commands::Clean => clean_state(json),
    }
}

#[derive(Debug, Serialize)]
struct CleanReport {
    roots: Vec<String>,
    removed_roots: usize,
    stopped_brokers: usize,
}

fn clean_state(json: bool) -> Result<i32> {
    let roots = clean_roots()?;
    let stopped_brokers = stop_brokers_for_roots(&roots)?;

    let mut removed_roots = 0usize;
    for root in &roots {
        if root.exists() {
            fs::remove_dir_all(root)
                .with_context(|| format!("failed to remove state root {}", root.display()))?;
            removed_roots += 1;
        }
    }

    let report = CleanReport {
        roots: roots
            .iter()
            .map(|p| p.to_string_lossy().to_string())
            .collect(),
        removed_roots,
        stopped_brokers,
    };

    if json {
        println!(
            "{}",
            serde_json::to_string(&report).expect("clean report serialization")
        );
    } else {
        println!("{}", style("cleaned", COLOR_CYAN, true));
        println!("stopped brokers: {}", report.stopped_brokers);
        println!("removed roots: {}", report.removed_roots);
        for root in report.roots {
            println!("root: {root}");
        }
    }

    Ok(0)
}

fn clean_roots() -> Result<Vec<PathBuf>> {
    let mut roots = Vec::new();
    let current = store::state_root()?;
    roots.push(current.clone());

    let default = default_state_root()?;
    if default != current {
        roots.push(default);
    }

    Ok(roots)
}

fn default_state_root() -> Result<PathBuf> {
    Ok(PathBuf::from("/tmp").join("contd"))
}

#[cfg(target_os = "linux")]
fn stop_brokers_for_roots(roots: &[PathBuf]) -> Result<usize> {
    let mut matched_pids = Vec::new();

    for entry in fs::read_dir("/proc").context("failed to read /proc")? {
        let entry = entry?;
        let name = entry.file_name();
        let Some(name) = name.to_str() else {
            continue;
        };
        let Ok(pid_raw) = name.parse::<i32>() else {
            continue;
        };

        let cmdline_path = entry.path().join("cmdline");
        let raw = match fs::read(&cmdline_path) {
            Ok(raw) => raw,
            Err(_) => continue,
        };
        let args = parse_proc_cmdline(&raw);

        if roots.iter().any(|root| is_broker_for_root(&args, root)) {
            matched_pids.push(pid_raw);
        }
    }

    for pid_raw in &matched_pids {
        let _ = kill(Pid::from_raw(*pid_raw), Signal::SIGTERM);
    }

    let deadline = Instant::now() + Duration::from_millis(700);
    while Instant::now() < deadline {
        if matched_pids
            .iter()
            .all(|pid_raw| !process_is_alive(*pid_raw))
        {
            break;
        }
        thread::sleep(Duration::from_millis(20));
    }

    for pid_raw in &matched_pids {
        if process_is_alive(*pid_raw) {
            let _ = kill(Pid::from_raw(*pid_raw), Signal::SIGKILL);
        }
    }

    Ok(matched_pids.len())
}

#[cfg(not(target_os = "linux"))]
fn stop_brokers_for_roots(_roots: &[PathBuf]) -> Result<usize> {
    Ok(0)
}

#[cfg(target_os = "linux")]
fn process_is_alive(pid_raw: i32) -> bool {
    kill(Pid::from_raw(pid_raw), None::<Signal>).is_ok()
}

#[cfg(target_os = "linux")]
fn parse_proc_cmdline(raw: &[u8]) -> Vec<String> {
    raw.split(|b| *b == 0)
        .filter(|segment| !segment.is_empty())
        .map(|segment| String::from_utf8_lossy(segment).to_string())
        .collect()
}

#[cfg(target_os = "linux")]
fn is_broker_for_root(args: &[String], root: &Path) -> bool {
    if args.len() < 4 {
        return false;
    }
    args[1] == "broker" && args[2] == "--state-root" && args[3] == root.to_string_lossy()
}

fn handle_response(
    response: BrokerResponse,
    json: bool,
    raw_pty: bool,
    verbose: bool,
    list_highlight_prefix: Option<&str>,
) -> i32 {
    if json {
        return handle_response_json(response);
    }

    handle_response_human(response, raw_pty, verbose, list_highlight_prefix)
}

fn handle_response_json(response: BrokerResponse) -> i32 {
    match response {
        BrokerResponse::Event { event } => {
            let code = match &event {
                Event::Error { code, .. }
                    if code == "session_not_found"
                        || code == "session_unrecoverable"
                        || code == "session_id_ambiguous" =>
                {
                    2
                }
                Event::Error { .. } => 1,
                _ => 0,
            };
            println!("{}", event.to_json_line());
            code
        }
        other => {
            println!(
                "{}",
                serde_json::to_string(&other).expect("response serialization")
            );
            0
        }
    }
}

fn handle_response_human(
    response: BrokerResponse,
    raw_pty: bool,
    verbose: bool,
    list_highlight_prefix: Option<&str>,
) -> i32 {
    match response {
        BrokerResponse::Event { event } => render_event_human(event, raw_pty, verbose),
        BrokerResponse::List { sessions } => {
            println!("{}", style("sessions", COLOR_CYAN, true));
            render_session_tree(&sessions, list_highlight_prefix);
            0
        }
        BrokerResponse::Gc { removed } => {
            println!("removed sessions: {removed}");
            0
        }
        BrokerResponse::Pong { version } => {
            println!("pong v{version}");
            0
        }
    }
}

fn render_event_human(event: Event, raw_pty: bool, verbose: bool) -> i32 {
    match event {
        Event::SessionWaitingInput {
            session_id,
            output,
            detector,
        } => {
            println!("{}", style("ready for input", COLOR_GREEN, true));
            println!("id: {session_id}");
            print_detector(&detector, verbose);
            print_pty_output(&output, raw_pty);
            0
        }
        Event::SessionSteadyInput {
            session_id,
            output,
            detector,
        } => {
            println!(
                "{}",
                style("screen stable, may be ready for input", COLOR_GREEN, true)
            );
            println!("id: {session_id}");
            print_detector(&detector, verbose);
            print_pty_output(&output, raw_pty);
            0
        }
        Event::SessionYieldTimeout {
            session_id,
            output,
            timeout_ms,
        } => {
            println!("{}", style("timed out", COLOR_YELLOW, true));
            println!("id: {session_id}");
            println!("timeout_ms: {timeout_ms}");
            print_pty_output(&output, raw_pty);
            0
        }
        Event::SessionExited {
            session_id,
            output,
            exit_code,
            signal,
        } => {
            let failed = exit_code.map(|code| code != 0).unwrap_or(false) || signal.is_some();
            if failed {
                println!("{}", style("failed", COLOR_RED, true));
                println!("reason: process ended with error");
            } else {
                println!("{}", style("exited", COLOR_MAGENTA, true));
            }
            println!("id: {session_id}");
            println!("exit_code: {}", format_option_i32(exit_code));
            println!("signal: {}", format_option_i32(signal));
            print_pty_output(&output, raw_pty);
            0
        }
        Event::Error { code, message } => {
            println!("{}", style("error", COLOR_RED, true));
            println!("code: {code}");
            println!("message: {message}");
            if code == "session_not_found"
                || code == "session_unrecoverable"
                || code == "session_id_ambiguous"
            {
                2
            } else {
                1
            }
        }
    }
}

fn send_to_broker(request: BrokerRequest) -> Result<BrokerResponse> {
    let root = store::state_root()?;
    // Client commands are stateless; broker keeps the long-lived session state.
    ensure_broker_available(&root)?;
    request_once(&root, &request)
}

fn ensure_broker_available(root: &Path) -> Result<()> {
    store::ensure_state_dirs(root)?;

    match ping_broker_status(root)? {
        BrokerPingStatus::Compatible => return Ok(()),
        BrokerPingStatus::Incompatible => {
            // A running but incompatible broker (older/newer wire format) must be removed,
            // otherwise a stale process can keep the root lock and prevent bootstrap.
            restart_incompatible_broker(root)?;
        }
        BrokerPingStatus::Unavailable => {}
    }

    // Lazy bootstrap: spawn hidden broker subcommand from the same binary.
    let exe = std::env::current_exe().context("failed to locate current executable")?;
    let broker_stderr =
        BrokerStderrLog::new().context("failed to allocate broker stderr capture")?;
    let broker_stderr_writer = broker_stderr
        .writer()
        .context("failed to prepare broker stderr capture")?;
    let mut child = Command::new(exe)
        .arg("broker")
        .arg("--state-root")
        .arg(root)
        .stdin(Stdio::null())
        .stdout(Stdio::null())
        .stderr(Stdio::from(broker_stderr_writer))
        .spawn()
        .context("failed to start broker process")?;

    for _ in 0..60 {
        if matches!(ping_broker_status(root)?, BrokerPingStatus::Compatible) {
            return Ok(());
        }
        if let Some(status) = child.try_wait().context("failed to poll broker process")? {
            let stderr = read_broker_stderr(&broker_stderr);
            if stderr.is_empty() {
                anyhow::bail!("broker exited before becoming ready ({status})");
            }
            anyhow::bail!("broker exited before becoming ready ({status}): {stderr}");
        }
        thread::sleep(Duration::from_millis(50));
    }

    let stderr = read_broker_stderr(&broker_stderr);
    if stderr.is_empty() {
        anyhow::bail!("broker failed to start");
    }
    anyhow::bail!("broker failed to start: {stderr}")
}

fn read_broker_stderr(stderr_log: &BrokerStderrLog) -> String {
    stderr_log.read()
}

struct BrokerStderrLog {
    path: PathBuf,
}

impl BrokerStderrLog {
    fn new() -> Result<Self> {
        let mut path = std::env::temp_dir();
        path.push(format!(
            "contd-broker-stderr-{}-{}.log",
            std::process::id(),
            random::<u64>()
        ));
        let _ = File::options()
            .create_new(true)
            .write(true)
            .open(&path)
            .with_context(|| format!("failed to create {}", path.display()))?;
        Ok(Self { path })
    }

    fn writer(&self) -> Result<File> {
        File::options()
            .append(true)
            .open(&self.path)
            .with_context(|| format!("failed to open {}", self.path.display()))
    }

    fn read(&self) -> String {
        let raw = match fs::read_to_string(&self.path) {
            Ok(raw) => raw,
            Err(_) => return String::new(),
        };
        raw.lines()
            .map(str::trim)
            .filter(|line| !line.is_empty())
            .collect::<Vec<_>>()
            .join(" | ")
    }
}

impl Drop for BrokerStderrLog {
    fn drop(&mut self) {
        let _ = fs::remove_file(&self.path);
    }
}

enum BrokerPingStatus {
    Compatible,
    Incompatible,
    Unavailable,
}

fn ping_broker_status(root: &Path) -> Result<BrokerPingStatus> {
    match request_once(root, &BrokerRequest::Ping) {
        Ok(BrokerResponse::Pong { version }) => {
            if version == BROKER_API_VERSION {
                Ok(BrokerPingStatus::Compatible)
            } else {
                Ok(BrokerPingStatus::Incompatible)
            }
        }
        Ok(_) => Ok(BrokerPingStatus::Incompatible),
        Err(err) => {
            let message = err.to_string();
            if message.contains("failed to decode broker response") {
                return Ok(BrokerPingStatus::Incompatible);
            }
            Ok(BrokerPingStatus::Unavailable)
        }
    }
}

fn restart_incompatible_broker(root: &Path) -> Result<()> {
    #[cfg(target_os = "linux")]
    {
        let _ = stop_brokers_for_roots(&[root.to_path_buf()])?;
    }
    let socket_path = store::broker_socket(root);
    if socket_path.exists() {
        let _ = fs::remove_file(&socket_path);
    }
    Ok(())
}

fn request_once(root: &Path, request: &BrokerRequest) -> Result<BrokerResponse> {
    store::ensure_state_dirs(root)?;
    let socket_path = store::broker_socket(root);
    let mut stream = connect_with_timeout(&socket_path, Duration::from_millis(300))?;
    let response_timeout = rpc_timeout_for(request);
    stream
        .set_read_timeout(Some(response_timeout))
        .context("failed to set broker read timeout")?;
    stream
        .set_write_timeout(Some(Duration::from_secs(2)))
        .context("failed to set broker write timeout")?;
    let encoded = serde_json::to_vec(request).context("failed to encode request")?;
    stream
        .write_all(&encoded)
        .with_context(|| format!("failed to write request to {}", socket_path.display()))?;
    stream
        .shutdown(Shutdown::Write)
        .with_context(|| format!("failed to shutdown write side of {}", socket_path.display()))?;

    let mut raw = Vec::new();
    stream
        .read_to_end(&mut raw)
        .with_context(|| format!("failed to read response from {}", socket_path.display()))?;
    let response = serde_json::from_slice::<BrokerResponse>(&raw)
        .context("failed to decode broker response")?;
    Ok(response)
}

fn connect_with_timeout(socket_path: &Path, timeout: Duration) -> Result<UnixStream> {
    let deadline = Instant::now() + timeout;
    loop {
        match UnixStream::connect(socket_path) {
            Ok(stream) => return Ok(stream),
            Err(err) => {
                if Instant::now() >= deadline {
                    return Err(err).with_context(|| {
                        format!(
                            "timed out connecting to broker socket {}",
                            socket_path.display()
                        )
                    });
                }
                thread::sleep(Duration::from_millis(20));
            }
        }
    }
}

fn rpc_timeout_for(request: &BrokerRequest) -> Duration {
    match request {
        // Boundary detection for interactive sessions can legitimately take longer
        // (for example, sleep before returning to stdin or natural process exit).
        BrokerRequest::Run {
            yield_timeout_ms: Some(timeout_ms),
            ..
        }
        | BrokerRequest::Attach {
            yield_timeout_ms: Some(timeout_ms),
            ..
        } => Duration::from_millis(*timeout_ms).saturating_add(Duration::from_secs(5)),
        BrokerRequest::Run { .. } | BrokerRequest::Attach { .. } => Duration::from_secs(300),
        BrokerRequest::Kill { .. }
        | BrokerRequest::List { .. }
        | BrokerRequest::Gc
        | BrokerRequest::Ping => Duration::from_secs(8),
    }
}

fn read_stdin_if_piped() -> Result<String> {
    if io::stdin().is_terminal() {
        return Ok(String::new());
    }

    let mut input = String::new();
    io::stdin()
        .read_to_string(&mut input)
        .context("failed to read stdin")?;
    Ok(input)
}

fn decode_cli_input_escapes(value: &str) -> Result<String> {
    let mut out = String::with_capacity(value.len());
    let mut chars = value.chars();

    while let Some(ch) = chars.next() {
        if ch != '\\' {
            out.push(ch);
            continue;
        }

        let Some(next) = chars.next() else {
            anyhow::bail!("invalid --input escape: trailing backslash");
        };

        match next {
            'n' => out.push('\n'),       // newline (LF, 0x0A)
            'r' => out.push('\r'),       // carriage return / Enter (CR, 0x0D)
            't' => out.push('\t'),       // horizontal tab (HT, 0x09)
            'b' => out.push('\u{0008}'), // backspace (BS, 0x08)
            'f' => out.push('\u{000c}'), // form feed (FF, 0x0C)
            'v' => out.push('\u{000b}'), // vertical tab (VT, 0x0B)
            '0' => out.push('\0'),       // NUL byte (0x00)
            'e' => out.push('\u{001b}'), // escape (ESC, 0x1B)
            '\\' => out.push('\\'),      // literal backslash
            '"' => out.push('"'),        // literal double quote
            '\'' => out.push('\''),      // literal single quote
            'x' => {
                // \xHH: one raw byte in hex (useful for ESC sequences like \x1b[A).
                let h1 = chars.next().ok_or_else(|| {
                    anyhow::anyhow!("invalid --input escape: \\x expects 2 hex digits")
                })?;
                let h2 = chars.next().ok_or_else(|| {
                    anyhow::anyhow!("invalid --input escape: \\x expects 2 hex digits")
                })?;
                let byte = parse_hex_byte(h1, h2)?;
                out.push(byte as char);
            }
            'u' => {
                // \u{...}: Unicode scalar, e.g. \u{2191} for ↑.
                let open = chars
                    .next()
                    .ok_or_else(|| anyhow::anyhow!("invalid --input escape: \\u expects '{{'"))?;
                if open != '{' {
                    anyhow::bail!("invalid --input escape: \\u expects '{{...}}'");
                }
                let mut hex = String::new();
                loop {
                    let c = chars.next().ok_or_else(|| {
                        anyhow::anyhow!("invalid --input escape: missing '}}' for \\u{{...}}")
                    })?;
                    if c == '}' {
                        break;
                    }
                    if !c.is_ascii_hexdigit() {
                        anyhow::bail!("invalid --input escape: \\u{{...}} only accepts hex digits");
                    }
                    if hex.len() >= 6 {
                        anyhow::bail!("invalid --input escape: \\u{{...}} too long");
                    }
                    hex.push(c);
                }
                if hex.is_empty() {
                    anyhow::bail!("invalid --input escape: \\u{{...}} cannot be empty");
                }
                let scalar = u32::from_str_radix(&hex, 16).with_context(|| {
                    format!("invalid --input escape: bad unicode scalar \\u{{{hex}}}")
                })?;
                let ch = char::from_u32(scalar).ok_or_else(|| {
                    anyhow::anyhow!("invalid --input escape: bad unicode scalar \\u{{{hex}}}")
                })?;
                out.push(ch);
            }
            other => anyhow::bail!(
                "invalid --input escape: \\{other} (supported: \\\\, \\n, \\r, \\t, \\b, \\f, \\v, \\e, \\0, \\xHH, \\u{{...}}, \\\", \\\')"
            ),
        }
    }

    Ok(out)
}

fn parse_hex_byte(h1: char, h2: char) -> Result<u8> {
    let hex = [h1, h2].iter().collect::<String>();
    u8::from_str_radix(&hex, 16)
        .with_context(|| format!("invalid --input escape: bad hex byte \\x{hex}"))
}

fn print_detector(detector: &crate::protocol::DetectorSnapshot, verbose: bool) {
    println!("pid: {}", detector.pid);
    if !verbose {
        return;
    }
    if let Some(stdin_target) = &detector.stdin_target {
        println!("stdin tty: {stdin_target}");
    }
    if let Some(state) = detector.process_state {
        println!("proc state: {state}");
    }
    if let Some(wchan) = &detector.wchan {
        println!("wchan: {wchan}");
    }
}

fn print_pty_output(output: &str, raw_pty: bool) {
    if output.is_empty() {
        return;
    }

    println!();
    println!("{}", style("output:", COLOR_CYAN, false));
    // Default mode sanitizes terminal-control sequences that can corrupt shell UI.
    // `--raw-pty` renders through a virtual VT screen so control sequences stay scoped.
    let rendered = if raw_pty {
        render_framed_vt100_output(output)
    } else {
        sanitize_terminal_output(output)
    };
    print!("{rendered}");
    if !rendered.ends_with('\n') {
        println!();
    }
    let _ = io::stdout().flush();
}

fn sanitize_terminal_output(input: &str) -> String {
    let bytes = input.as_bytes();
    let mut out = String::with_capacity(input.len());
    let mut i = 0usize;

    while i < bytes.len() {
        let b = bytes[i];
        if b == 0x1b {
            if let Some((seq_end, final_byte)) = parse_csi_sequence(bytes, i) {
                // Keep SGR color/style sequences (CSI ... m), escape screen-control sequences.
                if final_byte == b'm' {
                    out.push_str(&String::from_utf8_lossy(&bytes[i..seq_end]));
                } else {
                    out.push_str(&escape_bytes(&bytes[i..seq_end]));
                }
                i = seq_end;
                continue;
            }
            out.push_str("\\x1b");
            i += 1;
            continue;
        }

        // Keep common layout controls; escape the rest to avoid terminal corruption.
        if b == b'\n' || b == b'\r' || b == b'\t' || (0x20..=0x7e).contains(&b) {
            out.push(b as char);
        } else {
            out.push_str(&format!("\\x{b:02x}"));
        }
        i += 1;
    }

    out
}

fn render_framed_vt100_output(input: &str) -> String {
    const COLS: u16 = 80;
    let rows = estimate_vt_rows(input, COLS as usize);
    let mut parser = VtParser::new(rows, COLS, 0);
    parser.process(input.as_bytes());

    let screen = parser.screen();
    let lines = trim_trailing_blank_lines(screen.contents());
    frame_lines(&lines, COLS as usize)
}

fn estimate_vt_rows(input: &str, cols: usize) -> u16 {
    let mut plain_bytes = 0usize;
    let mut line_breaks = 1usize;
    for b in input.bytes() {
        if b == b'\n' {
            line_breaks += 1;
        } else if (0x20..=0x7e).contains(&b) || b == b'\t' {
            plain_bytes += 1;
        }
    }
    let wraps = plain_bytes / cols.max(1);
    let estimated = line_breaks.saturating_add(wraps).saturating_add(8);
    estimated.clamp(8, u16::MAX as usize) as u16
}

fn trim_trailing_blank_lines(contents: String) -> Vec<String> {
    let mut lines: Vec<String> = contents
        .lines()
        .map(|line| line.trim_end().to_string())
        .collect();
    while lines.last().is_some_and(|line| line.is_empty()) {
        lines.pop();
    }
    if lines.is_empty() {
        lines.push(String::new());
    }
    lines
}

fn frame_lines(lines: &[String], width: usize) -> String {
    let mut out = String::new();
    out.push('+');
    out.push_str(&"-".repeat(width));
    out.push('+');
    out.push('\n');

    for line in lines {
        out.push('|');
        out.push_str(&fit_to_width(line, width));
        out.push('|');
        out.push('\n');
    }

    out.push('+');
    out.push_str(&"-".repeat(width));
    out.push('+');
    out
}

fn fit_to_width(line: &str, width: usize) -> String {
    let mut out = String::new();
    for ch in line.chars().take(width) {
        out.push(ch);
    }
    let len = out.chars().count();
    if len < width {
        out.push_str(&" ".repeat(width - len));
    }
    out
}

fn parse_csi_sequence(bytes: &[u8], start: usize) -> Option<(usize, u8)> {
    if start + 1 >= bytes.len() || bytes[start + 1] != b'[' {
        return None;
    }

    let mut idx = start + 2;
    while idx < bytes.len() {
        let b = bytes[idx];
        // CSI final byte range in ANSI escape grammar.
        if (0x40..=0x7e).contains(&b) {
            return Some((idx + 1, b));
        }
        idx += 1;
    }
    None
}

fn escape_bytes(bytes: &[u8]) -> String {
    let mut escaped = String::new();
    for b in bytes {
        escaped.push_str(&format!("\\x{b:02x}"));
    }
    escaped
}

fn render_session_tree(sessions: &[SessionSummary], list_highlight_prefix: Option<&str>) {
    if sessions.is_empty() {
        println!("(none)");
        return;
    }

    for (index, session) in sessions.iter().enumerate() {
        let last = index + 1 == sessions.len();
        let branch = if last { "└─" } else { "├─" };
        let indent = if last { "   " } else { "│  " };
        let state = state_label(&session.state);

        println!(
            "{branch} {}  {}",
            render_session_id(&session.session_id, list_highlight_prefix),
            style(state, state_color(&session.state), true)
        );
        render_summary_block(session, indent);
    }
}

fn render_session_id(session_id: &str, list_highlight_prefix: Option<&str>) -> String {
    let Some(prefix) = list_highlight_prefix else {
        return style(session_id, COLOR_CYAN, true);
    };
    if prefix.is_empty() {
        return style(session_id, COLOR_CYAN, true);
    }

    if let Some(rest) = session_id.strip_prefix(prefix) {
        if !io::stdout().is_terminal() {
            return session_id.to_string();
        }
        let matched = style_underline(prefix, COLOR_RED);
        let rest = style(rest, COLOR_CYAN, true);
        return format!("{matched}{rest}");
    }

    style(session_id, COLOR_CYAN, true)
}

fn render_summary_block(session: &SessionSummary, indent: &str) {
    println!(
        "{indent}cmd: {}",
        compact_command(&session.command.join(" "), 96)
    );
    println!(
        "{indent}pid: {}  updated: {}",
        session.pid,
        format_local_timestamp_seconds(session.last_updated_ms)
    );
    println!(
        "{indent}wait: {}  exit: {}  sig: {}",
        session.last_wchan.as_deref().unwrap_or("-"),
        format_option_i32(session.exit_code),
        format_option_i32(session.signal)
    );
}

fn state_label(state: &SessionState) -> &'static str {
    match state {
        SessionState::Running => "running",
        SessionState::WaitingInput => "waiting_input",
        SessionState::SteadyInput => "steady_input",
        SessionState::Exited => "exited",
        SessionState::Unrecoverable => "unrecoverable",
    }
}

fn state_color(state: &SessionState) -> &'static str {
    match state {
        SessionState::Running => COLOR_BLUE,
        SessionState::WaitingInput => COLOR_GREEN,
        SessionState::SteadyInput => COLOR_GREEN,
        SessionState::Exited => COLOR_MAGENTA,
        SessionState::Unrecoverable => COLOR_RED,
    }
}

fn compact_command(command: &str, max_chars: usize) -> String {
    if command.chars().count() <= max_chars {
        return command.to_string();
    }
    let mut out = String::new();
    for ch in command.chars().take(max_chars.saturating_sub(3)) {
        out.push(ch);
    }
    out.push_str("...");
    out
}

fn format_option_i32(value: Option<i32>) -> String {
    value
        .map(|v| v.to_string())
        .unwrap_or_else(|| "-".to_string())
}

fn format_local_timestamp_seconds(epoch_ms: i64) -> String {
    Local
        .timestamp_millis_opt(epoch_ms)
        .single()
        .map(|dt| dt.format("%Y-%m-%d %H:%M:%S").to_string())
        .unwrap_or_else(|| "-".to_string())
}

fn parse_yield_timeout_ms(value: &str) -> Result<u64, String> {
    const MAX_TIMEOUT_MS: u64 = 24 * 60 * 60 * 1000;
    let raw = value.trim();
    if raw.is_empty() {
        return Err("yield_timeout cannot be empty".to_string());
    }

    let digits_len = raw.chars().take_while(|c| c.is_ascii_digit()).count();
    if digits_len == 0 {
        return Err(format!(
            "invalid yield_timeout '{value}': expected a number with unit, e.g. 100ms, 10s, 1min, 1h, 1d"
        ));
    }

    let number_str = &raw[..digits_len];
    let unit_str = raw[digits_len..].trim().to_ascii_lowercase();
    let number = number_str
        .parse::<u64>()
        .map_err(|_| format!("invalid yield_timeout number '{number_str}'"))?;

    let multiplier = match unit_str.as_str() {
        "" | "ms" => 1u64,
        "s" => 1000u64,
        "m" | "min" => 60_000u64,
        "h" | "hr" => 3_600_000u64,
        "d" => 86_400_000u64,
        _ => {
            return Err(format!(
                "unsupported yield_timeout unit in '{value}'; use one of: ms, s, min, m, h, hr, d"
            ));
        }
    };

    let timeout_ms = number
        .checked_mul(multiplier)
        .ok_or_else(|| format!("yield_timeout '{value}' is too large"))?;
    if timeout_ms > MAX_TIMEOUT_MS {
        return Err(format!("yield_timeout '{value}' exceeds maximum 1d"));
    }
    Ok(timeout_ms)
}

fn style(input: &str, color: &str, bold: bool) -> String {
    if !io::stdout().is_terminal() {
        return input.to_string();
    }

    if bold {
        format!("\x1b[1;{color}m{input}\x1b[0m")
    } else {
        format!("\x1b[{color}m{input}\x1b[0m")
    }
}

fn style_underline(input: &str, color: &str) -> String {
    if !io::stdout().is_terminal() {
        return input.to_string();
    }
    format!("\x1b[4;{color}m{input}\x1b[0m")
}

const COLOR_RED: &str = "31";
const COLOR_GREEN: &str = "32";
const COLOR_YELLOW: &str = "33";
const COLOR_BLUE: &str = "34";
const COLOR_MAGENTA: &str = "35";
const COLOR_CYAN: &str = "36";

#[cfg(test)]
mod tests {
    use super::{
        decode_cli_input_escapes, parse_yield_timeout_ms, render_framed_vt100_output,
        sanitize_terminal_output,
    };
    #[cfg(target_os = "linux")]
    use super::{is_broker_for_root, parse_proc_cmdline};
    #[cfg(target_os = "linux")]
    use std::path::Path;

    #[test]
    fn keeps_sgr_color_sequences() {
        // Purpose: preserve human-friendly colors in default (non-raw) mode.
        let raw = "\u{1b}[31mred\u{1b}[0m";
        assert_eq!(sanitize_terminal_output(raw), raw);
    }

    #[test]
    fn escapes_screen_control_sequences() {
        // Purpose: avoid shell corruption from clear-screen/cursor movement controls.
        let raw = "\u{1b}[2J\u{1b}[Hhello";
        let rendered = sanitize_terminal_output(raw);
        assert!(rendered.contains("\\x1b\\x5b\\x32\\x4a"));
        assert!(rendered.contains("\\x1b\\x5b\\x48"));
        assert!(rendered.ends_with("hello"));
    }

    #[test]
    fn raw_pty_renders_in_vt100_frame() {
        let raw = "\u{1b}[2J\u{1b}[HMinimal TUI Demo\r\nLast key: -\r\nInput:";
        let rendered = render_framed_vt100_output(raw);
        assert!(rendered.starts_with("+"));
        assert!(rendered.contains("|Minimal TUI Demo"));
        assert!(rendered.contains("|Last key: -"));
        assert!(rendered.contains("|Input:"));
        assert!(rendered.ends_with(
            "+--------------------------------------------------------------------------------+"
        ));
    }

    #[cfg(target_os = "linux")]
    #[test]
    fn parses_proc_cmdline() {
        let raw = b"/tmp/contd\0broker\0--state-root\0/tmp/state/contd\0";
        let args = parse_proc_cmdline(raw);
        assert_eq!(
            args,
            vec!["/tmp/contd", "broker", "--state-root", "/tmp/state/contd"]
        );
    }

    #[cfg(target_os = "linux")]
    #[test]
    fn detects_broker_for_root() {
        let args = vec![
            "/tmp/contd".to_string(),
            "broker".to_string(),
            "--state-root".to_string(),
            "/tmp/state/contd".to_string(),
        ];
        assert!(is_broker_for_root(&args, Path::new("/tmp/state/contd")));
        assert!(!is_broker_for_root(
            &args,
            Path::new("/tmp/state/another-contd")
        ));
    }

    #[test]
    fn parses_yield_timeout_units() {
        assert_eq!(parse_yield_timeout_ms("100ms").expect("100ms"), 100);
        assert_eq!(parse_yield_timeout_ms("10s").expect("10s"), 10_000);
        assert_eq!(parse_yield_timeout_ms("1min").expect("1min"), 60_000);
        assert_eq!(parse_yield_timeout_ms("1m").expect("1m"), 60_000);
        assert_eq!(parse_yield_timeout_ms("1h").expect("1h"), 3_600_000);
        assert_eq!(parse_yield_timeout_ms("1hr").expect("1hr"), 3_600_000);
        assert_eq!(parse_yield_timeout_ms("1d").expect("1d"), 86_400_000);
    }

    #[test]
    fn yield_timeout_rejects_invalid_or_too_large_values() {
        assert!(parse_yield_timeout_ms("abc").is_err());
        assert!(parse_yield_timeout_ms("10x").is_err());
        assert!(parse_yield_timeout_ms("2d").is_err());
    }

    #[test]
    fn decodes_cli_input_escapes() {
        let decoded =
            decode_cli_input_escapes(r#"line1\nline2\t\\\"\'\0\e\x41\u{2191}"#).expect("decode");
        assert_eq!(decoded, format!("line1\nline2\t\\\"\'\0\u{001b}A\u{2191}"));
    }

    #[test]
    fn decodes_arrow_key_escape_sequence() {
        let up = decode_cli_input_escapes(r#"\x1b[A"#).expect("decode up");
        assert_eq!(up.as_bytes(), b"\x1b[A");
    }

    #[test]
    fn rejects_invalid_cli_input_escape() {
        assert!(decode_cli_input_escapes(r"bad\q").is_err());
        assert!(decode_cli_input_escapes("bad\\").is_err());
        assert!(decode_cli_input_escapes(r"\x1").is_err());
        assert!(decode_cli_input_escapes(r"\xZZ").is_err());
        assert!(decode_cli_input_escapes(r"\u{}").is_err());
        assert!(decode_cli_input_escapes(r"\u{xyz}").is_err());
    }
}