use anyhow::{Result, anyhow};
use serde_json::{Value, json};
use tokio::io::{AsyncBufReadExt, AsyncRead, AsyncWrite, AsyncWriteExt, BufReader};

use crate::config::Paths;
use crate::daemon::request_or_start;
use crate::protocol::{HarnessConfig, Request, Response};

pub async fn serve_stdio(paths: Paths) -> Result<()> {
    serve(tokio::io::stdin(), tokio::io::stdout(), paths).await
}

pub async fn serve<R, W>(reader: R, mut writer: W, paths: Paths) -> Result<()>
where
    R: AsyncRead + Unpin,
    W: AsyncWrite + Unpin,
{
    let mut lines = BufReader::new(reader).lines();
    while let Some(line) = lines.next_line().await? {
        if line.trim().is_empty() {
            continue;
        }
        let req: Value = serde_json::from_str(&line)?;
        let resp = handle_jsonrpc(&paths, req).await;
        writer
            .write_all(format!("{}\n", serde_json::to_string(&resp)?).as_bytes())
            .await?;
        writer.flush().await?;
    }
    Ok(())
}

async fn handle_jsonrpc(paths: &Paths, req: Value) -> Value {
    let id = req.get("id").cloned().unwrap_or(Value::Null);
    let method = req.get("method").and_then(Value::as_str).unwrap_or("");
    let result = match method {
        "initialize" => Ok(json!({
            "protocolVersion": "2024-11-05",
            "capabilities": { "tools": {} },
            "serverInfo": { "name": "multi-brain", "version": env!("CARGO_PKG_VERSION") }
        })),
        "tools/list" => Ok(json!({ "tools": tools() })),
        "tools/call" => {
            handle_tool_call(paths, req.get("params").cloned().unwrap_or_default()).await
        }
        "notifications/initialized" => return json!({ "jsonrpc": "2.0", "id": id, "result": {} }),
        _ => Err(anyhow!("unknown MCP method '{method}'")),
    };

    match result {
        Ok(result) => json!({ "jsonrpc": "2.0", "id": id, "result": result }),
        Err(err) => json!({
            "jsonrpc": "2.0",
            "id": id,
            "error": { "code": -32000, "message": err.to_string() }
        }),
    }
}

async fn handle_tool_call(paths: &Paths, params: Value) -> Result<Value> {
    let name = params
        .get("name")
        .and_then(Value::as_str)
        .ok_or_else(|| anyhow!("missing tool name"))?;
    let args = params
        .get("arguments")
        .cloned()
        .unwrap_or_else(|| json!({}));
    let response = match name {
        "list_brains" => request_or_start(paths, Request::List).await?,
        "register_brain" => {
            let config: HarnessConfig = serde_json::from_value(args)?;
            request_or_start(paths, Request::Register { config }).await?
        }
        "spawn_brain" => {
            let brain = arg_string(&args, "name")?;
            request_or_start(paths, Request::Spawn { name: brain }).await?
        }
        "ensure_brain" => {
            let config: HarnessConfig = serde_json::from_value(args)?;
            request_or_start(paths, Request::Ensure { config }).await?
        }
        "ask_brain" => {
            request_or_start(
                paths,
                Request::Ask {
                    from: arg_string(&args, "from")?,
                    to: arg_string(&args, "to")?,
                    body: arg_string(&args, "body")?,
                },
            )
            .await?
        }
        "brain_events" => {
            request_or_start(
                paths,
                Request::Events {
                    name: arg_string(&args, "name")?,
                    max: args.get("max").and_then(Value::as_u64).unwrap_or(100) as usize,
                },
            )
            .await?
        }
        "inject_multi_brain_mcp" => {
            request_or_start(
                paths,
                Request::InjectMcp {
                    cwd: arg_string(&args, "cwd")?,
                    server_name: args
                        .get("server_name")
                        .and_then(Value::as_str)
                        .unwrap_or("multi-brain")
                        .to_string(),
                },
            )
            .await?
        }
        "stop_brain" => {
            let brain = arg_string(&args, "name")?;
            request_or_start(paths, Request::Stop { name: brain }).await?
        }
        _ => return Err(anyhow!("unknown multi-brain tool '{name}'")),
    };
    Ok(content(response))
}

fn arg_string(args: &Value, key: &str) -> Result<String> {
    args.get(key)
        .and_then(Value::as_str)
        .map(ToString::to_string)
        .ok_or_else(|| anyhow!("missing string argument '{key}'"))
}

fn content(response: Response) -> Value {
    json!({
        "content": [{
            "type": "text",
            "text": serde_json::to_string_pretty(&response).unwrap()
        }]
    })
}

