//! Minimal Language Server Protocol client for diagnostics after edits.
//!
//! Language servers are started on demand per (user, server, project root),
//! run in the user's sandbox like shell commands, and are stopped after a
//! while without use. After the agent changes a file, the file is synced to
//! the server and the errors it reports are appended to the tool result.

use crate::app::App;
use crate::config::LspServer;
use crate::sandbox::SandboxSpec;
use anyhow::{Context, Result, anyhow, bail};
use serde_json::{Value, json};
use std::collections::HashMap;
use std::path::{Path, PathBuf};
use std::process::Stdio;
use std::sync::atomic::{AtomicI64, Ordering};
use std::sync::{Arc, Mutex};
use std::time::{Duration, Instant};
use tokio::io::{AsyncBufReadExt, AsyncReadExt, AsyncWriteExt, BufReader};
use tokio::process::{Child, ChildStdin};
use tokio::sync::{Notify, oneshot};

const IDLE_TIMEOUT: Duration = Duration::from_secs(15 * 60);
const MAX_SHOWN: usize = 20;

/// Built-in servers (used when the config defines none).
pub fn default_servers() -> Vec<(String, LspServer)> {
    let s = |cmd: &str, args: &[&str], ext: &[&str], roots: &[&str]| LspServer {
        command: cmd.into(),
        args: args.iter().map(|a| a.to_string()).collect(),
        extensions: ext.iter().map(|a| a.to_string()).collect(),
        root_markers: roots.iter().map(|a| a.to_string()).collect(),
        init_options: Value::Null,
        enabled: true,
    };
    // Global TypeScript as fallback when the project has none in node_modules.
    let home = std::env::var("HOME").unwrap_or_default();
    let ts_lib = ["/usr/local/lib/node_modules/typescript/lib".to_string(), format!("{home}/.npm-global/lib/node_modules/typescript/lib")]
        .into_iter()
        .find(|p| Path::new(p).is_dir());
    let mut ts = s("typescript-language-server", &["--stdio"], &["ts", "tsx", "js", "jsx", "mjs", "cjs", "mts", "cts"], &["tsconfig.json", "jsconfig.json", "package.json"]);
    if let Some(lib) = ts_lib {
        ts.init_options = json!({"tsserver": {"fallbackPath": lib}});
    }
    vec![
        ("rust".into(), s("rust-analyzer", &[], &["rs"], &["Cargo.toml"])),
        ("typescript".into(), ts),
        ("python".into(), s("pyright-langserver", &["--stdio"], &["py", "pyi"], &["pyproject.toml", "setup.py", "setup.cfg", "requirements.txt", "pyrightconfig.json"])),
        ("go".into(), s("gopls", &[], &["go"], &["go.mod"])),
        ("c".into(), s("clangd", &[], &["c", "h", "cc", "cpp", "cxx", "hpp", "hh"], &["compile_commands.json", "CMakeLists.txt", "Makefile"])),
    ]
}

#[derive(Debug, Clone)]
pub struct Diagnostic {
    pub line: u64,
    pub col: u64,
    pub severity: u64,
    pub message: String,
    pub source: String,
}

struct Client {
    stdin: tokio::sync::Mutex<ChildStdin>,
    child: Mutex<Option<Child>>,
    next_id: AtomicI64,
    pending: Mutex<HashMap<i64, oneshot::Sender<Value>>>,
    /// uri -> (publish counter, diagnostics)
    diags: Mutex<HashMap<String, (u64, Vec<Diagnostic>)>>,
    published: Notify,
    /// uri -> version of opened documents
    open: Mutex<HashMap<String, i64>>,
    last_used: Mutex<Instant>,
    dead: std::sync::atomic::AtomicBool,
}

impl Drop for Client {
    fn drop(&mut self) {
        if let Some(mut c) = self.child.lock().unwrap().take() {
            let _ = c.start_kill();
        }
    }
}

#[derive(Default)]
pub struct Lsp {
    clients: Mutex<HashMap<String, Arc<Client>>>,
    /// Servers that failed to start (key -> time), not retried for a while.
    failed: Mutex<HashMap<String, Instant>>,
}

