use anyhow::{Context, Result, anyhow};
use serde_json::{Value, json};
use std::sync::Arc;
use tokio::fs::OpenOptions;
use tokio::io::{AsyncBufReadExt, AsyncWriteExt, BufReader};
use tokio::net::{UnixListener, UnixStream};
use tokio::process::Command;
use tokio::sync::Mutex;
use tokio::time::{Duration, sleep};

use crate::config::Paths;
use crate::harness::{ProcessTable, new_process_table, now_ms, spawn_harness, tail_events};
use crate::protocol::{Message, Request, Response};
use crate::state::State;

#[derive(Clone)]
pub struct Daemon {
    paths: Paths,
    state: Arc<Mutex<State>>,
    processes: ProcessTable,
}

impl Daemon {
    pub async fn new(paths: Paths) -> Result<Self> {
        paths.ensure().await?;
        let state = State::load(&paths.state_file).await?;
        Ok(Self {
            paths,
            state: Arc::new(Mutex::new(state)),
            processes: new_process_table(),
        })
    }

    pub async fn serve(self) -> Result<()> {
        if self.paths.socket_file.exists() {
            let _ = tokio::fs::remove_file(&self.paths.socket_file).await;
        }
        let listener = UnixListener::bind(&self.paths.socket_file)
            .with_context(|| format!("bind {}", self.paths.socket_file.display()))?;
        tracing::info!(socket = %self.paths.socket_file.display(), state = %self.paths.state_file.display(), "multi-brain daemon listening");
        loop {
            let (stream, _) = listener.accept().await?;
            let daemon = self.clone();
            tokio::spawn(async move {
                if let Err(err) = daemon.handle_stream(stream).await {
                    tracing::warn!(error = %err, "client handling failed");
                }
            });
        }
    }

    async fn handle_stream(&self, stream: UnixStream) -> Result<()> {
        let (read, mut write) = stream.into_split();
        let mut lines = BufReader::new(read).lines();
        while let Some(line) = lines.next_line().await? {
            let req: Request = serde_json::from_str(&line)?;
            let resp = self.handle(req).await;
            write
                .write_all(format!("{}\n", serde_json::to_string(&resp)?).as_bytes())
                .await?;
        }
        Ok(())
    }

    pub async fn handle(&self, req: Request) -> Response {
        match self.try_handle(req).await {
            Ok(resp) => resp,
            Err(err) => Response::Error {
                message: err.to_string(),
            },
        }
    }

    async fn try_handle(&self, req: Request) -> Result<Response> {
        match req {
            Request::Ping => Ok(Response::Pong),
            Request::Register { config } => {
                tracing::info!(harness = %config.name, kind = ?config.kind, cwd = %config.cwd, "register harness");
                let mut state = self.state.lock().await;
                let id = state.upsert(config);
                state.save(&self.paths.state_file).await?;
                Ok(Response::Registered { id })
            }
            Request::Ensure { config } => {
                let name = config.name.clone();
                {
                    let mut state = self.state.lock().await;
                    state.upsert(config);
                    state.save(&self.paths.state_file).await?;
                }
                self.spawn_by_name(&name).await
            }
            Request::Spawn { name } => self.spawn_by_name(&name).await,
            Request::List => {
                let mut table = self.processes.lock().await;
                let state = self.state.lock().await;
                let names: Vec<_> = state.harnesses.keys().cloned().collect();
                let mut harnesses = Vec::new();
                for name in names {
                    let (pid, alive) = match table.get_mut(&name) {
                        Some(proc) => (proc.pid(), proc.is_alive().await),
                        None => (None, false),
                    };
                    if let Some(info) = state.info_for(&name, pid, alive, self.log_path_for(&name))
                    {
                        harnesses.push(info);
                    }
                }
                Ok(Response::List { harnesses })
            }
            Request::Ask { from, to, body } => {
                let mut table = self.processes.lock().await;
                let proc = table
                    .get_mut(&to)
                    .ok_or_else(|| anyhow!("harness '{to}' is not running"))?;
                proc.write(&body).await?;
                tracing::info!(from = %from, to = %to, bytes = body.len(), "sent message to harness");
                let message = Message {
                    from,
                    to,
                    body,
                    at_unix_ms: now_ms(),
                };
                let mut state = self.state.lock().await;
                state.messages.push(message.clone());
                state.save(&self.paths.state_file).await?;
                Ok(Response::Asked { message })
            }
            Request::Messages { name } => {
                let state = self.state.lock().await;
                let messages = state
                    .messages
                    .iter()
                    .filter(|m| m.from == name || m.to == name)
                    .cloned()
                    .collect();
                Ok(Response::Messages { messages })
            }
            Request::Events { name, max } => {
                let log_path = self.paths.log_dir.join(format!("{name}.jsonl"));
                let events = tail_events(&log_path, max).await?;
                Ok(Response::Events { events })
            }
            Request::InjectMcp { cwd, server_name } => {
                let path = self.inject_mcp_config(&cwd, &server_name).await?;
                Ok(Response::InjectedMcp { path, server_name })
            }
            Request::Stop { name } => {
                let mut table = self.processes.lock().await;
                if let Some(mut proc) = table.remove(&name) {
                    proc.stop().await?;
                }
                tracing::info!(harness = %name, "stopped harness");
                Ok(Response::Stopped { name })
            }
        }
    }

