roundabout

use std::{io, os::unix::net::UnixStream, path::PathBuf};

use pound::Parse;

use crate::{

    artifacts, client,

    config::Config,

    daemon, history, hook, paths,

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

    run,

    units::{Amount, Bytes, Dur},

};

/// one build at a time, please

#[derive(Parse)]

#[pound(name = "rabt", version = env!("CARGO_PKG_VERSION"))]

pub enum Rabt {

    /// wait your turn, then run the thing

    Run {

        /// preset from config, s m l xl out of the box

        #[pound(short, long)]

        size: Option<String>,

        #[pound(long)]

        ram: Option<Bytes>,

        #[pound(long)]

        disk: Option<Bytes>,

        /// count or percentage of the machine

        #[pound(long)]

        cores: Option<Amount<u32>>,

        /// how long you reckon it'll take

        #[pound(short, long)]

        time: Option<Dur>,

        /// kill it past this, exit 124

        #[pound(long)]

        deadline: Option<Dur>,

        /// stop queueing past this, exit 75

        #[pound(long)]

        max_wait: Option<Dur>,

        #[pound(long, hidden, env = "CLAUDE_CODE_SESSION_ID")]

        session: Option<String>,

        /// exactly the cores asked for, no spares

        #[pound(long)]

        exact: bool,

        #[pound(trailing)]

        cmd: Vec<String>,

    },

    /// say you'll be building soon, so others leave you room

    Expect {

        #[pound(short, long)]

        size: Option<String>,

        #[pound(long)]

        ram: Option<Bytes>,

        #[pound(long)]

        cores: Option<Amount<u32>>,

        /// changed your mind

        #[pound(long)]

        done: bool,

        #[pound(long, hidden, env = "CLAUDE_CODE_SESSION_ID")]

        session: Option<String>,

    },

    /// who's running, who's waiting

    Status {

        #[pound(long)]

        json: bool,

    },

    /// a size guess from past runs

    Suggest {

        #[pound(trailing)]

        cmd: Vec<String>,

    },

    /// stop a lease's build

    Kill { lease: u64 },

    /// buy an overdue build more time

    Extend { lease: u64, by: Dur },

    /// exit when a lease goes overdue or finishes

    Watch { lease: u64 },

    /// build output and how much of it there is

    Artifacts {

        #[pound(subcommand)]

        action: Artifacts,

    },

    /// claude code hook entrypoint, json on stdin

    Hook {

        /// pre-tool-use post-tool-use prompt-submit stop subagent-stop session-end

        event: String,

    },

    #[pound(hidden)]

    Daemon,

}

#[derive(Parse)]

pub enum Artifacts {

    /// sizes against the caps, exits 1 when over

    Check {

        /// re-measure instead of trusting the cache

        #[pound(long)]

        fresh: bool,

    },

    /// everything tracked, biggest first

    List,

    /// delete idle artifacts

    Gc {

        /// keep going, oldest first, until under the caps

        #[pound(long)]

        cap: bool,

        #[pound(short = 'n', long)]

        dry_run: bool,

    },

    /// never gc this

    Pin {

        path: PathBuf,

    },

    Unpin {

        path: PathBuf,

    },

}

fn artifacts(cfg: &Config, action: Artifacts) -> i32 {

    let cwd = std::env::current_dir().unwrap_or_else(|_| "/".into());

    let root = paths::project_root(&cwd);

    match action {

        Artifacts::Check { fresh } => {

            let r = artifacts::check(cfg, Some(&root), fresh);

            if let Some((root, size)) = &r.project {

                println!(

                    "{}: {size} of {}",

                    root.display(),

                    cfg.artifacts.project_cap

                );

            }

            println!("everything: {} of {}", r.global, cfg.artifacts.global_cap);

            if r.over() {

                println!("over. `rabt artifacts gc --cap`");

                1

            } else {

                0

            }

        }

        Artifacts::List => {

            let r = artifacts::check(cfg, Some(&root), false);

            let mut rows: Vec<_> = r.entries.iter().collect();

            rows.sort_by_key(|(_, e)| std::cmp::Reverse(e.size));

            for (path, e) in rows {

                println!(

                    "{:>10}  {}{}",

                    e.size.to_string(),

                    path.display(),

                    if e.pinned { "  (pinned)" } else { "" }

                );

            }

            0

        }

        Artifacts::Gc { cap, dry_run } => {

            let g = artifacts::gc(cfg, cap, dry_run);

            for p in &g.removed {

                println!(

                    "{}{}",

                    if dry_run { "would remove " } else { "removed " },

                    p.display()

                );

            }

            for f in &g.failed {

                eprintln!("rabt: {f}");

            }

            println!(

                "{} {}",

                if dry_run { "would free" } else { "freed" },

                g.freed

            );

            i32::from(!g.failed.is_empty())

        }

        Artifacts::Pin { path } | Artifacts::Unpin { path } if !path.exists() => {

            eprintln!("rabt: {} doesn't exist", path.display());

            1

        }

        Artifacts::Pin { path } => report(artifacts::pin(&path, true)),

        Artifacts::Unpin { path } => report(artifacts::pin(&path, false)),

    }

}

fn report(r: Result<(), String>) -> i32 {

    r.map_or_else(

        |e| {

            eprintln!("rabt: {e}");

            1

        },

        |()| 0,

    )

}