fn file_uri(p: &Path) -> String {
    let mut s = String::from("file://");
    for b in p.display().to_string().bytes() {
        match b {
            b'A'..=b'Z' | b'a'..=b'z' | b'0'..=b'9' | b'/' | b'-' | b'_' | b'.' | b'~' => s.push(b as char),
            _ => s.push_str(&format!("%{b:02X}")),
        }
    }
    s
}

fn language_id(ext: &str) -> &'static str {
    match ext {
        "rs" => "rust",
        "ts" | "mts" | "cts" => "typescript",
        "tsx" => "typescriptreact",
        "js" | "mjs" | "cjs" => "javascript",
        "jsx" => "javascriptreact",
        "py" | "pyi" => "python",
        "go" => "go",
        "c" | "h" => "c",
        _ => "cpp",
    }
}

/// Find a command on the PATH the process will see (sandbox home or server home).
fn find_command(cmd: &str, sandbox: Option<&SandboxSpec>) -> Option<PathBuf> {
    if cmd.contains('/') {
        return Path::new(cmd).exists().then(|| PathBuf::from(cmd));
    }
    let mut dirs: Vec<PathBuf> = Vec::new();
    match sandbox {
        Some(s) => {
            for d in [".local/bin", ".cargo/bin", ".npm-global/bin", "go/bin"] {
                dirs.push(s.home.join(d));
            }
        }
        None => {
            if let Ok(p) = std::env::var("PATH") {
                dirs.extend(std::env::split_paths(&p));
            }
            if let Ok(h) = std::env::var("HOME") {
                dirs.push(Path::new(&h).join("go/bin"));
            }
        }
    }
    dirs.extend(["/usr/local/bin", "/usr/bin", "/bin"].map(PathBuf::from));
    dirs.into_iter().map(|d| d.join(cmd)).find(|p| p.is_file())
}

/// Nearest directory (from the file upwards, not above `limit`) containing a marker.
fn project_root(file: &Path, markers: &[String], limit: &Path) -> PathBuf {
    let mut dir = file.parent();
    while let Some(d) = dir {
        if markers.iter().any(|m| d.join(m).exists()) {
            return d.to_path_buf();
        }
        if d == limit || !d.starts_with(limit) {
            break;
        }
        dir = d.parent();
    }
    if file.starts_with(limit) { limit.to_path_buf() } else { file.parent().unwrap_or(limit).to_path_buf() }
}

async fn write_msg(stdin: &tokio::sync::Mutex<ChildStdin>, msg: &Value) -> Result<()> {
    let body = serde_json::to_vec(msg)?;
    let mut w = stdin.lock().await;
    w.write_all(format!("Content-Length: {}\r\n\r\n", body.len()).as_bytes()).await?;
    w.write_all(&body).await?;
    w.flush().await?;
    Ok(())
}

impl Client {
    async fn notify(&self, method: &str, params: Value) -> Result<()> {
        write_msg(&self.stdin, &json!({"jsonrpc": "2.0", "method": method, "params": params})).await
    }

    async fn request(&self, method: &str, params: Value, timeout: Duration) -> Result<Value> {
        let id = self.next_id.fetch_add(1, Ordering::SeqCst);
        let (tx, rx) = oneshot::channel();
        self.pending.lock().unwrap().insert(id, tx);
        write_msg(&self.stdin, &json!({"jsonrpc": "2.0", "id": id, "method": method, "params": params})).await?;
        let v = tokio::time::timeout(timeout, rx).await.map_err(|_| anyhow!("{method}: no answer"))??;
        if let Some(e) = v.get("error") {
            bail!("{method}: {e}");
        }
        Ok(v["result"].clone())
    }

