codive_lsp/

client.rs

//! LSP client implementation
//!
//! This module provides the LSP client that communicates with language servers
//! using JSON-RPC 2.0 over stdio.

use crate::language::language_id_for_path;
use crate::server::LspServerHandle;
use crate::types::*;
use anyhow::{anyhow, Context, Result};
use lsp_types::request::{
    CallHierarchyIncomingCalls, CallHierarchyOutgoingCalls, CallHierarchyPrepare,
    DocumentSymbolRequest, GotoDefinition, GotoImplementation, HoverRequest, References,
    WorkspaceSymbolRequest,
};
use lsp_types::{
    CallHierarchyIncomingCallsParams, CallHierarchyOutgoingCallsParams, ClientCapabilities,
    DidChangeTextDocumentParams, DidOpenTextDocumentParams, DocumentSymbolParams,
    GotoDefinitionParams, HoverParams, InitializeParams, InitializedParams, PartialResultParams,
    ReferenceContext, ReferenceParams, TextDocumentContentChangeEvent, TextDocumentItem,
    VersionedTextDocumentIdentifier, WorkDoneProgressParams, WorkspaceSymbolParams,
};
// GotoImplementation uses the same params as GotoDefinition
use lsp_types::GotoDefinitionParams as GotoImplementationParams;
use serde::{de::DeserializeOwned, Deserialize, Serialize};
use std::collections::HashMap;
use std::io::{BufRead, BufReader, Write};
use std::path::{Path, PathBuf};

/// Helper function to convert a file path to LSP Uri
fn path_to_uri(path: &Path) -> Result<Uri> {
    let url = Url::from_file_path(path).map_err(|_| anyhow!("Invalid file path: {:?}", path))?;
    url.as_str()
        .parse()
        .map_err(|e| anyhow!("Failed to parse URI: {}", e))
}

/// Helper function to convert LSP Uri to file path
fn uri_to_path(uri: &Uri) -> Option<PathBuf> {
    let url: url::Url = uri.as_str().parse().ok()?;
    url.to_file_path().ok()
}
use std::process::{Child, ChildStdin, ChildStdout};
use std::sync::atomic::{AtomicI64, Ordering};
use std::sync::{Arc, Mutex, RwLock};
use tokio::sync::{mpsc, oneshot};
use tracing::{debug, error, info, trace, warn};

/// Timeout for LSP requests in milliseconds
const REQUEST_TIMEOUT_MS: u64 = 10_000;
/// Timeout for initialization in milliseconds
const INIT_TIMEOUT_MS: u64 = 45_000;

/// JSON-RPC request
#[derive(Debug, Serialize)]
struct JsonRpcRequest<T: Serialize> {
    jsonrpc: &'static str,
    id: i64,
    method: &'static str,
    params: T,
}

/// JSON-RPC notification (no id)
#[derive(Debug, Serialize)]
struct JsonRpcNotification<T: Serialize> {
    jsonrpc: &'static str,
    method: &'static str,
    params: T,
}

/// JSON-RPC response
#[derive(Debug, Deserialize)]
struct JsonRpcResponse<T> {
    #[allow(dead_code)]
    jsonrpc: String,
    id: Option<i64>,
    result: Option<T>,
    error: Option<JsonRpcError>,
}

/// JSON-RPC error
#[derive(Debug, Deserialize)]
struct JsonRpcError {
    code: i64,
    message: String,
}

/// Generic JSON-RPC message for routing
#[derive(Debug, Deserialize)]
struct JsonRpcMessage {
    #[allow(dead_code)]
    jsonrpc: String,
    id: Option<i64>,
    method: Option<String>,
    #[serde(default)]
    params: serde_json::Value,
    result: Option<serde_json::Value>,
    error: Option<JsonRpcError>,
}

/// Pending request awaiting response
struct PendingRequest {
    sender: oneshot::Sender<Result<serde_json::Value>>,
}