    async fn spawn_by_name(&self, name: &str) -> Result<Response> {
        let config = {
            let state = self.state.lock().await;
            state
                .harnesses
                .get(name)
                .ok_or_else(|| anyhow!("unknown harness '{name}'"))?
                .config
                .clone()
        };

        let mut table = self.processes.lock().await;
        if let Some(proc) = table.get_mut(name) {
            if proc.is_alive().await {
                let state = self.state.lock().await;
                let info = state
                    .info_for(name, proc.pid(), true, proc.log_path().to_string())
                    .unwrap();
                return Ok(Response::Spawned { info });
            }
        }
        let proc = spawn_harness(&config, &self.paths.log_dir).await?;
        let pid = proc.pid();
        let log_path = proc.log_path().to_string();
        table.insert(name.to_string(), proc);
        tracing::info!(harness = %name, pid = ?pid, log_path = %log_path, "spawned harness");
        let state = self.state.lock().await;
        let info = state.info_for(name, pid, true, log_path).unwrap();
        Ok(Response::Spawned { info })
    }

    fn log_path_for(&self, name: &str) -> String {
        self.paths
            .log_dir
            .join(format!("{name}.jsonl"))
            .display()
            .to_string()
    }

    async fn inject_mcp_config(&self, cwd: &str, server_name: &str) -> Result<String> {
        inject_mcp_config(
            cwd,
            &server_name,
            &self.paths.data_dir.display().to_string(),
        )
        .await
    }
}

pub async fn request(paths: &Paths, req: Request) -> Result<Response> {
    let stream = UnixStream::connect(&paths.socket_file)
        .await
        .with_context(|| format!("connect {}", paths.socket_file.display()))?;
    let (read, mut write) = stream.into_split();
    write
        .write_all(format!("{}\n", serde_json::to_string(&req)?).as_bytes())
        .await?;
    let mut lines = BufReader::new(read).lines();
    let line = lines
        .next_line()
        .await?
        .ok_or_else(|| anyhow!("daemon closed connection without a response"))?;
    Ok(serde_json::from_str(&line)?)
}

pub async fn request_or_start(paths: &Paths, req: Request) -> Result<Response> {
    match request(paths, Request::Ping).await {
        Ok(Response::Pong) => {}
        _ => start_daemon_headless(paths).await?,
    }
    request(paths, req).await
}

pub async fn start_daemon_headless(paths: &Paths) -> Result<()> {
    paths.ensure().await?;
    let daemon_log = paths.log_dir.join("daemon.log");
    let stdout = OpenOptions::new()
        .create(true)
        .append(true)
        .open(&daemon_log)
        .await?
        .into_std()
        .await;
    let stderr = stdout.try_clone()?;
    let mut child = Command::new(std::env::current_exe()?)
        .arg("--home")
        .arg(&paths.data_dir)
        .arg("daemon")
        .stdin(std::process::Stdio::null())
        .stdout(std::process::Stdio::from(stdout))
        .stderr(std::process::Stdio::from(stderr))
        .spawn()?;
    tracing::info!(pid = ?child.id(), log_path = %daemon_log.display(), "started daemon headlessly");

    for _ in 0..100 {
        if matches!(request(paths, Request::Ping).await, Ok(Response::Pong)) {
            return Ok(());
        }
        if let Some(status) = child.try_wait()? {
            return Err(anyhow!("daemon exited during startup: {status}"));
        }
        sleep(Duration::from_millis(50)).await;
    }
    Err(anyhow!(
        "daemon did not become ready at {}",
        paths.socket_file.display()
    ))
}

