//! App-owned subagents, the way T3 Code does them: the agent calls
//! `delegate_task`, and the app starts a child thread for the task, shows it
//! working and hands its result back. Any agent can run any task.
//!
//! The tools reach each agent its own way. Claude and Pi call an MCP server
//! the app runs on loopback: Claude through `--mcp-config`, Pi through an
//! extension the app writes and loads with `--extension`. That extension
//! also switches off Pi's other subagent tools before each run. jcode runs
//! them as session tools registered over its API, with its swarm switched
//! off. Nothing needs installing or setting up.
//!
//! Each thread's agent gets a token of its own, which tells the server who's
//! calling. Calls go to the workspace on the UI thread, which answers them.

use std::{
    collections::HashMap,
    io::{BufRead, BufReader, Read, Write},
    net::{TcpListener, TcpStream},
    path::PathBuf,
    sync::{Mutex, OnceLock},
    thread,
    time::Duration,
};

use serde_json::{Value, json};

/// The MCP server's name, as Claude's tool names carry it.
pub const SERVER: &str = "vectrum-code";
pub const DELEGATE: &str = "delegate_task";
pub const TASK_STATUS: &str = "task_status";
pub const TASK_CANCEL: &str = "task_cancel";
pub const CAPABILITIES: &str = "orchestrator_capabilities";

/// The environment that tells the Pi extension where the server is.
const URL_ENV: &str = "VECTRUM_CODE_MCP_URL";
const TOKEN_ENV: &str = "VECTRUM_CODE_MCP_TOKEN";

/// What the tools are for, for agents whose instructions the app adds to:
/// the MCP server's instructions (Claude), the system prompt (Pi) and a
/// reminder with the first message (jcode).
pub const INSTRUCTIONS: &str = "\
## Vectrum Code subagents

Vectrum Code runs this conversation and offers app-owned subagents through its `delegate_task`, `task_status`, `task_cancel` and `orchestrator_capabilities` tools. Tool names may carry a prefix, such as `mcp__vectrum-code__delegate_task`; they work the same.

- `delegate_task` starts a child agent of this thread, working in the same folder. It sees only the task you give it, none of this conversation, so make the task self-contained: the goal, the files and facts that matter, and what to report back. The user watches it work in Vectrum Code's subagents panel.
- The child runs this thread's agent and model unless `target` names others. `orchestrator_capabilities` lists every agent (provider instance) and model the app has, so a task can go to another agent entirely.
- Use `delegate_task` for subagent work; other subagent tools may be switched off. For independent pieces of work, start several tasks so they run in parallel.
- `mode` defaults to `async`: the call returns a `taskId` at once, and the child's result arrives in this thread as a message once it finishes. Carry on with other work, or end your turn instead of polling. Use `task_status` only when you need a result mid-turn. `mode: \"wait\"` blocks until the child is done or `timeoutMs` runs out; the child keeps working after a timeout.
- `task_cancel` stops a task. For a further round of work, call `delegate_task` again with the full context: the brief, what came back and what's still open.";