fn tools() -> Value {
    json!([
        {
            "name": "list_brains",
            "description": "List registered harnesses and whether each one is alive.",
            "inputSchema": { "type": "object", "properties": {} }
        },
        {
            "name": "register_brain",
            "description": "Register or update a harness personality without starting it.",
            "inputSchema": harness_schema()
        },
        {
            "name": "spawn_brain",
            "description": "Start a registered harness and keep it alive in the daemon.",
            "inputSchema": named_schema()
        },
        {
            "name": "ensure_brain",
            "description": "Register or update a harness and start it if it is not alive.",
            "inputSchema": harness_schema()
        },
        {
            "name": "ask_brain",
            "description": "Send an advice request to a running harness over the global bus.",
            "inputSchema": {
                "type": "object",
                "properties": {
                    "from": { "type": "string" },
                    "to": { "type": "string" },
                    "body": { "type": "string" }
                },
                "required": ["from", "to", "body"]
            }
        },
        {
            "name": "brain_events",
            "description": "Read recent structured stdin/stdout/stderr/status events for a harness.",
            "inputSchema": {
                "type": "object",
                "properties": {
                    "name": { "type": "string" },
                    "max": { "type": "number" }
                },
                "required": ["name"]
            }
        },
        {
            "name": "inject_multi_brain_mcp",
            "description": "Write a .mcp.json entry that points a repository back to this multi-brain daemon.",
            "inputSchema": {
                "type": "object",
                "properties": {
                    "cwd": { "type": "string" },
                    "server_name": { "type": "string" }
                },
                "required": ["cwd"]
            }
        },
        {
            "name": "stop_brain",
            "description": "Stop a running harness.",
            "inputSchema": named_schema()
        }
    ])
}

fn named_schema() -> Value {
    json!({
        "type": "object",
        "properties": { "name": { "type": "string" } },
        "required": ["name"]
    })
}