/// Active LSP client connection
pub struct LspClient {
    /// Server ID
    server_id: String,
    /// Project root
    root: PathBuf,
    /// Request ID counter
    request_id: AtomicI64,
    /// Pending requests
    pending: Arc<RwLock<HashMap<i64, PendingRequest>>>,
    /// Writer to server stdin
    writer: Arc<Mutex<ChildStdin>>,
    /// Published diagnostics per file
    diagnostics: Arc<RwLock<HashMap<PathBuf, Vec<Diagnostic>>>>,
    /// File versions for text synchronization
    file_versions: Arc<RwLock<HashMap<PathBuf, i32>>>,
    /// Channel to signal shutdown
    shutdown_tx: Option<mpsc::Sender<()>>,
    /// Child process handle
    #[allow(dead_code)]
    process: Child,
}

impl LspClient {
    /// Create and initialize a new LSP client
    pub async fn new(
        server_id: impl Into<String>,
        mut handle: LspServerHandle,
        root: PathBuf,
    ) -> Result<Self> {
        let server_id = server_id.into();
        info!(server_id = %server_id, root = ?root, "Initializing LSP client");

        let stdin = handle
            .process
            .stdin
            .take()
            .context("Failed to get stdin")?;
        let stdout = handle
            .process
            .stdout
            .take()
            .context("Failed to get stdout")?;

        let pending: Arc<RwLock<HashMap<i64, PendingRequest>>> =
            Arc::new(RwLock::new(HashMap::new()));
        let diagnostics: Arc<RwLock<HashMap<PathBuf, Vec<Diagnostic>>>> =
            Arc::new(RwLock::new(HashMap::new()));
        let file_versions: Arc<RwLock<HashMap<PathBuf, i32>>> =
            Arc::new(RwLock::new(HashMap::new()));

        let (shutdown_tx, shutdown_rx) = mpsc::channel::<()>(1);

        // Spawn reader task
        let pending_clone = pending.clone();
        let diagnostics_clone = diagnostics.clone();
        let server_id_clone = server_id.clone();
        std::thread::spawn(move || {
            Self::reader_loop(stdout, pending_clone, diagnostics_clone, server_id_clone);
        });

        let writer = Arc::new(Mutex::new(stdin));

        let mut client = Self {
            server_id,
            root: root.clone(),
            request_id: AtomicI64::new(1),
            pending,
            writer,
            diagnostics,
            file_versions,
            shutdown_tx: Some(shutdown_tx),
            process: handle.process,
        };

        // Initialize the server
        client.initialize(&root, handle.initialization).await?;

        Ok(client)
    }

    /// Get the server ID
    pub fn server_id(&self) -> &str {
        &self.server_id
    }

    /// Get the project root
    pub fn root(&self) -> &Path {
        &self.root
    }

    /// Get current diagnostics
    pub fn diagnostics(&self) -> HashMap<PathBuf, Vec<Diagnostic>> {
        self.diagnostics.read().unwrap().clone()
    }

    /// Get diagnostics for a specific file
    pub fn diagnostics_for_file(&self, path: &Path) -> Vec<Diagnostic> {
        self.diagnostics
            .read()
            .unwrap()
            .get(path)
            .cloned()
            .unwrap_or_default()
    }

    // ========================================================================
    // LSP Protocol Methods
    // ========================================================================

    /// Initialize the LSP server
    async fn initialize(
        &mut self,
        root: &Path,
        initialization_options: Option<serde_json::Value>,
    ) -> Result<()> {
        let root_uri = path_to_uri(root)?;

        let params = InitializeParams {
            process_id: Some(std::process::id()),
            root_uri: Some(root_uri.clone()),
            root_path: None,
            initialization_options,
            capabilities: Self::client_capabilities(),
            trace: None,
            workspace_folders: Some(vec![lsp_types::WorkspaceFolder {
                uri: root_uri,
                name: root
                    .file_name()
                    .and_then(|n| n.to_str())
                    .unwrap_or("workspace")
                    .to_string(),
            }]),
            client_info: Some(lsp_types::ClientInfo {
                name: "codive-lsp".to_string(),
                version: Some(env!("CARGO_PKG_VERSION").to_string()),
            }),
            locale: None,
            work_done_progress_params: WorkDoneProgressParams::default(),
        };

        let _result: lsp_types::InitializeResult =
            self.request::<lsp_types::request::Initialize>(params).await?;

        // Send initialized notification
        self.notify::<lsp_types::notification::Initialized>(InitializedParams {})?;

        info!(server_id = %self.server_id, "LSP server initialized");
        Ok(())
    }