/// What the tools take, as MCP describes tools.
pub fn tool_definitions() -> Value {
    json!([
        {
            "name": DELEGATE,
            "description": "Delegate one task to a subagent: a child agent of this thread that Vectrum Code runs in a thread of its own, in this thread's working folder. It gets only the task text, not this conversation, so make it self-contained. The provider instance and model default to this thread's; orchestrator_capabilities lists the others. mode 'async' (the default) returns a taskId at once, and the result arrives in this thread as a message when the child finishes; end your turn rather than polling. mode 'wait' blocks until the child finishes or timeoutMs passes (the child carries on after a timeout; its result then arrives as a message).",
            "inputSchema": {
                "type": "object",
                "properties": {
                    "task": {"type": "string", "description": "The complete, self-contained task for the child agent."},
                    "title": {"type": "string", "description": "A short title for the subagent, shown to the user."},
                    "target": {
                        "type": "object",
                        "description": "Which agent and model run the task. Omitted fields inherit this thread's.",
                        "properties": {
                            "providerInstanceId": {"type": "string", "description": "A provider instance id from orchestrator_capabilities."},
                            "model": {"type": "string", "description": "A model id that instance offers."}
                        }
                    },
                    "mode": {"type": "string", "enum": ["async", "wait"], "description": "async (default) returns at once; wait blocks for the result."},
                    "timeoutMs": {"type": "number", "description": "How long mode 'wait' waits. The child isn't stopped when it runs out."}
                },
                "required": ["task"]
            },
            "annotations": {"title": "Delegate a task", "destructiveHint": true, "openWorldHint": true}
        },
        {
            "name": TASK_STATUS,
            "description": "Where a delegated task stands: running, completed, failed or cancelled, and once it's done, its result.",
            "inputSchema": {
                "type": "object",
                "properties": {"taskId": {"type": "string"}},
                "required": ["taskId"]
            },
            "annotations": {"title": "Get a task's status", "readOnlyHint": true}
        },
        {
            "name": TASK_CANCEL,
            "description": "Stop a delegated task that's still running. Its result so far stays readable with task_status.",
            "inputSchema": {
                "type": "object",
                "properties": {"taskId": {"type": "string"}, "reason": {"type": "string"}},
                "required": ["taskId"]
            },
            "annotations": {"title": "Cancel a task", "destructiveHint": true}
        },
        {
            "name": CAPABILITIES,
            "description": "The agents (provider instances) and models Vectrum Code can run a delegated task with, and which ones this thread uses.",
            "inputSchema": {"type": "object", "properties": {}},
            "annotations": {"title": "List agents and models", "readOnlyHint": true}
        }
    ])
}

/// Whether a tool call, by name, is one of the app's: plain for Pi and
/// jcode, prefixed for Claude (`mcp__vectrum-code__delegate_task`).
pub fn own_tool(name: &str) -> Option<&'static str> {
    let bare = name.strip_prefix("mcp__vectrum-code__").or_else(|| name.strip_prefix("mcp__vectrum_code__")).unwrap_or(name);
    [DELEGATE, TASK_STATUS, TASK_CANCEL, CAPABILITIES].into_iter().find(|tool| *tool == bare)
}

/// A call from an agent, for the workspace to answer.
pub struct ToolRequest {
    /// The calling thread's id.
    pub thread: String,
    pub tool: String,
    pub args: Value,
    /// Takes the answer: the result, or what went wrong.
    pub reply: smol::channel::Sender<Result<Value, String>>,
}

struct Server {
    url: String,
    /// Thread id by token, and token by thread id.
    tokens: Mutex<(HashMap<String, String>, HashMap<String, String>)>,
}

static LISTENER: OnceLock<Option<Server>> = OnceLock::new();
static REQUESTS: OnceLock<(smol::channel::Sender<ToolRequest>, smol::channel::Receiver<ToolRequest>)> = OnceLock::new();

fn channel() -> &'static (smol::channel::Sender<ToolRequest>, smol::channel::Receiver<ToolRequest>) {
    REQUESTS.get_or_init(smol::channel::unbounded)
}

/// The calls agents make, for the workspace to answer.
pub fn requests() -> smol::channel::Receiver<ToolRequest> {
    channel().1.clone()
}

/// Starts the server the first time it's needed. `None` if it can't listen.
fn server() -> Option<&'static Server> {
    LISTENER
        .get_or_init(|| {
            let listener = TcpListener::bind(("127.0.0.1", 0)).ok()?;
            let port = listener.local_addr().ok()?.port();
            thread::Builder::new()
                .name("mcp-listener".into())
                .spawn(move || {
                    for stream in listener.incoming() {
                        let Ok(stream) = stream else { continue };
                        let _ = thread::Builder::new().name("mcp-connection".into()).spawn(move || serve(stream));
                    }
                })
                .ok()?;
            Some(Server { url: format!("http://127.0.0.1:{port}/mcp"), tokens: Mutex::default() })
        })
        .as_ref()
}

