roundabout

use std::{

    io::{self, Read},

    path::{Path, PathBuf},

};

use serde::Deserialize;

use serde_json::{Value, json};

use crate::{

    artifacts, client,

    config::{Config, StopPolicy, Unwrapped},

    history, paths,

    proto::{self, Event, LeaseInfo, Reader, Request},

    units::Dur,

};

#[derive(Deserialize)]

struct Input {

    #[serde(default)]

    session_id: String,

    #[serde(default)]

    cwd: Option<PathBuf>,

    #[serde(default)]

    tool_name: Option<String>,

    #[serde(default, rename = "tool_input")]

    tool_args: Option<Value>,

    #[serde(default)]

    stop_hook_active: bool,

}

fn read_input() -> Option<Input> {

    let mut buf = String::new();

    io::stdin().read_to_string(&mut buf).ok()?;

    serde_json::from_str(&buf).ok()

}

fn emit(v: &Value) {

    println!("{v}");

}

/// quote-aware split on `&& || ; | & \n`. subshells, heredocs and friends get

/// treated as words, which errs toward seeing a build that isn't there

fn segments(cmd: &str) -> Vec<Vec<String>> {

    let mut out = Vec::new();

    let mut words = Vec::new();

    let mut word = String::new();

    let mut quote: Option<char> = None;

    let mut chars = cmd.chars().peekable();

    let flush_word = |word: &mut String, words: &mut Vec<String>| {

        if !word.is_empty() {

            words.push(std::mem::take(word));

        }

    };

    while let Some(c) = chars.next() {

        match (quote, c) {

            (Some(q), c) if c == q => quote = None,

            (Some('"') | None, '\\') => {

                if let Some(n) = chars.next() {

                    word.push(n);

                }

            }

            (None, '\'' | '"') => quote = Some(c),

            (None, '&') if word.ends_with('>') || chars.peek() == Some(&'>') => word.push(c),

            (None, '&' | '|' | ';' | '\n' | '(' | ')') => {

                flush_word(&mut word, &mut words);

                if !words.is_empty() {

                    out.push(std::mem::take(&mut words));

                }

            }

            (None, c) if c.is_whitespace() => flush_word(&mut word, &mut words),

            (_, c) => word.push(c),

        }

    }

    flush_word(&mut word, &mut words);

    if !words.is_empty() {

        out.push(words);

    }

    out

}

/// peels env assignments and wrappers that don't change what runs

fn head(seg: &[String]) -> &[String] {

    let mut s = seg;

    loop {

        match s.first().map(String::as_str) {

            Some(w) if w.contains('=') && !w.starts_with('-') && !w.starts_with('=') => s = &s[1..],

            Some("time" | "nice" | "nohup" | "command" | "exec" | "env" | "sudo" | "doas") => {

                s = &s[1..];

            }

            Some("timeout") => {

                s = &s[1..];

                while s.first().is_some_and(|w| w.starts_with('-')) {

                    s = &s[1..];

                }

                s = s.get(1..).unwrap_or_default();

            }

            _ => return s,

        }

    }

}

fn is_rabt(word: &str) -> bool {

    word == "rabt" || word.ends_with("/rabt")

}

fn matches_pattern(seg: &[String], pattern: &str) -> bool {

    let want: Vec<&str> = pattern.split_whitespace().collect();

    let first_ok = |w: &str, p: &str| w == p || w.rsplit('/').next() == Some(p);

    !want.is_empty()

        && seg.len() >= want.len()

        && first_ok(&seg[0], want[0])

        && seg[1..want.len()]

            .iter()

            .zip(&want[1..])

            .all(|(w, p)| w == p)

}

enum Scan {

    Clean,

    Unwrapped(Vec<String>),

}

fn scan(cfg: &Config, cmd: &str) -> Scan {

    let mut found = Vec::new();

    for seg in segments(cmd) {

        let seg = head(&seg);

        if seg.first().is_some_and(|w| is_rabt(w)) {

            continue;

        }

        if cfg.detect_patterns().any(|p| matches_pattern(seg, p)) {

            let words: Vec<&str> = seg

                .iter()

                .map(String::as_str)

                .take_while(|w| !w.contains('>') && !w.contains('<'))

                .collect();

            found.push(words.join(" "));

        }

    }

    if found.is_empty() {

        Scan::Clean

    } else {

        Scan::Unwrapped(found)

    }

}

fn guess(cfg: &Config, cwd: &Path, cmd: &[String]) -> String {

    let key = history::key(&paths::project_root(cwd), cmd);

    history::lookup(&key).map_or_else(

        || format!("--size {}", cfg.enforce.default_size),

        |rec| {

            crate::run::base(cfg).map_or_else(

                |_| format!("--size {}", cfg.enforce.default_size),

                |b| {

                    let n = history::suggest(&rec, b, cmd);

                    format!("--ram {} --cores {} --time {}", n.ram, n.cores, Dur(n.time))

                },

            )

        },

    )

}