    /// Notify the server about an opened file
    pub async fn open_file(&self, path: &Path) -> Result<()> {
        let path = path.canonicalize().unwrap_or_else(|_| path.to_path_buf());

        // Check if already open (before reading file content)
        let existing_version = {
            let versions = self.file_versions.read().unwrap();
            versions.get(&path).copied()
        };

        let content = tokio::fs::read_to_string(&path).await?;

        if let Some(version) = existing_version {
            // File is already open, send didChange instead
            return self.change_file(&path, &content, version + 1);
        }

        let uri = path_to_uri(&path)?;
        let language_id = language_id_for_path(&path);

        debug!(path = ?path, language_id = %language_id, "Opening file in LSP");

        let params = DidOpenTextDocumentParams {
            text_document: TextDocumentItem {
                uri,
                language_id: language_id.to_string(),
                version: 0,
                text: content,
            },
        };

        self.notify::<lsp_types::notification::DidOpenTextDocument>(params)?;

        // Track the file version
        self.file_versions.write().unwrap().insert(path, 0);

        Ok(())
    }

    /// Notify the server about a file change
    fn change_file(&self, path: &Path, content: &str, version: i32) -> Result<()> {
        let uri = path_to_uri(path)?;

        let params = DidChangeTextDocumentParams {
            text_document: VersionedTextDocumentIdentifier { uri, version },
            content_changes: vec![TextDocumentContentChangeEvent {
                range: None,
                range_length: None,
                text: content.to_string(),
            }],
        };

        self.notify::<lsp_types::notification::DidChangeTextDocument>(params)?;

        // Update version
        self.file_versions.write().unwrap().insert(path.to_path_buf(), version);

        Ok(())
    }

    /// Get hover information
    pub async fn hover(&self, path: &Path, line: u32, character: u32) -> Result<Option<Hover>> {
        let uri = path_to_uri(path)?;

        let params = HoverParams {
            text_document_position_params: TextDocumentPositionParams {
                text_document: TextDocumentIdentifier { uri },
                position: Position { line, character },
            },
            work_done_progress_params: WorkDoneProgressParams::default(),
        };

        self.request::<HoverRequest>(params).await
    }

    /// Go to definition
    pub async fn definition(&self, path: &Path, line: u32, character: u32) -> Result<Vec<Location>> {
        let uri = path_to_uri(path)?;

        let params = GotoDefinitionParams {
            text_document_position_params: TextDocumentPositionParams {
                text_document: TextDocumentIdentifier { uri },
                position: Position { line, character },
            },
            work_done_progress_params: WorkDoneProgressParams::default(),
            partial_result_params: PartialResultParams::default(),
        };

        let result: Option<GotoDefinitionResponse> =
            self.request::<GotoDefinition>(params).await?;

        Ok(match result {
            Some(GotoDefinitionResponse::Scalar(loc)) => vec![loc],
            Some(GotoDefinitionResponse::Array(locs)) => locs,
            Some(GotoDefinitionResponse::Link(links)) => links
                .into_iter()
                .map(|l| Location {
                    uri: l.target_uri,
                    range: l.target_selection_range,
                })
                .collect(),
            None => vec![],
        })
    }

    /// Find references
    pub async fn references(
        &self,
        path: &Path,
        line: u32,
        character: u32,
        include_declaration: bool,
    ) -> Result<Vec<Location>> {
        let uri = path_to_uri(path)?;

        let params = ReferenceParams {
            text_document_position: TextDocumentPositionParams {
                text_document: TextDocumentIdentifier { uri },
                position: Position { line, character },
            },
            work_done_progress_params: WorkDoneProgressParams::default(),
            partial_result_params: PartialResultParams::default(),
            context: ReferenceContext {
                include_declaration,
            },
        };

        let result: Option<Vec<Location>> = self.request::<References>(params).await?;
        Ok(result.unwrap_or_default())
    }

