zesdex/src/app/lsp/client.rs

use std::io::{BufRead, BufReader, Read, Write};
use std::process::{Command, Stdio};
use std::time::{Duration, Instant};

use serde_json::{json, Value};

const LSP_INIT_TIMEOUT_MS: u64 = 60_000;
const LSP_CALL_TIMEOUT_MS: u64 = 30_000;
const LSP_DIAGNOSTICS_TIMEOUT_MS: u64 = 10_000;

pub struct LspClient {
    stdin: std::process::ChildStdin,
    stdout: BufReader<std::process::ChildStdout>,
    next_id: u64,
    server_capabilities: Value,
}

fn file_path_to_uri(path: &str) -> String {
    let abs_path = std::path::Path::new(path);
    let abs_path = if abs_path.is_relative() {
        match std::env::current_dir() {
            Ok(cwd) => cwd.join(path),
            Err(_) => abs_path.to_path_buf(),
        }
    } else {
        abs_path.to_path_buf()
    };
    let canonical = abs_path.canonicalize().unwrap_or(abs_path);
    let path_str = canonical.to_string_lossy();
    if cfg!(windows) {
        let path_str = path_str.replace('\\', "/");
        if path_str.starts_with('/') {
            format!("file://{}", path_str)
        } else {
            format!("file:///{}", path_str)
        }
    } else {
        format!("file://{}", path_str)
    }
}

impl LspClient {
    pub fn spawn(command: &str, args: &[String]) -> anyhow::Result<Self> {
        let mut cmd = Command::new(command);
        cmd.args(args);
        cmd.stdin(Stdio::piped());
        cmd.stdout(Stdio::piped());
        cmd.stderr(Stdio::piped());

        let mut child = cmd.spawn()
            .map_err(|e| anyhow::anyhow!("failed to spawn LSP server '{}': {}", command, e))?;

        let stdin = child.stdin.take()
            .ok_or_else(|| anyhow::anyhow!("failed to capture stdin for LSP server"))?;
        let stdout = BufReader::new(child.stdout.take()
            .ok_or_else(|| anyhow::anyhow!("failed to capture stdout for LSP server"))?);

        let mut client = LspClient {
            stdin,
            stdout,
            next_id: 0,
            server_capabilities: Value::Null,
        };

        let init_params = json!({
            "processId": std::process::id(),
            "clientInfo": {
                "name": "zesdex",
                "version": "0.1.0"
            },
            "capabilities": {
                "textDocument": {
                    "synchronization": {
                        "dynamicRegistration": true,
                        "willSave": false,
                        "willSaveWaitUntil": false,
                        "didSave": false
                    },
                    "hover": {
                        "dynamicRegistration": true,
                        "contentFormat": ["plaintext", "markdown"]
                    },
                    "completion": {
                        "dynamicRegistration": true,
                        "completionItem": {
                            "snippetSupport": false
                        }
                    },
                    "definition": {
                        "dynamicRegistration": true
                    },
                    "references": {
                        "dynamicRegistration": true
                    },
                    "documentSymbol": {
                        "dynamicRegistration": true,
                        "hierarchicalDocumentSymbolSupport": true
                    }
                },
                "workspace": {
                    "workspaceFolders": true
                },
                "general": {
                    "positionEncodings": ["utf-16"]
                }
            }
        });

        let result = client.call_with_timeout("initialize", init_params, Duration::from_millis(LSP_INIT_TIMEOUT_MS))?;
        client.server_capabilities = result.get("capabilities").cloned().unwrap_or_default();

        client.notify("initialized", json!({}))?;

        Ok(client)
    }

    pub fn server_capabilities(&self) -> &Value {
        &self.server_capabilities
    }

    pub fn call(&mut self, method: &str, params: Value) -> anyhow::Result<Value> {
        self.call_with_timeout(method, params, Duration::from_millis(LSP_CALL_TIMEOUT_MS))
    }

    fn call_with_timeout(&mut self, method: &str, params: Value, timeout: Duration) -> anyhow::Result<Value> {
        self.next_id += 1;
        let id = self.next_id;
        let req = json!({
            "jsonrpc": "2.0",
            "id": id,
            "method": method,
            "params": params
        });
        self.send_frame(&req)?;
        self.read_response(id, timeout)
    }

    pub fn notify(&mut self, method: &str, params: Value) -> anyhow::Result<()> {
        let req = json!({
            "jsonrpc": "2.0",
            "method": method,
            "params": params
        });
        self.send_frame(&req)
    }

    fn send_frame(&mut self, msg: &Value) -> anyhow::Result<()> {
        let body = serde_json::to_string(msg)
            .map_err(|e| anyhow::anyhow!("failed to serialize LSP message: {}", e))?;
        let header = format!("Content-Length: {}\r\n\r\n", body.len());
        self.stdin.write_all(header.as_bytes())
            .map_err(|e| anyhow::anyhow!("failed to write LSP frame header: {}", e))?;
        self.stdin.write_all(body.as_bytes())
            .map_err(|e| anyhow::anyhow!("failed to write LSP frame body: {}", e))?;
        self.stdin.flush()
            .map_err(|e| anyhow::anyhow!("failed to flush LSP stdin: {}", e))?;
        Ok(())
    }