/// The server's address and the token a thread's agent calls it with.
pub fn credentials(thread: &str) -> Option<(String, String)> {
    let server = server()?;
    let mut tokens = server.tokens.lock().unwrap_or_else(|poisoned| poisoned.into_inner());
    let token = match tokens.1.get(thread) {
        Some(token) => token.clone(),
        None => {
            let token = crate::remote::new_key();
            tokens.0.insert(token.clone(), thread.to_string());
            tokens.1.insert(thread.to_string(), token.clone());
            token
        }
    };
    Some((server.url.clone(), token))
}

fn thread_for(token: &str) -> Option<String> {
    let server = server()?;
    let tokens = server.tokens.lock().unwrap_or_else(|poisoned| poisoned.into_inner());
    tokens.0.get(token).cloned()
}

/// Claude's `--mcp-config` for a thread.
pub fn claude_config(thread: &str) -> Option<String> {
    let (url, token) = credentials(thread)?;
    let config = json!({
        "mcpServers": {
            SERVER: {
                "type": "http",
                "url": url,
                "headers": {"Authorization": format!("Bearer {token}")},
                // Long enough for a `mode: "wait"` call.
                "timeout": 3_900_000,
            }
        }
    });
    Some(config.to_string())
}

/// Answers one connection's requests, one at a time, for as long as it stays
/// open. Plain HTTP/1.1 with JSON replies: all of MCP's streamable HTTP
/// transport a tools-only server needs.
fn serve(stream: TcpStream) {
    let _ = stream.set_nodelay(true);
    let Ok(mut writer) = stream.try_clone() else { return };
    let mut reader = BufReader::new(stream);
    loop {
        let mut request_line = String::new();
        match reader.read_line(&mut request_line) {
            Ok(0) | Err(_) => return,
            Ok(_) => {}
        }
        let method = request_line.split_whitespace().next().unwrap_or_default().to_string();
        let mut length = 0usize;
        let mut token = None;
        loop {
            let mut line = String::new();
            if reader.read_line(&mut line).unwrap_or(0) == 0 {
                return;
            }
            let line = line.trim_end();
            if line.is_empty() {
                break;
            }
            let Some((name, value)) = line.split_once(':') else { continue };
            let value = value.trim();
            match name.trim().to_ascii_lowercase().as_str() {
                "content-length" => length = value.parse().unwrap_or(0),
                "authorization" => token = value.strip_prefix("Bearer ").map(str::to_string),
                _ => {}
            }
        }
        let mut body = vec![0; length.min(16 << 20)];
        if reader.read_exact(&mut body).is_err() {
            return;
        }
        let (status, reply) = if method != "POST" {
            // No server-sent event stream: everything's answered in the reply.
            ("405 Method Not Allowed", None)
        } else {
            match token.as_deref().and_then(thread_for) {
                None => ("401 Unauthorized", Some(json!({"error": "unknown token"}))),
                Some(thread) => match serde_json::from_slice::<Value>(&body) {
                    Ok(message) => match handle(&thread, &message) {
                        Some(reply) => ("200 OK", Some(reply)),
                        // A notification.
                        None => ("202 Accepted", None),
                    },
                    Err(err) => ("400 Bad Request", Some(json!({"error": err.to_string()}))),
                },
            }
        };
        let body = reply.map(|reply| reply.to_string()).unwrap_or_default();
        let head = format!(
            "HTTP/1.1 {status}\r\nContent-Type: application/json\r\nContent-Length: {}\r\n\r\n",
            body.len()
        );
        if writer.write_all(head.as_bytes()).and_then(|_| writer.write_all(body.as_bytes())).is_err() {
            return;
        }
        let _ = writer.flush();
    }
}