fn ask(req: &Request) -> io::Result<Option<Event>> {

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

        return Ok(None);

    };

    proto::send(&mut s, req)?;

    Reader::new(&s)?.recv()

}

fn reply(ev: io::Result<Option<Event>>) -> i32 {

    match ev {

        Ok(Some(Event::Ok)) => 0,

        Ok(Some(Event::Err { msg })) => {

            eprintln!("rabt: {msg}");

            1

        }

        Ok(None) => {

            eprintln!("rabt: nothing running");

            1

        }

        Ok(Some(_)) => 1,

        Err(e) => {

            eprintln!("rabt: {e}");

            1

        }

    }

}

fn expect(cfg: &Config, session: Option<String>, done: bool, a: &run::Args) -> i32 {

    let Some(session) = session.filter(|s| !s.is_empty()) else {

        eprintln!("rabt: no CLAUDE_CODE_SESSION_ID, so nobody to expect on behalf of");

        return 2;

    };

    let need = if done {

        None

    } else {

        match run::need(cfg, a) {

            Ok(n) => Some(n),

            Err(e) => {

                eprintln!("rabt: {e}");

                return 2;

            }

        }

    };

    let mut s = match client::connect_or_spawn() {

        Ok(s) => s,

        Err(e) => {

            eprintln!("rabt: can't reach the daemon: {e}");

            return 1;

        }

    };

    if let Err(e) = proto::send(&mut s, &Request::Expect { session, need }) {

        eprintln!("rabt: {e}");

        return 1;

    }

    reply(Reader::new(&s).and_then(|mut r| r.recv()))

}

fn status(json: bool) -> i32 {

    let Ok(Some(Event::Status(st))) = ask(&Request::Status) else {

        if json {

            println!("null");

        } else {

            println!("idle");

        }

        return 0;

    };

    if json {

        println!("{}", serde_json::to_string(&st).unwrap_or_default());

        return 0;

    }

    println!(

        "ram {} free of {}, cores {} free of {}",

        st.ram_free, st.ram_budget, st.cores_free, st.cores

    );

    for l in &st.leases {

        let state = if l.running {

            format!(

                "running {}{}",

                Dur::secs(l.elapsed),

                if l.overdue { " OVERDUE" } else { "" }

            )

        } else {

            format!(

                "#{} {}",

                l.position,

                l.waiting.map(|r| r.to_string()).unwrap_or_default()

            )

        };

        let cores = if l.granted > l.need.cores {

            format!("{}+{}", l.need.cores, l.granted - l.need.cores)

        } else {

            l.need.cores.to_string()

        };

        println!(

            "{:>4} {}{} {} ram {} cores {} rss {}  {}",

            l.id,

            if l.heavy { "H " } else { "" },

            state,

            Dur(l.need.time),

            l.need.ram,

            cores,

            l.rss,

            l.cmd.join(" ")

        );

    }

    for i in &st.intents {

        println!(

            "     expect {} ram {} cores {}, {} left",

            &i.session[..i.session.len().min(8)],

            i.need.ram,

            i.need.cores,

            Dur::secs(i.expires_in)

        );

    }

    0

}

fn suggest(cfg: &Config, cmd: &[String]) -> i32 {

    let cwd = std::env::current_dir().unwrap_or_else(|_| "/".into());

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

    let Some(rec) = history::lookup(&key) else {

        println!("never seen it, pick a --size");

        return 1;

    };

    let base = match run::base(cfg) {

        Ok(b) => b,

        Err(e) => {

            eprintln!("rabt: {e}");

            return 2;

        }

    };

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

    println!(

        "--ram {} --cores {} --time {} --disk {}",

        n.ram,

        n.cores,

        Dur(n.time),

        n.disk

    );

    0

}

fn watch(lease: u64) -> i32 {

    let Some(mut s): Option<UnixStream> = client::connect() else {

        return 0;

    };

    if proto::send(&mut s, &Request::Watch { lease }).is_err() {

        return 0;

    }

    if matches!(

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

        Ok(Some(Event::Fired { overdue: true }))

    ) {

        println!("lease {lease} is overdue, `rabt kill {lease}` or `rabt extend {lease} <dur>`");

        3

    } else {

        println!("lease {lease} is done");

        0

    }

}

pub fn dispatch(cfg: &Config, cmd: Rabt) -> i32 {

    match cmd {

        Rabt::Run {

            size,

            ram,

            disk,

            cores,

            time,

            deadline,

            max_wait,

            session,

            exact,

            cmd,

        } => run::run(

            cfg,

            &run::Args {

                size,

                ram,

                disk,

                cores,

                time,

                deadline,

                max_wait,

                session,

                exact,

                cmd,

            },

        ),

        Rabt::Expect {

            size,

            ram,

            cores,

            done,

            session,

        } => expect(

            cfg,

            session,

            done,

            &run::Args {

                size,

                ram,

                cores,

                ..run::Args::default()

            },

        ),

        Rabt::Status { json } => status(json),

        Rabt::Suggest { cmd } => suggest(cfg, &cmd),

        Rabt::Kill { lease } => reply(ask(&Request::Kill { lease })),

        Rabt::Extend { lease, by } => reply(ask(&Request::Extend { lease, by })),

        Rabt::Watch { lease } => watch(lease),

        Rabt::Artifacts { action } => artifacts(cfg, action),

        Rabt::Hook { event } => hook::main(cfg, &event),

        Rabt::Daemon => daemon::main(Config::load().unwrap_or_default()),

    }

}