fn shell_quote(s: &str) -> String {

    format!("'{}'", s.replace('\'', r"'\''"))

}

fn pre_tool_use(cfg: &Config, input: Input) {

    if input.tool_name.as_deref() != Some("Bash") {

        return;

    }

    let Some(mut tool_input) = input.tool_args else {

        return;

    };

    let Some(cmd) = tool_input

        .get("command")

        .and_then(Value::as_str)

        .map(str::to_owned)

    else {

        return;

    };

    let cwd = input.cwd.unwrap_or_else(|| "/".into());

    match (scan(cfg, &cmd), cfg.enforce.unwrapped) {

        (Scan::Clean, _) | (_, Unwrapped::Allow) => {}

        (Scan::Unwrapped(found), Unwrapped::Deny) => {

            let first: Vec<String> = found[0].split_whitespace().map(Into::into).collect();

            let hint = guess(cfg, &cwd, &first);

            let reason = format!(

                "builds go through roundabout here so parallel agents don't oom the box. wrap each build \

                 in `rabt run <size> -- <cmd>`, e.g. `rabt run {hint} -- {}`. sizes are {}, or \

                 give --ram/--disk/--cores/--time yourself. for anything long use \

                 run_in_background and start `rabt watch <lease>` in the background too so you \

                 hear if it drags. see the roundabout skill",

                found[0],

                cfg.sizes.keys().cloned().collect::<Vec<_>>().join("/"),

            );

            emit(&json!({

                "hookSpecificOutput": {

                    "hookEventName": "PreToolUse",

                    "permissionDecision": "deny",

                    "permissionDecisionReason": reason,

                }

            }));

        }

        (Scan::Unwrapped(_), Unwrapped::Wrap) => {

            let hint = guess(cfg, &cwd, std::slice::from_ref(&cmd));

            tool_input["command"] =

                Value::String(format!("rabt run {hint} -- bash -c {}", shell_quote(&cmd)));

            // no permissionDecision on purpose. allow would skip the user's

            // prompt for whatever the agent wrote, rewriting shouldn't

            emit(&json!({

                "hookSpecificOutput": {

                    "hookEventName": "PreToolUse",

                    "updatedInput": tool_input,

                }

            }));

        }

    }

}

fn codex_pre_tool_use(cfg: &Config, input: Input) -> Option<Value> {

    if input.tool_name.as_deref() != Some("Bash") {

        return None;

    }

    let command = input.tool_args?.get("command")?.as_str()?.to_owned();

    let cwd = input.cwd.unwrap_or_else(|| "/".into());

    let Scan::Unwrapped(found) = scan(cfg, &command) else {

        return None;

    };

    if cfg.enforce.unwrapped == Unwrapped::Allow {

        return None;

    }

    let first: Vec<String> = found[0].split_whitespace().map(Into::into).collect();

    let hint = guess(cfg, &cwd, &first);

    let reason = format!(

        "builds go through roundabout here so parallel agents don't oom the box. run `rabt run {hint} -- {}` instead. see the roundabout skill",

        found[0]

    );

    Some(json!({

        "hookSpecificOutput": {

            "hookEventName": "PreToolUse",

            "permissionDecision": "deny",

            "permissionDecisionReason": reason,

        }

    }))

}

fn overdue(session: &str) -> Vec<LeaseInfo> {

    if session.is_empty() {

        return Vec::new();

    }

    let Some(mut s) = client::connect() else {

        return Vec::new();

    };

    if proto::send(

        &mut s,

        &Request::Overdue {

            session: session.to_owned(),

        },

    )

    .is_err()

    {

        return Vec::new();

    }

    match Reader::new(&s).and_then(|mut r| r.recv::<Event>()) {

        Ok(Some(Event::Overdue { leases })) => leases,

        _ => Vec::new(),

    }

}

/// same source `rabt run` reads. after `--continue` the env var can keep the

/// startup id while the json moves on, and we need both ends to agree

fn session(input: &Input, codex: bool) -> String {

    if codex {

        return input.session_id.clone();

    }

    std::env::var("CLAUDE_CODE_SESSION_ID")

        .ok()

        .filter(|s| !s.is_empty())

        .unwrap_or_else(|| input.session_id.clone())

}

/// both ids, since around /clear the env var and the json can disagree about

/// which session just ended. never spawns, a missing daemon holds no intents