    /// Go to implementation
    pub async fn implementation(
        &self,
        path: &Path,
        line: u32,
        character: u32,
    ) -> Result<Vec<Location>> {
        let uri = path_to_uri(path)?;

        let params = GotoImplementationParams {
            text_document_position_params: TextDocumentPositionParams {
                text_document: TextDocumentIdentifier { uri },
                position: Position { line, character },
            },
            work_done_progress_params: WorkDoneProgressParams::default(),
            partial_result_params: PartialResultParams::default(),
        };

        let result: Option<GotoDefinitionResponse> =
            self.request::<GotoImplementation>(params).await?;

        Ok(match result {
            Some(GotoDefinitionResponse::Scalar(loc)) => vec![loc],
            Some(GotoDefinitionResponse::Array(locs)) => locs,
            Some(GotoDefinitionResponse::Link(links)) => links
                .into_iter()
                .map(|l| Location {
                    uri: l.target_uri,
                    range: l.target_selection_range,
                })
                .collect(),
            None => vec![],
        })
    }

    /// Get document symbols
    pub async fn document_symbols(&self, path: &Path) -> Result<DocumentSymbolResponse> {
        let uri = path_to_uri(path)?;

        let params = DocumentSymbolParams {
            text_document: TextDocumentIdentifier { uri },
            work_done_progress_params: WorkDoneProgressParams::default(),
            partial_result_params: PartialResultParams::default(),
        };

        let result: Option<DocumentSymbolResponse> =
            self.request::<DocumentSymbolRequest>(params).await?;
        Ok(result.unwrap_or(DocumentSymbolResponse::Flat(vec![])))
    }

    /// Search workspace symbols
    pub async fn workspace_symbols(&self, query: &str) -> Result<Vec<SymbolInformation>> {
        let params = WorkspaceSymbolParams {
            query: query.to_string(),
            work_done_progress_params: WorkDoneProgressParams::default(),
            partial_result_params: PartialResultParams::default(),
        };

        let result: Option<WorkspaceSymbolResponse> =
            self.request::<WorkspaceSymbolRequest>(params).await?;

        Ok(match result {
            Some(WorkspaceSymbolResponse::Flat(symbols)) => symbols,
            Some(WorkspaceSymbolResponse::Nested(symbols)) => {
                // Convert WorkspaceSymbol to SymbolInformation
                symbols
                    .into_iter()
                    .filter_map(|s| {
                        let location = match s.location {
                            lsp_types::OneOf::Left(loc) => loc,
                            lsp_types::OneOf::Right(doc_id) => Location {
                                uri: doc_id.uri,
                                range: Range::default(),
                            },
                        };
                        Some(SymbolInformation {
                            name: s.name,
                            kind: s.kind,
                            tags: s.tags,
                            deprecated: None,
                            location,
                            container_name: s.container_name,
                        })
                    })
                    .collect()
            }
            None => vec![],
        })
    }

    /// Prepare call hierarchy
    pub async fn prepare_call_hierarchy(
        &self,
        path: &Path,
        line: u32,
        character: u32,
    ) -> Result<Vec<CallHierarchyItem>> {
        let uri = path_to_uri(path)?;

        let params = lsp_types::CallHierarchyPrepareParams {
            text_document_position_params: TextDocumentPositionParams {
                text_document: TextDocumentIdentifier { uri },
                position: Position { line, character },
            },
            work_done_progress_params: WorkDoneProgressParams::default(),
        };

        let result: Option<Vec<CallHierarchyItem>> =
            self.request::<CallHierarchyPrepare>(params).await?;
        Ok(result.unwrap_or_default())
    }

    /// Get incoming calls
    pub async fn incoming_calls(
        &self,
        item: CallHierarchyItem,
    ) -> Result<Vec<CallHierarchyIncomingCall>> {
        let params = CallHierarchyIncomingCallsParams {
            item,
            work_done_progress_params: WorkDoneProgressParams::default(),
            partial_result_params: PartialResultParams::default(),
        };

        let result: Option<Vec<CallHierarchyIncomingCall>> =
            self.request::<CallHierarchyIncomingCalls>(params).await?;
        Ok(result.unwrap_or_default())
    }

    /// Get outgoing calls
    pub async fn outgoing_calls(
        &self,
        item: CallHierarchyItem,
    ) -> Result<Vec<CallHierarchyOutgoingCall>> {
        let params = CallHierarchyOutgoingCallsParams {
            item,
            work_done_progress_params: WorkDoneProgressParams::default(),
            partial_result_params: PartialResultParams::default(),
        };

        let result: Option<Vec<CallHierarchyOutgoingCall>> =
            self.request::<CallHierarchyOutgoingCalls>(params).await?;
        Ok(result.unwrap_or_default())
    }