    fn read_response(&mut self, expected_id: u64, timeout: Duration) -> anyhow::Result<Value> {
        let deadline = Instant::now() + timeout;
        loop {
            if Instant::now() > deadline {
                anyhow::bail!("LSP call timed out after {}ms", timeout.as_millis());
            }
            let frame = self.read_frame()?;
            if frame.get("id") == Some(&json!(expected_id)) {
                if let Some(err) = frame.get("error") {
                    let code = err.get("code").and_then(|c| c.as_i64()).unwrap_or(0);
                    let msg = err.get("message").and_then(|m| m.as_str()).unwrap_or("unknown error");
                    anyhow::bail!("LSP error {}: {}", code, msg);
                }
                return Ok(frame.get("result").cloned().unwrap_or(Value::Null));
            }
        }
    }

    pub fn read_notification(&mut self, method: &str, timeout: Duration) -> anyhow::Result<Value> {
        let deadline = Instant::now() + timeout;
        loop {
            if Instant::now() > deadline {
                anyhow::bail!("timed out waiting for LSP notification '{}'", method);
            }
            let frame = self.read_frame()?;
            if frame.get("method") == Some(&json!(method)) {
                return Ok(frame.get("params").cloned().unwrap_or(Value::Null));
            }
        }
    }

    fn read_frame(&mut self) -> anyhow::Result<Value> {
        let mut content_length: Option<usize> = None;
        loop {
            let mut line = String::new();
            match self.stdout.read_line(&mut line) {
                Ok(0) => anyhow::bail!("LSP server closed the connection"),
                Ok(_) => {}
                Err(e) => anyhow::bail!("LSP read error: {}", e),
            }
            let trimmed = line.trim();
            if trimmed.is_empty() {
                break;
            }
            if let Some(len_str) = trimmed.strip_prefix("Content-Length: ") {
                content_length = Some(
                    len_str.trim().parse::<usize>()
                        .map_err(|e| anyhow::anyhow!("invalid Content-Length '{}': {}", len_str.trim(), e))?,
                );
            }
        }

        let length = content_length
            .ok_or_else(|| anyhow::anyhow!("missing Content-Length header in LSP response"))?;

        let mut body = vec![0u8; length];
        self.stdout.read_exact(&mut body)
            .map_err(|e| anyhow::anyhow!("failed to read LSP body ({} bytes): {}", length, e))?;

        let json_str = String::from_utf8(body)
            .map_err(|e| anyhow::anyhow!("invalid UTF-8 in LSP response: {}", e))?;

        serde_json::from_str(&json_str)
            .map_err(|e| anyhow::anyhow!("invalid JSON in LSP response: {}", e))
    }

    pub fn did_open(&mut self, uri: &str, language_id: &str, version: i32, text: &str) -> anyhow::Result<()> {
        self.notify("textDocument/didOpen", json!({
            "textDocument": {
                "uri": uri,
                "languageId": language_id,
                "version": version,
                "text": text
            }
        }))
    }

    #[allow(dead_code)]
    pub fn did_change(&mut self, uri: &str, version: i32, text: &str) -> anyhow::Result<()> {
        self.notify("textDocument/didChange", json!({
            "textDocument": {
                "uri": uri,
                "version": version
            },
            "contentChanges": [{
                "text": text
            }]
        }))
    }

    pub fn did_close(&mut self, uri: &str) -> anyhow::Result<()> {
        self.notify("textDocument/didClose", json!({
            "textDocument": {
                "uri": uri
            }
        }))
    }

    pub fn hover(&mut self, uri: &str, line: u32, character: u32) -> anyhow::Result<Value> {
        self.call("textDocument/hover", json!({
            "textDocument": { "uri": uri },
            "position": { "line": line, "character": character }
        }))
    }

    pub fn completion(&mut self, uri: &str, line: u32, character: u32) -> anyhow::Result<Value> {
        self.call("textDocument/completion", json!({
            "textDocument": { "uri": uri },
            "position": { "line": line, "character": character }
        }))
    }

    pub fn goto_definition(&mut self, uri: &str, line: u32, character: u32) -> anyhow::Result<Value> {
        self.call("textDocument/definition", json!({
            "textDocument": { "uri": uri },
            "position": { "line": line, "character": character }
        }))
    }

    pub fn references(&mut self, uri: &str, line: u32, character: u32) -> anyhow::Result<Value> {
        self.call("textDocument/references", json!({
            "textDocument": { "uri": uri },
            "position": { "line": line, "character": character },
            "context": {
                "includeDeclaration": true
            }
        }))
    }

    #[allow(dead_code)]
    pub fn document_symbols(&mut self, uri: &str) -> anyhow::Result<Value> {
        self.call("textDocument/documentSymbol", json!({
            "textDocument": { "uri": uri }
        }))
    }

    pub fn collect_diagnostics(
        &mut self,
        uri: &str,
        language_id: &str,
        text: &str,
    ) -> anyhow::Result<Value> {
        self.did_open(uri, language_id, 1, text)?;
        let result = self.read_notification(
            "textDocument/publishDiagnostics",
            Duration::from_millis(LSP_DIAGNOSTICS_TIMEOUT_MS),
        );
        self.did_close(uri)?;
        match result {
            Ok(params) => Ok(params.get("diagnostics").cloned().unwrap_or_else(|| json!([]))),
            Err(e) => Err(e),
        }
    }

    pub fn shutdown(&mut self) -> anyhow::Result<()> {
        let _ = self.call_with_timeout("shutdown", json!({}), Duration::from_secs(5));
        let _ = self.notify("exit", json!({}));
        Ok(())
    }
}

impl Drop for LspClient {
    fn drop(&mut self) {
        let _ = self.notify("exit", json!({}));
    }
}

pub fn path_to_lsp_uri(path: &str) -> String {
    file_path_to_uri(path)
}