//! `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() as i64;
        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(|m| m.len()).unwrap_or(0);
    let err_len = std::fs::metadata(&err_file).map(|m| m.len()).unwrap_or(0);

    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 {
                if 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:#}"),
    }
}

/// Turn a `DaemonResponse` from a `BashStatus` call into the string
/// claude sees as the tool result.
fn render_bash_status(id: &str, resp: Result<DaemonResponse>) -> String {
    match resp {
        Ok(DaemonResponse::Ok { payload }) => format_task(id, &payload),
        Ok(DaemonResponse::Error { message }) => message,
        Err(e) => format!("bash bridge error: {e:#}"),
    }
}

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

#[derive(Debug, Deserialize, JsonSchema)]
struct BashRunArgs {
    /// Shell command to run (passed to `sh -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>,
}

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(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."
    )]
    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,
        };
        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), 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 req = DaemonRequest::BashStatus {
            id: args.id,
            wait_seconds: args.wait_seconds,
        };
        render_bash_status(&id, 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(())
}