    /// Shutdown the client
    pub async fn shutdown(mut self) {
        info!(server_id = %self.server_id, "Shutting down LSP client");

        // Send shutdown request
        let _ = self.request::<lsp_types::request::Shutdown>(()).await;

        // Send exit notification
        let _ = self.notify::<lsp_types::notification::Exit>(());

        // Signal reader thread to stop
        if let Some(tx) = self.shutdown_tx.take() {
            let _ = tx.send(()).await;
        }
    }

    // ========================================================================
    // Internal Methods
    // ========================================================================

    /// Send a request and wait for response
    async fn request<R>(&self, params: R::Params) -> Result<R::Result>
    where
        R: lsp_types::request::Request,
        R::Params: Serialize,
        R::Result: DeserializeOwned,
    {
        let id = self.request_id.fetch_add(1, Ordering::SeqCst);

        let request = JsonRpcRequest {
            jsonrpc: "2.0",
            id,
            method: R::METHOD,
            params,
        };

        let message = serde_json::to_string(&request)?;
        let header = format!("Content-Length: {}\r\n\r\n", message.len());

        trace!(id = id, method = R::METHOD, "Sending LSP request");

        // Create response channel
        let (tx, rx) = oneshot::channel();

        // Register pending request
        {
            let mut pending = self.pending.write().unwrap();
            pending.insert(id, PendingRequest { sender: tx });
        }

        // Send the request
        {
            let mut writer = self.writer.lock().unwrap();
            writer.write_all(header.as_bytes())?;
            writer.write_all(message.as_bytes())?;
            writer.flush()?;
        }

        // Wait for response with timeout
        let result = tokio::time::timeout(
            std::time::Duration::from_millis(REQUEST_TIMEOUT_MS),
            rx,
        )
        .await
        .map_err(|_| anyhow!("LSP request timed out"))??;

        let value = result?;
        let result: R::Result = serde_json::from_value(value)?;
        Ok(result)
    }

    /// Send a notification (no response expected)
    fn notify<N>(&self, params: N::Params) -> Result<()>
    where
        N: lsp_types::notification::Notification,
        N::Params: Serialize,
    {
        let notification = JsonRpcNotification {
            jsonrpc: "2.0",
            method: N::METHOD,
            params,
        };

        let message = serde_json::to_string(&notification)?;
        let header = format!("Content-Length: {}\r\n\r\n", message.len());

        trace!(method = N::METHOD, "Sending LSP notification");

        let mut writer = self.writer.lock().unwrap();
        writer.write_all(header.as_bytes())?;
        writer.write_all(message.as_bytes())?;
        writer.flush()?;

        Ok(())
    }

