//! `hive-bash-mcp` binary — stdio MCP server claude spawns per turn.
//! Thin protocol bridge: every tool call → connect to the daemon's
//! unix socket → write a JSON request line → read the JSON response →
//! return the result to claude.
//!
//! No subprocess management at this entrypoint — the daemon owns that.
//! Cold-starts in milliseconds.

use anyhow::{Context, Result};
use rmcp::{
    ServerHandler, ServiceExt,
    handler::server::wrapper::Parameters,
    schemars::{self, JsonSchema},
    tool, tool_handler, tool_router,
    transport::stdio,
};
use serde::Deserialize;
use std::fmt::Write as _;
use tokio::io::{AsyncBufReadExt as _, AsyncWriteExt as _, BufReader};
use tokio::net::UnixStream;

use hive_bash_mcp::paths;
use hive_bash_mcp::protocol::{DaemonRequest, DaemonResponse};
use hive_bash_mcp::runner::SUMMARY_BYTES;

/// Send `req` to the daemon and read back the response. Each call is a
/// fresh unix-socket connection — short-lived (single round-trip) so
/// connection pooling is unnecessary.
async fn round_trip(req: DaemonRequest) -> Result<DaemonResponse> {
    let socket = paths::daemon_socket();
    let stream = UnixStream::connect(&socket)
        .await
        .with_context(|| format!("connect bash daemon socket {}", socket.display()))?;
    let (reader, mut writer) = stream.into_split();
    let mut line = serde_json::to_string(&req)?;
    line.push('\n');
    writer
        .write_all(line.as_bytes())
        .await
        .context("write request to bash daemon socket")?;
    writer.shutdown().await.ok();
    let mut buf = String::new();
    BufReader::new(reader)
        .read_line(&mut buf)
        .await
        .context("read response from bash daemon socket")?;
    serde_json::from_str(&buf).context("parse bash daemon response")
}

/// Format a `TaskFile` JSON value as a human-readable status string.
/// Mirrors `format_bash_status` in the old hive-ag3nt, adapted to work
/// from the daemon's JSON payload.
fn format_task(id: &str, task: &serde_json::Value) -> String {
    let status = task["status"].as_str().unwrap_or("unknown");
    let mut out = format!("task `{id}`: status={status}");

    if let Some(code) = task["exit_code"].as_i64() {
        let _ = write!(out, ", exit={code}");
    }
    if let (Some(started), None) = (
        task["started_at"].as_i64(),
        task["completed_at"].as_i64().map(|_| ()),
    ) {
        // Running — show age.
        let now = std::time::SystemTime::now()
            .duration_since(std::time::UNIX_EPOCH)
            .unwrap_or_default()
            .as_secs()
            .cast_signed();
        let _ = write!(out, ", running for {}s", now - started);
    }
    if let (Some(completed), Some(started)) =
        (task["completed_at"].as_i64(), task["started_at"].as_i64())
    {
        let _ = write!(out, ", took {}s", completed - started);
    }

    let out_file = paths::task_out(id);
    let err_file = paths::task_err(id);
    let out_len = std::fs::metadata(&out_file).map_or(0, |m| m.len());
    let err_len = std::fs::metadata(&err_file).map_or(0, |m| m.len());

    if let Some(stdout) = task["stdout_tail"].as_str() {
        let s = stdout.trim();
        if !s.is_empty() {
            let _ = write!(out, "\n\nstdout:\n```\n{s}\n```");
        }
    }
    if out_len > SUMMARY_BYTES as u64 {
        let _ = write!(out, "\n\nFull stdout lives in `{}`", out_file.display());
    }

    if let Some(stderr) = task["stderr_tail"].as_str() {
        let s = stderr.trim();
        if !s.is_empty() {
            let _ = write!(out, "\n\nstderr:\n```\n{s}\n```");
        }
    }
    if err_len > SUMMARY_BYTES as u64 {
        let _ = write!(out, "\n\nFull stderr lives in `{}`", err_file.display());
    }

    out
}

/// Turn a `DaemonResponse` from a `BashRun` call into the string
/// claude sees as the tool result.
fn render_bash_run(id: &str, resp: Result<DaemonResponse>) -> String {
    match resp {
        Ok(DaemonResponse::Ok { payload }) => {
            let finished = payload["finished"].as_bool().unwrap_or(false);
            if finished && let Some(task) = payload.get("task") {
                return format_task(id, task);
            }
            format!("task started: id={id}")
        }
        Ok(DaemonResponse::Error { message }) => format!("bash_run error: {message}"),
        Err(e) => format!("bash bridge error: {e:#}"),
    }
}

/// Appended to a `status` result when the caller parked on a
/// `wait_seconds` poll that expired while the task was still running —
/// nudges the model to spend the idle time on other work instead of
/// immediately re-waiting on the same task.
const BASH_IDLE_WAIT_HINT: &str = "\n\nThe task is still running — your wait timed out before it \
finished. If you have other useful work, do that and check back later (the task keeps running, and \
a wake fires when it completes) rather than immediately re-waiting.";