    fn handle(self: &Arc<Self>, msg: Value) {
        let method = msg["method"].as_str().unwrap_or("");
        if msg.get("id").is_some() && !method.is_empty() {
            // Request from the server: answer the few that need it.
            let result = match method {
                "workspace/configuration" => Value::Array(vec![Value::Null; msg["params"]["items"].as_array().map_or(0, |a| a.len())]),
                _ => Value::Null,
            };
            let me = self.clone();
            let id = msg["id"].clone();
            tokio::spawn(async move {
                let _ = write_msg(&me.stdin, &json!({"jsonrpc": "2.0", "id": id, "result": result})).await;
            });
            return;
        }
        if let Some(id) = msg["id"].as_i64() {
            if let Some(tx) = self.pending.lock().unwrap().remove(&id) {
                let _ = tx.send(msg);
            }
            return;
        }
        if method == "textDocument/publishDiagnostics" {
            let uri = msg["params"]["uri"].as_str().unwrap_or("").to_string();
            let list = msg["params"]["diagnostics"]
                .as_array()
                .map(|a| {
                    a.iter()
                        .map(|d| Diagnostic {
                            line: d["range"]["start"]["line"].as_u64().unwrap_or(0) + 1,
                            col: d["range"]["start"]["character"].as_u64().unwrap_or(0) + 1,
                            severity: d["severity"].as_u64().unwrap_or(1),
                            message: d["message"].as_str().unwrap_or("").to_string(),
                            source: d["source"].as_str().unwrap_or("").to_string(),
                        })
                        .collect()
                })
                .unwrap_or_default();
            let mut diags = self.diags.lock().unwrap();
            let n = diags.get(&uri).map_or(0, |e| e.0) + 1;
            diags.insert(uri, (n, list));
            drop(diags);
            self.published.notify_waiters();
        }
    }

    fn publish_count(&self, uri: &str) -> u64 {
        self.diags.lock().unwrap().get(uri).map_or(0, |e| e.0)
    }
}

async fn start(app: &App, server: &LspServer, root: &Path, sandbox: Option<&SandboxSpec>) -> Result<Arc<Client>> {
    let program = find_command(&server.command, sandbox).ok_or_else(|| anyhow!("{} is not installed", server.command))?;
    let program = program.display().to_string();
    let args: Vec<&str> = server.args.iter().map(String::as_str).collect();
    let mut cmd = match sandbox {
        Some(spec) => {
            if !crate::sandbox::available() {
                bail!("sandbox unavailable");
            }
            // Inside the sandbox the home directory is mounted elsewhere.
            let inside = program.replace(&spec.home.display().to_string(), &spec.home_inside());
            crate::sandbox::command(spec, root, &inside, &args)
        }
        None => {
            let mut c = tokio::process::Command::new(&program);
            c.args(&args).current_dir(root);
            c
        }
    };
    let _ = app;
    cmd.stdin(Stdio::piped()).stdout(Stdio::piped()).stderr(Stdio::null()).kill_on_drop(true);
    let mut child = cmd.spawn().with_context(|| format!("starting {}", server.command))?;
    let stdin = child.stdin.take().unwrap();
    let stdout = child.stdout.take().unwrap();
    let client = Arc::new(Client {
        stdin: tokio::sync::Mutex::new(stdin),
        child: Mutex::new(Some(child)),
        next_id: AtomicI64::new(1),
        pending: Mutex::default(),
        diags: Mutex::default(),
        published: Notify::new(),
        open: Mutex::default(),
        last_used: Mutex::new(Instant::now()),
        dead: false.into(),
    });

    // Reader: Content-Length framed JSON messages.
    let weak = Arc::downgrade(&client);
    tokio::spawn(async move {
        let mut reader = BufReader::new(stdout);
        loop {
            let mut len = 0usize;
            loop {
                let mut line = String::new();
                match reader.read_line(&mut line).await {
                    Ok(0) | Err(_) => {
                        if let Some(c) = weak.upgrade() {
                            c.dead.store(true, Ordering::SeqCst);
                            c.published.notify_waiters();
                        }
                        return;
                    }
                    Ok(_) => {}
                }
                let line = line.trim();
                if line.is_empty() {
                    break;
                }
                if let Some(v) = line.strip_prefix("Content-Length:") {
                    len = v.trim().parse().unwrap_or(0);
                }
            }
            if len == 0 || len > 64 * 1024 * 1024 {
                continue;
            }
            let mut buf = vec![0u8; len];
            if reader.read_exact(&mut buf).await.is_err() {
                return;
            }
            let Some(c) = weak.upgrade() else { return };
            if let Ok(msg) = serde_json::from_slice::<Value>(&buf) {
                c.handle(msg);
            }
        }
    });

    let root_uri = file_uri(root);
    client
        .request(
            "initialize",
            json!({
                "processId": null,
                "rootUri": root_uri,
                "rootPath": root.display().to_string(),
                "workspaceFolders": [{"uri": root_uri, "name": root.file_name().and_then(|n| n.to_str()).unwrap_or("project")}],
                "capabilities": {
                    "textDocument": {
                        "synchronization": {"didSave": true, "dynamicRegistration": false},
                        "publishDiagnostics": {"relatedInformation": false, "versionSupport": true}
                    },
                    "workspace": {"configuration": true, "workspaceFolders": true}
                },
                "initializationOptions": if server.init_options.is_null() { json!({}) } else { server.init_options.clone() }
            }),
            Duration::from_secs(60),
        )
        .await?;
    client.notify("initialized", json!({})).await?;
    Ok(client)
}