/// Answers one JSON-RPC message; `None` for a notification.
fn handle(thread: &str, message: &Value) -> Option<Value> {
    let id = message.get("id").cloned()?;
    let params = message.get("params").cloned().unwrap_or(Value::Null);
    let result = match message.get("method").and_then(Value::as_str).unwrap_or_default() {
        "initialize" => Ok(json!({
            "protocolVersion": params.get("protocolVersion").and_then(Value::as_str).unwrap_or("2025-06-18"),
            "capabilities": {"tools": {"listChanged": false}},
            "serverInfo": {"name": SERVER, "version": env!("CARGO_PKG_VERSION")},
            "instructions": INSTRUCTIONS,
        })),
        "ping" => Ok(json!({})),
        "tools/list" => Ok(json!({"tools": tool_definitions()})),
        "tools/call" => {
            let tool = params.get("name").and_then(Value::as_str).unwrap_or_default().to_string();
            let args = params.get("arguments").cloned().unwrap_or_else(|| json!({}));
            Ok(tool_result(call(thread, tool, args)))
        }
        method => Err(json!({"code": -32601, "message": format!("unknown method {method}")})),
    };
    Some(match result {
        Ok(result) => json!({"jsonrpc": "2.0", "id": id, "result": result}),
        Err(error) => json!({"jsonrpc": "2.0", "id": id, "error": error}),
    })
}

/// Hands a call to the workspace and waits for its answer.
fn call(thread: &str, tool: String, args: Value) -> Result<Value, String> {
    let (reply, answer) = smol::channel::bounded(1);
    let request = ToolRequest { thread: thread.to_string(), tool, args, reply };
    channel().0.send_blocking(request).map_err(|_| "Vectrum Code is shutting down.".to_string())?;
    answer.recv_blocking().unwrap_or_else(|_| Err("Vectrum Code dropped the call.".into()))
}

/// An answer as an MCP tool result: the data as text for the model, and
/// structured for scripts.
pub fn tool_result(outcome: Result<Value, String>) -> Value {
    match outcome {
        Ok(value) => json!({
            "content": [{"type": "text", "text": serde_json::to_string_pretty(&value).unwrap_or_default()}],
            "structuredContent": value,
        }),
        Err(error) => json!({"content": [{"type": "text", "text": error}], "isError": true}),
    }
}

/// The Pi extension: it bridges the app's tools into Pi (when the thread
/// may delegate), adds the instructions to the system prompt, and switches
/// off other subagent tools before every run.
const PI_EXTENSION: &str = r#"// Written by Vectrum Code, which loads it into the Pi agents it runs.
// Changes are overwritten.
const URL = process.env.__URL_ENV__;
const TOKEN = process.env.__TOKEN_ENV__;
const TOOLS = __TOOLS__;
const INSTRUCTIONS = __INSTRUCTIONS__;
const OWN = new Set(TOOLS.map((tool) => tool.name));
// Tools that start agents of their own: Pi's subagent extension and the like.
const OTHER_SUBAGENTS = /subagent|^(task|tasks|agent|agents|swarm|spawn_agent|spawn_agents|delegate)$/i;

async function schema(inputSchema) {
  for (const name of ["typebox", "@sinclair/typebox"]) {
    try {
      const { Type } = await import(name);
      if (typeof Type?.Unsafe === "function") return Type.Unsafe(inputSchema);
    } catch {}
  }
  return inputSchema;
}

let nextId = 1;
async function callTool(name, args, signal) {
  const response = await fetch(URL, {
    method: "POST",
    headers: { "content-type": "application/json", accept: "application/json", authorization: `Bearer ${TOKEN}` },
    body: JSON.stringify({ jsonrpc: "2.0", id: nextId++, method: "tools/call", params: { name, arguments: args ?? {} } }),
    signal,
  });
  const body = await response.text();
  if (!response.ok) throw new Error(`Vectrum Code answered ${response.status}: ${body.slice(0, 400)}`);
  const message = JSON.parse(body);
  if (message.error) throw new Error(message.error.message ?? "Vectrum Code returned an error");
  return message.result;
}