/// Turn a `DaemonResponse` from a `BashStatus` call into the string
/// claude sees as the tool result. When `waited` is set (the call
/// parked on a `wait_seconds` poll) and the task is still non-terminal,
/// [`BASH_IDLE_WAIT_HINT`] is appended.
fn render_bash_status(id: &str, resp: Result<DaemonResponse>, waited: bool) -> String {
    match resp {
        Ok(DaemonResponse::Ok { payload }) => {
            let mut out = format_task(id, &payload);
            if waited && matches!(payload["status"].as_str(), Some("pending" | "running")) {
                out.push_str(BASH_IDLE_WAIT_HINT);
            }
            out
        }
        Ok(DaemonResponse::Error { message }) => message,
        Err(e) => format!("bash bridge error: {e:#}"),
    }
}

/// Turn a `DaemonResponse` from a `BashKill` call into the string claude sees.
fn render_bash_kill(resp: Result<DaemonResponse>) -> String {
    match resp {
        Ok(DaemonResponse::Ok { payload }) => {
            let id = payload["id"].as_str().unwrap_or("unknown");
            if payload["was_running"].as_bool().unwrap_or(false) {
                let sig = payload["signal"].as_str().unwrap_or("SIGINT");
                format!(
                    "task `{id}`: {sig} sent to its process group; it transitions to `killed` \
                     once the process exits and the usual completion wake fires."
                )
            } else {
                format!("task `{id}` was pending — cancelled before it started.")
            }
        }
        Ok(DaemonResponse::Error { message }) => format!("kill error: {message}"),
        Err(e) => format!("bash bridge error: {e:#}"),
    }
}

// ---------------------------------------------------------------------------
// MCP server
// ---------------------------------------------------------------------------

#[derive(Debug, Deserialize, JsonSchema)]
struct BashRunArgs {
    /// Shell command to run (passed to `bash -c`).
    cmd: String,
    /// Timeout in seconds. Defaults to `None` (no timeout) — task runs until
    /// natural exit. Pass an explicit value to kill the task after N seconds
    /// and mark it `timed_out`.
    #[serde(default)]
    timeout_secs: Option<u64>,
    /// Optional inline wait: `run` polls for up to `wait_seconds`
    /// (capped at 30) before returning. When the task finishes within the
    /// window the full status is returned immediately and no wake is fired;
    /// when the timeout expires the task keeps running and the normal
    /// `task started: id=<id>` response is returned. Defaults to 3s. Pass
    /// `0` to disable inline waiting and always get the immediate response.
    #[serde(default = "default_wait")]
    wait_seconds: Option<u64>,
    /// Optional task name. When set it becomes the task id, so it appears
    /// in the completion wake (`from: "bash-task-<name>"`), in `status`
    /// lookups, and in the loose-ends list — handy for recognising a task
    /// later instead of an opaque hex id. A name can be reused once its
    /// previous task has finished; reusing a name whose task is still
    /// running is rejected. Allowed chars: ASCII letters, digits, `.`,
    /// `_`, `-` (max 64). Omit to get the auto-generated id.
    #[serde(default)]
    name: Option<String>,
}

#[allow(
    clippy::unnecessary_wraps,
    reason = "serde `#[serde(default = ...)]` requires the default fn's return type to match the field's `Option<u64>`, so the `Some` wrap is mandatory even though the value is constant"
)]
fn default_wait() -> Option<u64> {
    Some(3)
}

#[derive(Debug, Deserialize, JsonSchema)]
struct BashStatusArgs {
    /// Task ID returned by `run`.
    id: String,
    /// Optional inline wait: `status` polls for up to `wait_seconds`
    /// (capped at 30) before returning. Useful to avoid a separate
    /// round-trip when the task is expected to finish soon.
    #[serde(default)]
    wait_seconds: Option<u64>,
}

#[derive(Debug, Deserialize, JsonSchema)]
struct BashKillArgs {
    /// Task ID (from `run`) to kill.
    id: String,
    /// `false` (default): SIGINT to the task's process group — graceful, the
    /// process can clean up. `true`: SIGKILL. Either way the whole process
    /// group is signalled, so children of the shell (e.g. a `cargo`/`nix`
    /// invocation) are stopped too, not just the shell. Fire-and-forget: the
    /// call sends the signal and returns without waiting. If a SIGINT'd task
    /// doesn't exit, call `kill` again with `force: true` to send SIGKILL.
    #[serde(default)]
    force: bool,
}

#[derive(Clone)]
struct BashMcp;