impl Lsp {
    /// Stop servers that were not used for a while.
    pub fn reap(&self) {
        self.clients.lock().unwrap().retain(|_, c| c.last_used.lock().unwrap().elapsed() < IDLE_TIMEOUT && !c.dead.load(Ordering::SeqCst));
    }
}

/// The configured servers (config, else built-in defaults).
fn servers(app: &App) -> Vec<(String, LspServer)> {
    let cfg = app.config();
    if !cfg.lsp.enabled {
        return vec![];
    }
    let mut list = default_servers();
    for (name, s) in &cfg.lsp.servers {
        list.retain(|(n, _)| n != name);
        list.push((name.clone(), s.clone()));
    }
    list.retain(|(_, s)| s.enabled);
    list
}

/// Sync a file to its language server and return the diagnostics it reports.
/// `None` when no server handles the file (or it is not installed).
pub async fn diagnostics(app: &Arc<App>, file: &Path, limit: &Path, sandbox: Option<&SandboxSpec>, wait: Duration) -> Option<Result<Vec<Diagnostic>>> {
    let ext = file.extension()?.to_str()?.to_ascii_lowercase();
    let (name, server) = servers(app).into_iter().find(|(_, s)| s.extensions.iter().any(|e| e.eq_ignore_ascii_case(&ext)))?;
    let root = project_root(file, &server.root_markers, limit);
    let user = sandbox.map(|s| s.username.as_str()).unwrap_or("-");
    let key = format!("{user}|{name}|{}", root.display());

    if let Some(t) = app.lsp.failed.lock().unwrap().get(&key) {
        if t.elapsed() < Duration::from_secs(300) {
            return None;
        }
    }
    let existing = app.lsp.clients.lock().unwrap().get(&key).cloned().filter(|c| !c.dead.load(Ordering::SeqCst));
    let (client, fresh) = match existing {
        Some(c) => (c, false),
        None => match start(app, &server, &root, sandbox).await {
            Ok(c) => {
                app.lsp.clients.lock().unwrap().insert(key.clone(), c.clone());
                (c, true)
            }
            Err(e) => {
                tracing::debug!("lsp {name}: {e:#}");
                app.lsp.failed.lock().unwrap().insert(key, Instant::now());
                return None;
            }
        },
    };
    *client.last_used.lock().unwrap() = Instant::now();
    Some(sync_and_wait(&client, file, &ext, if fresh { wait.max(Duration::from_secs(20)) } else { wait }).await)
}

async fn sync_and_wait(client: &Arc<Client>, file: &Path, ext: &str, wait: Duration) -> Result<Vec<Diagnostic>> {
    let text = tokio::fs::read_to_string(file).await.unwrap_or_default();
    let uri = file_uri(file);
    let before = client.publish_count(&uri);
    let version = {
        let mut open = client.open.lock().unwrap();
        let v = open.get(&uri).map_or(1, |v| v + 1);
        open.insert(uri.clone(), v);
        v
    };
    if version == 1 {
        client
            .notify("textDocument/didOpen", json!({"textDocument": {"uri": uri, "languageId": language_id(ext), "version": 1, "text": text}}))
            .await?;
    } else {
        client
            .notify("textDocument/didChange", json!({"textDocument": {"uri": uri, "version": version}, "contentChanges": [{"text": text}]}))
            .await?;
    }
    client.notify("textDocument/didSave", json!({"textDocument": {"uri": uri}, "text": text})).await?;

    // Wait for a fresh publish; servers often send an early empty one and the
    // real result shortly after, so keep listening briefly after the first.
    let deadline = Instant::now() + wait;
    let mut settle: Option<Instant> = None;
    loop {
        if client.dead.load(Ordering::SeqCst) {
            bail!("language server exited");
        }
        let now = Instant::now();
        let count = client.publish_count(&uri);
        if count > before && settle.is_none() {
            settle = Some(now + Duration::from_millis(700));
        }
        let until = settle.map_or(deadline, |s| s.min(deadline));
        if now >= until {
            break;
        }
        let _ = tokio::time::timeout(until - now, client.published.notified()).await;
        if let Some(s) = settle {
            if client.publish_count(&uri) > count {
                settle = Some(Instant::now().max(s));
            }
        }
    }
    if client.publish_count(&uri) == before {
        bail!("no diagnostics within {}s (server still indexing?)", wait.as_secs());
    }
    Ok(client.diags.lock().unwrap().get(&uri).map(|e| e.1.clone()).unwrap_or_default())
}