    /// Reader loop that processes messages from the server
    fn reader_loop(
        stdout: ChildStdout,
        pending: Arc<RwLock<HashMap<i64, PendingRequest>>>,
        diagnostics: Arc<RwLock<HashMap<PathBuf, Vec<Diagnostic>>>>,
        server_id: String,
    ) {
        let mut reader = BufReader::new(stdout);
        let mut headers = String::new();

        loop {
            headers.clear();

            // Read headers
            let mut content_length: Option<usize> = None;
            loop {
                let mut line = String::new();
                match reader.read_line(&mut line) {
                    Ok(0) => {
                        debug!(server_id = %server_id, "LSP server stdout closed");
                        return;
                    }
                    Ok(_) => {
                        if line == "\r\n" {
                            break;
                        }
                        if line.to_lowercase().starts_with("content-length:") {
                            if let Some(len_str) = line.split(':').nth(1) {
                                content_length = len_str.trim().parse().ok();
                            }
                        }
                    }
                    Err(e) => {
                        error!(server_id = %server_id, error = ?e, "Error reading from LSP server");
                        return;
                    }
                }
            }

            // Read content
            let content_length = match content_length {
                Some(len) => len,
                None => {
                    warn!(server_id = %server_id, "No Content-Length header");
                    continue;
                }
            };

            let mut content = vec![0u8; content_length];
            if let Err(e) = std::io::Read::read_exact(&mut reader, &mut content) {
                error!(server_id = %server_id, error = ?e, "Error reading LSP content");
                continue;
            }

            // Parse message
            let message: JsonRpcMessage = match serde_json::from_slice(&content) {
                Ok(msg) => msg,
                Err(e) => {
                    warn!(server_id = %server_id, error = ?e, "Failed to parse LSP message");
                    continue;
                }
            };

            // Handle response
            if let Some(id) = message.id {
                if message.method.is_none() {
                    // This is a response
                    let mut pending = pending.write().unwrap();
                    if let Some(req) = pending.remove(&id) {
                        let result = if let Some(error) = message.error {
                            Err(anyhow!("LSP error {}: {}", error.code, error.message))
                        } else {
                            Ok(message.result.unwrap_or(serde_json::Value::Null))
                        };
                        let _ = req.sender.send(result);
                    }
                    continue;
                }
            }

            // Handle notification
            if let Some(method) = &message.method {
                match method.as_str() {
                    "textDocument/publishDiagnostics" => {
                        if let Ok(params) =
                            serde_json::from_value::<lsp_types::PublishDiagnosticsParams>(
                                message.params,
                            )
                        {
                            if let Some(path) = uri_to_path(&params.uri) {
                                debug!(
                                    server_id = %server_id,
                                    path = ?path,
                                    count = params.diagnostics.len(),
                                    "Received diagnostics"
                                );
                                diagnostics.write().unwrap().insert(path, params.diagnostics);
                            }
                        }
                    }
                    "window/logMessage" | "window/showMessage" => {
                        // Log server messages
                        if let Ok(params) =
                            serde_json::from_value::<lsp_types::LogMessageParams>(message.params)
                        {
                            debug!(server_id = %server_id, message = %params.message, "LSP server message");
                        }
                    }
                    _ => {
                        trace!(server_id = %server_id, method = %method, "Unhandled LSP notification");
                    }
                }
            }
        }
    }

    /// Build client capabilities
    fn client_capabilities() -> ClientCapabilities {
        ClientCapabilities {
            text_document: Some(lsp_types::TextDocumentClientCapabilities {
                synchronization: Some(lsp_types::TextDocumentSyncClientCapabilities {
                    dynamic_registration: Some(false),
                    will_save: Some(false),
                    will_save_wait_until: Some(false),
                    did_save: Some(true),
                }),
                hover: Some(lsp_types::HoverClientCapabilities {
                    dynamic_registration: Some(false),
                    content_format: Some(vec![MarkupKind::Markdown, MarkupKind::PlainText]),
                }),
                definition: Some(lsp_types::GotoCapability {
                    dynamic_registration: Some(false),
                    link_support: Some(true),
                }),
                references: Some(lsp_types::DynamicRegistrationClientCapabilities {
                    dynamic_registration: Some(false),
                }),
                implementation: Some(lsp_types::GotoCapability {
                    dynamic_registration: Some(false),
                    link_support: Some(true),
                }),
                document_symbol: Some(lsp_types::DocumentSymbolClientCapabilities {
                    dynamic_registration: Some(false),
                    symbol_kind: None,
                    hierarchical_document_symbol_support: Some(true),
                    tag_support: None,
                }),
                publish_diagnostics: Some(lsp_types::PublishDiagnosticsClientCapabilities {
                    related_information: Some(true),
                    tag_support: None,
                    version_support: Some(true),
                    code_description_support: Some(true),
                    data_support: Some(true),
                }),
                call_hierarchy: Some(lsp_types::CallHierarchyClientCapabilities {
                    dynamic_registration: Some(false),
                }),
                ..Default::default()
            }),
            workspace: Some(lsp_types::WorkspaceClientCapabilities {
                workspace_folders: Some(true),
                symbol: Some(lsp_types::WorkspaceSymbolClientCapabilities {
                    dynamic_registration: Some(false),
                    symbol_kind: None,
                    tag_support: None,
                    resolve_support: None,
                }),
                ..Default::default()
            }),
            window: Some(lsp_types::WindowClientCapabilities {
                work_done_progress: Some(true),
                show_message: None,
                show_document: None,
            }),
            ..Default::default()
        }
    }
}