#[tool_router]
impl BashMcp {
    #[tool(
        description = "Run a shell command in the background. Returns a task ID immediately — \
            do NOT wait inline. When the command finishes, the harness fires a wake with \
            `from: \"bash-task-<id>\"` and the exit code + last stdout lines in the body; \
            handle it on a future turn. Use `status` to poll the task status within \
            the same turn if needed. `timeout_secs` defaults to `None` (no timeout) — \
            task runs until natural exit; pass an explicit value to kill after N seconds. \
            Pass `wait_seconds` (capped at 30) to wait inline for fast commands: when the \
            task finishes within the window the full status is returned immediately and no \
            wake is fired; when the timeout expires the task keeps running and the normal \
            `task started: id=<id>` response is returned. `wait_seconds` defaults to 3; \
            pass `wait_seconds: 0` to disable inline waiting and always get the immediate \
            response. Pass `name` to label the task with a memorable id (used in the wake \
            `from`, `status` lookups, and the loose-ends list) instead of an opaque hex id; \
            a name is reusable once its prior task has finished, and rejected while one is \
            still running."
    )]
    async fn run(&self, Parameters(args): Parameters<BashRunArgs>) -> String {
        let req = DaemonRequest::BashRun {
            cmd: args.cmd,
            timeout_secs: args.timeout_secs,
            wait_seconds: args.wait_seconds,
            name: args.name,
        };
        let resp = round_trip(req).await;
        // Extract the id from the response to format the result.
        match &resp {
            Ok(DaemonResponse::Ok { payload }) => {
                let id = payload["id"].as_str().unwrap_or("unknown").to_owned();
                render_bash_run(&id, resp)
            }
            _ => render_bash_run("unknown", resp),
        }
    }

    #[tool(
        description = "Check the status of a background bash task by its ID (from `run`). \
            Returns the current status (pending/running/done/timed_out/interrupted/killed), exit \
            code if finished, and a tail of stdout/stderr. Full output lives in \
            `harness/bash-tasks/<id>.out` / `.err`. \
            Pass `wait_seconds` (capped at 30) to wait inline for the task to finish: when the \
            task finishes within the window the full status is returned immediately. Useful to \
            avoid a separate round-trip when the task is expected to finish soon."
    )]
    async fn status(&self, Parameters(args): Parameters<BashStatusArgs>) -> String {
        let id = args.id.clone();
        let waited = args.wait_seconds.is_some_and(|w| w > 0);
        let req = DaemonRequest::BashStatus {
            id: args.id,
            wait_seconds: args.wait_seconds,
        };
        render_bash_status(&id, round_trip(req).await, waited)
    }

    #[tool(
        description = "Kill a background bash task you started (by its ID from `run`). \
            `force: false` (default) sends SIGINT — graceful, lets the process clean up; \
            `force: true` sends SIGKILL. Either way the task's whole process group is \
            signalled, so a runaway child (cargo/nix/etc.) is stopped too, not just the shell. \
            Fire-and-forget: sends the signal and returns without waiting. If a SIGINT'd task \
            doesn't exit, call kill again with `force: true` to SIGKILL. A still-pending task \
            is cancelled before it starts. The task ends as `killed` and fires the usual \
            completion wake."
    )]
    async fn kill(&self, Parameters(args): Parameters<BashKillArgs>) -> String {
        let req = DaemonRequest::BashKill {
            id: args.id,
            force: args.force,
        };
        render_bash_kill(round_trip(req).await)
    }
}

#[tool_handler]
impl ServerHandler for BashMcp {}

#[tokio::main]
async fn main() -> Result<()> {
    tracing_subscriber::fmt()
        .with_env_filter(
            tracing_subscriber::EnvFilter::try_from_env("RUST_LOG")
                .unwrap_or_else(|_| tracing_subscriber::EnvFilter::new("warn")),
        )
        .with_writer(std::io::stderr)
        .init();

    let service = BashMcp.serve(stdio()).await?;
    service.waiting().await?;
    Ok(())
}

#[cfg(test)]
mod status_hint_tests {
    use super::{BASH_IDLE_WAIT_HINT, render_bash_status};
    use hive_bash_mcp::protocol::DaemonResponse;

    fn status_resp(status: &str) -> DaemonResponse {
        DaemonResponse::Ok {
            payload: serde_json::json!({ "status": status, "started_at": 1 }),
        }
    }

    #[test]
    fn running_task_after_wait_appends_idle_hint() {
        let out = render_bash_status("t1", Ok(status_resp("running")), true);
        assert!(out.contains(BASH_IDLE_WAIT_HINT));
    }

    #[test]
    fn running_task_without_wait_has_no_hint() {
        let out = render_bash_status("t1", Ok(status_resp("running")), false);
        assert!(!out.contains(BASH_IDLE_WAIT_HINT));
    }

    #[test]
    fn finished_task_after_wait_has_no_hint() {
        let out = render_bash_status("t1", Ok(status_resp("done")), true);
        assert!(!out.contains(BASH_IDLE_WAIT_HINT));
    }
}