fn session_end(input: &Input, codex: bool) {

    let mut ids = vec![session(input, codex), input.session_id.clone()];

    ids.dedup();

    for session in ids.into_iter().filter(|s| !s.is_empty()) {

        if let Some(mut s) = client::connect() {

            let _ = proto::send(

                &mut s,

                &Request::Expect {

                    session,

                    need: None,

                },

            );

            let _ = Reader::new(&s).and_then(|mut r| r.recv::<Event>());

        }

    }

}

fn remind(event: &str, input: &Input, codex: bool) {

    let late = overdue(&session(input, codex));

    if late.is_empty() {

        return;

    }

    let lines: Vec<String> = late

        .iter()

        .map(|l| {

            format!(

                "roundabout lease {} (`{}`) has run {} against a {} estimate. `rabt kill {}` to stop it, \

                 `rabt extend {} <dur>` if it's worth the wait",

                l.id,

                l.cmd.join(" "),

                Dur::secs(l.elapsed),

                Dur(l.need.time),

                l.id,

                l.id

            )

        })

        .collect();

    emit(&json!({

        "hookSpecificOutput": {

            "hookEventName": event,

            "additionalContext": lines.join("\n"),

        }

    }));

}

fn stop(cfg: &Config, event: &str, input: &Input) {

    if input.stop_hook_active || cfg.artifacts.stop == StopPolicy::Off {

        return;

    }

    let root = input.cwd.as_deref().map(paths::project_root);

    let report = artifacts::check(cfg, root.as_deref(), false);

    if !report.over() {

        return;

    }

    let msg = format!(

        "{}. run `rabt artifacts gc --cap`, or `rabt artifacts list` and clear out what you \

         don't need",

        report.describe(cfg)

    );

    let out = match cfg.artifacts.stop {

        StopPolicy::Block => json!({ "decision": "block", "reason": msg }),

        StopPolicy::Nudge => json!({

            "hookSpecificOutput": { "hookEventName": event, "additionalContext": msg }

        }),

        StopPolicy::Warn => json!({ "systemMessage": msg }),

        StopPolicy::Off => return,

    };

    emit(&out);

}

/// hooks never fail loudly. a broken roundabout shouldn't wedge the agent

pub fn main(cfg: &Config, event: &str) -> i32 {

    let Some(input) = read_input() else {

        return 0;

    };

    match event {

        "codex-pre-tool-use" => {

            if let Some(output) = codex_pre_tool_use(cfg, input) {

                emit(&output);

            }

        }

        "codex-post-tool-use" | "omp-post-tool-use" => remind("PostToolUse", &input, true),

        "codex-prompt-submit" | "omp-prompt-submit" => remind("UserPromptSubmit", &input, true),

        "codex-stop" | "omp-stop" | "stop" => stop(cfg, "Stop", &input),

        "codex-subagent-stop" | "omp-subagent-stop" | "subagent-stop" => {

            stop(cfg, "SubagentStop", &input);

        }

        "codex-session-end" | "omp-session-end" => session_end(&input, true),

        "pre-tool-use" | "omp-pre-tool-use" => pre_tool_use(cfg, input),

        "post-tool-use" => remind("PostToolUse", &input, false),

        "prompt-submit" => remind("UserPromptSubmit", &input, false),

        "session-end" => session_end(&input, false),

        other => eprintln!("rabt: unknown hook {other}"),

    }

    0

}

#[cfg(test)]

mod tests {

    use super::*;

    fn bash(command: &str) -> Input {

        Input {

            session_id: String::new(),

            cwd: Some("/".into()),

            tool_name: Some("Bash".into()),

            tool_args: Some(json!({ "command": command })),

            stop_hook_active: false,

        }

    }

    #[test]

    fn codex_denies_unwrapped_builds() {

        let output = codex_pre_tool_use(&Config::default(), bash("cargo build")).unwrap();

        assert_eq!(output["hookSpecificOutput"]["permissionDecision"], "deny");

        assert!(

            output["hookSpecificOutput"]["permissionDecisionReason"]

                .as_str()

                .unwrap()

                .contains("rabt run")

        );

    }

    #[test]

    fn codex_does_not_skip_approval_to_wrap_builds() {

        let mut cfg = Config::default();

        cfg.enforce.unwrapped = Unwrapped::Wrap;

        let output = codex_pre_tool_use(&cfg, bash("nix build .#example")).unwrap();

        assert_eq!(output["hookSpecificOutput"]["permissionDecision"], "deny");

        assert!(output["hookSpecificOutput"].get("updatedInput").is_none());

    }

    #[test]

    fn codex_ignores_other_tools_and_wrapped_builds() {

        assert!(

            codex_pre_tool_use(&Config::default(), bash("rabt run -s m -- cargo build")).is_none()

        );

        let mut input = bash("cargo build");

        input.tool_name = Some("apply_patch".into());

        assert!(codex_pre_tool_use(&Config::default(), input).is_none());

    }

}