fn harness_schema() -> Value {
    json!({
        "type": "object",
        "properties": {
            "name": { "type": "string" },
            "kind": { "type": "string", "enum": ["claude", "codex", "open_claw", "generic"] },
            "command": { "type": "string" },
            "args": { "type": "array", "items": { "type": "string" } },
            "cwd": { "type": "string" },
            "personality": { "type": "string" },
            "env": {
                "type": "array",
                "items": {
                    "type": "array",
                    "prefixItems": [{ "type": "string" }, { "type": "string" }],
                    "minItems": 2,
                    "maxItems": 2
                }
            }
        },
        "required": ["name", "kind", "command", "cwd", "personality"]
    })
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::daemon::{Daemon, request};
    use crate::protocol::Response;
    use std::os::unix::fs::PermissionsExt;
    use tokio::time::{Duration, sleep};

    #[tokio::test]
    async fn initialize_returns_server_info() {
        let paths = Paths::from_root(tempfile::tempdir().unwrap().path().to_path_buf());
        let req = json!({"jsonrpc":"2.0","id":1,"method":"initialize"});
        let resp = handle_jsonrpc(&paths, req).await;
        assert_eq!(resp["result"]["serverInfo"]["name"], "multi-brain");
    }

    #[tokio::test]
    async fn list_tools_includes_ask_brain() {
        let paths = Paths::from_root(tempfile::tempdir().unwrap().path().to_path_buf());
        let req = json!({"jsonrpc":"2.0","id":1,"method":"tools/list"});
        let resp = handle_jsonrpc(&paths, req).await;
        assert!(
            resp["result"]["tools"]
                .as_array()
                .unwrap()
                .iter()
                .any(|t| t["name"] == "ask_brain")
        );
    }

    #[tokio::test]
    async fn list_tools_includes_observability_and_injection_tools() {
        let paths = Paths::from_root(tempfile::tempdir().unwrap().path().to_path_buf());
        let req = json!({"jsonrpc":"2.0","id":1,"method":"tools/list"});
        let resp = handle_jsonrpc(&paths, req).await;
        let tools = resp["result"]["tools"].as_array().unwrap();
        assert!(tools.iter().any(|t| t["name"] == "brain_events"));
        assert!(tools.iter().any(|t| t["name"] == "inject_multi_brain_mcp"));
    }

    #[tokio::test]
    async fn jsonrpc_handles_notifications_and_unknown_methods() {
        let paths = Paths::from_root(tempfile::tempdir().unwrap().path().to_path_buf());
        let initialized = handle_jsonrpc(
            &paths,
            json!({"jsonrpc":"2.0","id":2,"method":"notifications/initialized"}),
        )
        .await;
        assert_eq!(initialized["result"], json!({}));

        let unknown = handle_jsonrpc(&paths, json!({"jsonrpc":"2.0","id":3,"method":"wat"})).await;
        assert!(
            unknown["error"]["message"]
                .as_str()
                .unwrap()
                .contains("unknown MCP method")
        );
    }

    #[tokio::test]
    async fn tool_call_validates_tool_names_and_arguments_before_daemon_io() {
        let paths = Paths::from_root(tempfile::tempdir().unwrap().path().to_path_buf());
        let missing = handle_jsonrpc(
            &paths,
            json!({"jsonrpc":"2.0","id":1,"method":"tools/call","params":{}}),
        )
        .await;
        assert!(
            missing["error"]["message"]
                .as_str()
                .unwrap()
                .contains("missing tool name")
        );

        let unknown = handle_jsonrpc(
            &paths,
            json!({"jsonrpc":"2.0","id":1,"method":"tools/call","params":{"name":"nope","arguments":{}}}),
        )
        .await;
        assert!(
            unknown["error"]["message"]
                .as_str()
                .unwrap()
                .contains("unknown multi-brain tool")
        );

        let bad_args = handle_jsonrpc(
            &paths,
            json!({"jsonrpc":"2.0","id":1,"method":"tools/call","params":{"name":"spawn_brain","arguments":{}}}),
        )
        .await;
        assert!(
            bad_args["error"]["message"]
                .as_str()
                .unwrap()
                .contains("missing string argument")
        );
    }

    #[tokio::test]
    async fn serve_skips_blank_lines_and_writes_json_response() {
        let paths = Paths::from_root(tempfile::tempdir().unwrap().path().to_path_buf());
        let input = tokio::io::BufReader::new(
            b"\n{\"jsonrpc\":\"2.0\",\"id\":1,\"method\":\"initialize\"}\n".as_slice(),
        );
        let mut output = Vec::new();
        serve(input, &mut output, paths).await.unwrap();
        let text = String::from_utf8(output).unwrap();
        assert_eq!(text.lines().count(), 1);
        assert!(text.contains("multi-brain"));
    }

    #[tokio::test]
    async fn tool_calls_round_trip_through_real_daemon() {
        let dir = tempfile::tempdir().unwrap();
        let home = dir.path().join("home");
        let cwd = dir.path().join("repo");
        tokio::fs::create_dir_all(&cwd).await.unwrap();
        let paths = Paths::from_root(home);
        let daemon = Daemon::new(paths.clone()).await.unwrap();
        let server = tokio::spawn(async move { daemon.serve().await.unwrap() });
        wait_for_daemon(&paths).await;

        let bin = dir.path().join("agent");
        std::fs::write(
            &bin,
            r#"#!/usr/bin/env bash
echo ready
while IFS= read -r line; do
  echo "reply:$line"
done
"#,
        )
        .unwrap();
        let mut perms = std::fs::metadata(&bin).unwrap().permissions();
        perms.set_mode(0o755);
        std::fs::set_permissions(&bin, perms).unwrap();

        let config = json!({
            "name": "mcp-unit",
            "kind": "generic",
            "command": bin.display().to_string(),
            "args": [],
            "cwd": cwd.display().to_string(),
            "personality": "p",
            "env": []
        });
        let register = call_tool(&paths, "register_brain", config.clone()).await;
        assert!(register.to_string().contains("registered"));
        let spawn = call_tool(&paths, "spawn_brain", json!({"name":"mcp-unit"})).await;
        assert!(spawn.to_string().contains("spawned"));
        let list = call_tool(&paths, "list_brains", json!({})).await;
        assert!(list.to_string().contains("mcp-unit"));
        let ask = call_tool(
            &paths,
            "ask_brain",
            json!({"from":"tester","to":"mcp-unit","body":"hello"}),
        )
        .await;
        assert!(ask.to_string().contains("hello"));
        let events = call_tool(&paths, "brain_events", json!({"name":"mcp-unit","max":20})).await;
        assert!(events.to_string().contains("stdin"));
        let inject = call_tool(
            &paths,
            "inject_multi_brain_mcp",
            json!({"cwd":cwd.display().to_string(),"server_name":"multi-brain"}),
        )
        .await;
        assert!(inject.to_string().contains("injected_mcp"));
        let stop = call_tool(&paths, "stop_brain", json!({"name":"mcp-unit"})).await;
        assert!(stop.to_string().contains("stopped"));

        server.abort();
    }

    async fn call_tool(paths: &Paths, name: &str, arguments: Value) -> Value {
        let response = handle_jsonrpc(
            paths,
            json!({"jsonrpc":"2.0","id":1,"method":"tools/call","params":{"name":name,"arguments":arguments}}),
        )
        .await;
        assert!(response.get("error").is_none(), "{response}");
        response
    }

    async fn wait_for_daemon(paths: &Paths) {
        tokio::time::timeout(Duration::from_secs(5), async {
            loop {
                if matches!(
                    request(paths, crate::protocol::Request::Ping).await,
                    Ok(Response::Pong)
                ) {
                    return;
                }
                sleep(Duration::from_millis(10)).await;
            }
        })
        .await
        .unwrap();
    }
}