/// Format errors (and a few warnings) for a tool result.
pub fn format(file: &str, diags: &[Diagnostic], warnings: bool) -> String {
    let mut list: Vec<&Diagnostic> = diags.iter().filter(|d| d.severity == 1 || (warnings && d.severity == 2)).collect();
    list.sort_by_key(|d| (d.severity, d.line, d.col));
    if list.is_empty() {
        return String::new();
    }
    let errors = diags.iter().filter(|d| d.severity == 1).count();
    let mut out = format!("\n\n<diagnostics file=\"{file}\" errors=\"{errors}\">\n");
    for d in list.iter().take(MAX_SHOWN) {
        let sev = if d.severity == 1 { "error" } else { "warning" };
        let src = if d.source.is_empty() { String::new() } else { format!(" [{}]", d.source) };
        out.push_str(&format!("{}:{} {sev}{src}: {}\n", d.line, d.col, d.message.lines().next().unwrap_or("")));
    }
    if list.len() > MAX_SHOWN {
        out.push_str(&format!("… {} more\n", list.len() - MAX_SHOWN));
    }
    out.push_str("</diagnostics>");
    out
}

/// Background task: stop idle servers.
pub fn spawn_reaper(app: Arc<App>) {
    tokio::spawn(async move {
        loop {
            tokio::time::sleep(Duration::from_secs(60)).await;
            app.lsp.reap();
        }
    });
}

// ---------------------------------------------------------------------------

/// Tool: diagnostics of a file on demand.
pub struct Diagnostics;

#[async_trait::async_trait]
impl crate::tools::Tool for Diagnostics {
    fn name(&self) -> &str {
        "diagnostics"
    }
    fn description(&self) -> String {
        "Get compiler/linter diagnostics (errors and warnings) for a source file from its language server \
         (rust-analyzer, typescript-language-server, pyright, gopls, clangd – if installed). Edits already \
         report new errors automatically; use this to check a file you did not edit."
            .into()
    }
    fn schema(&self) -> Value {
        json!({"type": "object", "properties": {"path": {"type": "string"}}, "required": ["path"]})
    }
    fn read_only(&self) -> bool {
        true
    }
    fn title(&self, input: &Value) -> String {
        format!("Diagnostics {}", input["path"].as_str().unwrap_or(""))
    }
    async fn run(&self, input: Value, ctx: &crate::tools::ToolCtx) -> Result<crate::tools::ToolOutput> {
        let p = crate::tools::arg_str(&input, "path")?;
        let path = ctx.resolve_read(p)?;
        match diagnostics(&ctx.app, &path, &ctx.cwd, ctx.sandbox.as_ref(), Duration::from_secs(10)).await {
            None => Ok(crate::tools::ToolOutput::text(
                "No language server available for this file type. Install one (e.g. `npm i -g typescript-language-server typescript`, \
                 `npm i -g pyright`, `rustup component add rust-analyzer`, `go install golang.org/x/tools/gopls@latest`).",
            )),
            Some(Err(e)) => Ok(crate::tools::ToolOutput::error(format!("Language server: {e:#}"))),
            Some(Ok(d)) => {
                let s = format(p, &d, true);
                Ok(crate::tools::ToolOutput::text(if s.is_empty() { "No errors or warnings.".to_string() } else { s.trim().to_string() }))
            }
        }
    }
}