export default async function vectrumCode(pi) {
  // Subagents go through the app, so other subagent tools stay off. Checked
  // before every run, since extensions can turn tools back on.
  const switchOffOthers = () => {
    if (typeof pi.getActiveTools !== "function" || typeof pi.setActiveTools !== "function") return;
    const active = pi.getActiveTools();
    const kept = active.filter((name) => OWN.has(name) || !OTHER_SUBAGENTS.test(name));
    if (kept.length !== active.length) pi.setActiveTools(kept);
  };
  pi.on("session_start", switchOffOthers);
  pi.on("before_agent_start", (event) => {
    switchOffOthers();
    if (URL && TOKEN) return { systemPrompt: event.systemPrompt + "\n\n" + INSTRUCTIONS };
  });
  if (!URL || !TOKEN) return;
  for (const tool of TOOLS) {
    pi.registerTool({
      name: tool.name,
      label: tool.annotations?.title ?? tool.name,
      description: tool.description,
      parameters: await schema(tool.inputSchema),
      async execute(_toolCallId, params, signal) {
        const result = await callTool(tool.name, params, signal);
        const text = (result?.content ?? []).filter((part) => part?.type === "text").map((part) => part.text).join("\n");
        if (result?.isError) throw new Error(text || `${tool.name} failed`);
        return { content: [{ type: "text", text }], details: result?.structuredContent };
      },
    });
  }
}
"#;

/// Writes the Pi extension where Pi can load it, and returns its path.
fn pi_extension() -> Option<PathBuf> {
    static PATH: OnceLock<Option<PathBuf>> = OnceLock::new();
    PATH.get_or_init(|| {
        let source = PI_EXTENSION
            .replace("__URL_ENV__", URL_ENV)
            .replace("__TOKEN_ENV__", TOKEN_ENV)
            .replace("__TOOLS__", &tool_definitions().to_string())
            .replace("__INSTRUCTIONS__", &Value::String(INSTRUCTIONS.into()).to_string());
        let path = crate::settings::data_dir()?.join("pi").join("vectrum-code.js");
        std::fs::create_dir_all(path.parent()?).ok()?;
        if std::fs::read_to_string(&path).ok().as_deref() != Some(source.as_str()) {
            std::fs::write(&path, source).ok()?;
        }
        Some(path)
    })
    .clone()
}

/// Loads the extension into a Pi command. `thread` gets the tools; without
/// it (a subagent's own thread) Pi only has its other subagent tools turned off.
pub fn prepare_pi(command: &mut std::process::Command, thread: Option<&str>) {
    let Some(extension) = pi_extension() else { return };
    command.arg("--extension").arg(extension);
    if let Some((url, token)) = thread.and_then(credentials) {
        command.env(URL_ENV, url).env(TOKEN_ENV, token);
    } else {
        command.env_remove(URL_ENV).env_remove(TOKEN_ENV);
    }
}

/// jcode's session tool configuration: the app's tools (when the thread may
/// delegate) and its own agent-spawning tools off.
pub fn jcode_tools(delegate: bool) -> Value {
    let custom: Vec<Value> = if delegate {
        tool_definitions()
            .as_array()
            .into_iter()
            .flatten()
            .map(|tool| json!({"name": tool["name"], "description": tool["description"], "parameters": tool["inputSchema"]}))
            .collect()
    } else {
        Vec::new()
    };
    // `schedule` can spawn a session too.
    json!({"disabled": ["swarm", "schedule"], "custom": custom})
}