async fn inject_mcp_config(cwd: &str, server_name: &str, home: &str) -> Result<String> {
    let path = std::path::Path::new(cwd).join(".mcp.json");
    let mut config: Value = match tokio::fs::read_to_string(&path).await {
        Ok(raw) => serde_json::from_str(&raw).unwrap_or_else(|_| json!({ "mcpServers": {} })),
        Err(err) if err.kind() == std::io::ErrorKind::NotFound => json!({ "mcpServers": {} }),
        Err(err) => return Err(err.into()),
    };

    if !config.get("mcpServers").is_some_and(Value::is_object) {
        config["mcpServers"] = json!({});
    }
    config["mcpServers"][server_name] = json!({
        "command": std::env::current_exe()?.display().to_string(),
        "args": ["mcp"],
        "env": {
            "MULTI_BRAIN_HOME": home,
            "PATH": std::env::var("PATH").unwrap_or_default()
        }
    });

    tokio::fs::write(
        &path,
        format!("{}\n", serde_json::to_string_pretty(&config)?),
    )
    .await?;
    tracing::info!(path = %path.display(), server_name = %server_name, "injected MCP config");
    Ok(path.display().to_string())
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::protocol::{HarnessConfig, HarnessKind};
    use std::os::unix::fs::PermissionsExt;

    fn mock_config(dir: &tempfile::TempDir, name: &str) -> HarnessConfig {
        let bin = dir.path().join(format!("{name}-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();
        HarnessConfig {
            name: name.into(),
            kind: HarnessKind::Generic,
            command: bin.display().to_string(),
            args: vec![],
            cwd: dir.path().display().to_string(),
            personality: "p".into(),
            env: vec![],
        }
    }

    #[tokio::test]
    async fn daemon_handle_covers_registry_lifecycle() {
        let dir = tempfile::tempdir().unwrap();
        let paths = Paths::from_root(dir.path().join("home"));
        let daemon = Daemon::new(paths.clone()).await.unwrap();
        let config = mock_config(&dir, "brain");

        let registered = daemon
            .handle(Request::Register {
                config: config.clone(),
            })
            .await;
        assert!(matches!(registered, Response::Registered { .. }));

        let spawned = daemon
            .handle(Request::Spawn {
                name: "brain".into(),
            })
            .await;
        assert!(matches!(spawned, Response::Spawned { .. }));

        let asked = daemon
            .handle(Request::Ask {
                from: "tester".into(),
                to: "brain".into(),
                body: "hello".into(),
            })
            .await;
        assert!(matches!(asked, Response::Asked { .. }));

        let listed = daemon.handle(Request::List).await;
        match listed {
            Response::List { harnesses } => {
                assert_eq!(harnesses.len(), 1);
                assert!(harnesses[0].alive);
                assert!(harnesses[0].log_path.ends_with("brain.jsonl"));
            }
            other => panic!("unexpected response: {other:?}"),
        }

        let messages = daemon
            .handle(Request::Messages {
                name: "brain".into(),
            })
            .await;
        match messages {
            Response::Messages { messages } => assert_eq!(messages[0].body, "hello"),
            other => panic!("unexpected response: {other:?}"),
        }

        let events = wait_for_daemon_events(&daemon, "brain", |events| {
            events.iter().any(|e| e.line.contains("reply:hello"))
        })
        .await;
        assert!(events.iter().any(|e| e.line == "hello"));

        let stopped = daemon
            .handle(Request::Stop {
                name: "brain".into(),
            })
            .await;
        assert!(matches!(stopped, Response::Stopped { .. }));
        assert!(paths.state_file.exists());
    }

    #[tokio::test]
    async fn daemon_ensure_upserts_and_unknown_spawn_errors() {
        let dir = tempfile::tempdir().unwrap();
        let paths = Paths::from_root(dir.path().join("home"));
        let daemon = Daemon::new(paths).await.unwrap();

        let unknown = daemon
            .handle(Request::Spawn {
                name: "missing".into(),
            })
            .await;
        assert!(matches!(unknown, Response::Error { .. }));

        let ensured = daemon
            .handle(Request::Ensure {
                config: mock_config(&dir, "ensured"),
            })
            .await;
        assert!(matches!(ensured, Response::Spawned { .. }));
        let ensured_again = daemon
            .handle(Request::Spawn {
                name: "ensured".into(),
            })
            .await;
        assert!(matches!(ensured_again, Response::Spawned { .. }));
        let _ = daemon
            .handle(Request::Stop {
                name: "ensured".into(),
            })
            .await;

        let ask_missing = daemon
            .handle(Request::Ask {
                from: "a".into(),
                to: "missing".into(),
                body: "b".into(),
            })
            .await;
        assert!(matches!(ask_missing, Response::Error { .. }));

        let stop_missing = daemon
            .handle(Request::Stop {
                name: "missing".into(),
            })
            .await;
        assert!(matches!(stop_missing, Response::Stopped { .. }));
    }

    #[tokio::test]
    async fn handle_stream_serves_json_lines_over_unix_pair() {
        let dir = tempfile::tempdir().unwrap();
        let daemon = Daemon::new(Paths::from_root(dir.path().join("home")))
            .await
            .unwrap();
        let (mut client, server) = UnixStream::pair().unwrap();
        let task = tokio::spawn(async move { daemon.handle_stream(server).await.unwrap() });
        client.write_all(br#"{"type":"ping"}"#).await.unwrap();
        client.write_all(b"\n").await.unwrap();
        let mut lines = BufReader::new(client).lines();
        let line = lines.next_line().await.unwrap().unwrap();
        assert!(line.contains("pong"));
        drop(lines);
        task.abort();
    }

    async fn wait_for_daemon_events<F>(
        daemon: &Daemon,
        name: &str,
        predicate: F,
    ) -> Vec<crate::protocol::HarnessEvent>
    where
        F: Fn(&[crate::protocol::HarnessEvent]) -> bool,
    {
        tokio::time::timeout(Duration::from_secs(5), async {
            loop {
                let response = daemon
                    .handle(Request::Events {
                        name: name.into(),
                        max: 50,
                    })
                    .await;
                let Response::Events { events } = response else {
                    panic!("unexpected response: {response:?}");
                };
                if predicate(&events) {
                    return events;
                }
                tokio::time::sleep(Duration::from_millis(10)).await;
            }
        })
        .await
        .unwrap()
    }

    #[tokio::test]
    async fn inject_mcp_config_preserves_existing_servers() {
        let dir = tempfile::tempdir().unwrap();
        let path = dir.path().join(".mcp.json");
        tokio::fs::write(
            &path,
            r#"{"mcpServers":{"other":{"command":"old","args":[],"env":{}}}}"#,
        )
        .await
        .unwrap();

        inject_mcp_config(dir.path().to_str().unwrap(), "multi-brain", "/tmp/home")
            .await
            .unwrap();
        let raw = tokio::fs::read_to_string(&path).await.unwrap();
        let parsed: Value = serde_json::from_str(&raw).unwrap();
        assert_eq!(parsed["mcpServers"]["other"]["command"], "old");
        assert!(parsed["mcpServers"]["multi-brain"]["command"].is_string());
    }

    #[tokio::test]
    async fn inject_mcp_config_sets_explicit_home() {
        let dir = tempfile::tempdir().unwrap();
        let path = dir.path().join(".mcp.json");
        inject_mcp_config(
            dir.path().to_str().unwrap(),
            "multi-brain",
            "/tmp/brain-home",
        )
        .await
        .unwrap();
        let raw = tokio::fs::read_to_string(&path).await.unwrap();
        let parsed: Value = serde_json::from_str(&raw).unwrap();
        assert_eq!(
            parsed["mcpServers"]["multi-brain"]["env"]["MULTI_BRAIN_HOME"],
            "/tmp/brain-home"
        );
    }
}