/// How long a call's `mode: "wait"` may block before the agent's own client
/// gives up on it: jcode allows a session tool two minutes, and Node's
/// fetch (Pi) waits five for a reply's headers.
pub fn wait_limit(provider: crate::models::Provider) -> Duration {
    match provider {
        crate::models::Provider::Claude => Duration::from_secs(30 * 60),
        crate::models::Provider::Pi => Duration::from_secs(270),
        crate::models::Provider::Jcode => Duration::from_secs(100),
    }
}

#[cfg(test)]
mod tests {
    use serde_json::{Value, json};

    use super::{DELEGATE, credentials, own_tool, requests};

    fn post(url: &str, token: &str, body: Value) -> (u16, String) {
        let response = ureq::post(url).set("Authorization", &format!("Bearer {token}")).set("Content-Type", "application/json").send_string(&body.to_string());
        match response {
            Ok(response) => (response.status(), response.into_string().unwrap_or_default()),
            Err(ureq::Error::Status(code, response)) => (code, response.into_string().unwrap_or_default()),
            Err(err) => panic!("{err}"),
        }
    }

    #[::core::prelude::v1::test]
    fn serves_mcp_to_the_threads_it_knows() {
        let (url, token) = credentials("thread-a").expect("listens");
        assert_eq!(credentials("thread-a").unwrap().1, token, "a thread keeps its token");
        assert_ne!(credentials("thread-b").unwrap().1, token);

        let (status, _) = post(&url, "nope", json!({"jsonrpc": "2.0", "id": 1, "method": "tools/list"}));
        assert_eq!(status, 401);

        let (status, body) = post(&url, &token, json!({"jsonrpc": "2.0", "id": 1, "method": "initialize", "params": {"protocolVersion": "2025-03-26"}}));
        assert_eq!(status, 200);
        let reply: Value = serde_json::from_str(&body).unwrap();
        assert_eq!(reply["result"]["protocolVersion"], "2025-03-26");
        assert!(reply["result"]["instructions"].as_str().unwrap().contains("delegate_task"));

        let (status, body) = post(&url, &token, json!({"jsonrpc": "2.0", "method": "notifications/initialized"}));
        assert_eq!((status, body.as_str()), (202, ""));

        let (_, body) = post(&url, &token, json!({"jsonrpc": "2.0", "id": 2, "method": "tools/list"}));
        let reply: Value = serde_json::from_str(&body).unwrap();
        assert_eq!(reply["result"]["tools"][0]["name"], DELEGATE);

        // The workspace's side, answering with who called and what.
        let incoming = requests();
        std::thread::spawn(move || {
            while let Ok(request) = incoming.recv_blocking() {
                let outcome = match request.tool.as_str() {
                    "task_status" => Ok(json!({"thread": request.thread, "args": request.args})),
                    _ => Err("no such task".to_string()),
                };
                let _ = request.reply.send_blocking(outcome);
            }
        });
        let call = |name: &str| {
            let body = json!({"jsonrpc": "2.0", "id": 3, "method": "tools/call", "params": {"name": name, "arguments": {"taskId": "t1"}}});
            let (_, body) = post(&url, &token, body);
            serde_json::from_str::<Value>(&body).unwrap()["result"].clone()
        };
        let result = call("task_status");
        assert_eq!(result["structuredContent"], json!({"thread": "thread-a", "args": {"taskId": "t1"}}));
        assert!(result.get("isError").is_none());
        let result = call("task_cancel");
        assert_eq!(result["isError"], true);
        assert_eq!(result["content"][0]["text"], "no such task");
    }

    #[::core::prelude::v1::test]
    fn recognizes_its_tools_under_any_prefix() {
        assert_eq!(own_tool("delegate_task"), Some(DELEGATE));
        assert_eq!(own_tool("mcp__vectrum-code__delegate_task"), Some(DELEGATE));
        assert_eq!(own_tool("mcp__vectrum_code__task_status"), Some("task_status"));
        assert_eq!(own_tool("mcp__other__delegate_task"), None);
        assert_eq!(own_tool("Task"), None);
    }
}