lsp-mcp 0.1.0

use crate::error::{Language, LspMcpError, Result};
use crate::lsp::config::LanguageServerConfig;
use crate::lsp::transport::StdioTransport;
use lsp_types::*;
use parking_lot::RwLock;
use std::path::PathBuf;
use std::sync::Arc;
use tracing::{debug, info};
use url::Url;

/// Wrapper around a single language server process
pub struct LspClient {
    /// Language this server handles
    language: Language,
    /// Workspace root path
    workspace_root: PathBuf,
    /// Transport for communication
    transport: Arc<StdioTransport>,
    /// Server capabilities (after initialization)
    capabilities: RwLock<Option<ServerCapabilities>>,
    /// Server info
    server_info: RwLock<Option<ServerInfo>>,
    /// Whether the server has been initialized
    initialized: RwLock<bool>,
}

impl LspClient {
    /// Create a new LSP client by spawning the language server
    pub fn new(
        config: &LanguageServerConfig,
        language: Language,
        workspace_root: PathBuf,
    ) -> Result<Self> {
        info!(
            "Starting {} for workspace: {:?}",
            config.name, workspace_root
        );

        let transport = StdioTransport::spawn(
            &config.command,
            &config.args,
            &config.env,
            &workspace_root,
        )?;

        Ok(Self {
            language,
            workspace_root,
            transport: Arc::new(transport),
            capabilities: RwLock::new(None),
            server_info: RwLock::new(None),
            initialized: RwLock::new(false),
        })
    }

    /// Initialize the language server
    pub fn initialize(&self) -> Result<InitializeResult> {
        if *self.initialized.read() {
            if let Some(caps) = self.capabilities.read().clone() {
                return Ok(InitializeResult {
                    capabilities: caps,
                    server_info: self.server_info.read().clone(),
                });
            }
        }

        debug!("Initializing language server for {:?}", self.workspace_root);

        let workspace_uri = Url::from_file_path(&self.workspace_root)
            .map_err(|_| LspMcpError::InvalidPosition { line: 0, character: 0 })?;

        let params = InitializeParams {
            process_id: Some(std::process::id()),
            root_uri: Some(workspace_uri.clone()),
            root_path: Some(self.workspace_root.to_string_lossy().to_string()),
            capabilities: ClientCapabilities {
                text_document: Some(TextDocumentClientCapabilities {
                    hover: Some(HoverClientCapabilities {
                        content_format: Some(vec![MarkupKind::Markdown, MarkupKind::PlainText]),
                        ..Default::default()
                    }),
                    completion: Some(CompletionClientCapabilities {
                        completion_item: Some(CompletionItemCapability {
                            documentation_format: Some(vec![MarkupKind::Markdown, MarkupKind::PlainText]),
                            ..Default::default()
                        }),
                        ..Default::default()
                    }),
                    definition: Some(GotoCapability {
                        link_support: Some(true),
                        ..Default::default()
                    }),
                    references: Some(DynamicRegistrationClientCapabilities {
                        dynamic_registration: Some(false),
                    }),
                    document_symbol: Some(DocumentSymbolClientCapabilities {
                        hierarchical_document_symbol_support: Some(true),
                        ..Default::default()
                    }),
                    publish_diagnostics: Some(PublishDiagnosticsClientCapabilities {
                        related_information: Some(true),
                        ..Default::default()
                    }),
                    rename: Some(RenameClientCapabilities {
                        prepare_support: Some(true),
                        ..Default::default()
                    }),
                    ..Default::default()
                }),
                workspace: Some(WorkspaceClientCapabilities {
                    workspace_folders: Some(true),
                    symbol: Some(WorkspaceSymbolClientCapabilities {
                        ..Default::default()
                    }),
                    ..Default::default()
                }),
                ..Default::default()
            },
            workspace_folders: Some(vec![WorkspaceFolder {
                uri: workspace_uri,
                name: self.workspace_root
                    .file_name()
                    .map(|n| n.to_string_lossy().to_string())
                    .unwrap_or_else(|| "workspace".to_string()),
            }]),
            ..Default::default()
        };

        let result: InitializeResult = self.transport.request("initialize", params)?;

        // Store capabilities
        *self.capabilities.write() = Some(result.capabilities.clone());
        *self.server_info.write() = result.server_info.clone();

        // Send initialized notification
        self.transport.notify("initialized", InitializedParams {})?;

        *self.initialized.write() = true;

        info!(
            "Language server initialized: {:?}",
            result.server_info.as_ref().map(|i| &i.name)
        );

        Ok(result)
    }

    /// Open a document in the language server
    pub fn open_document(&self, uri: &Url, text: &str, language_id: &str) -> Result<()> {
        debug!("Opening document: {}", uri);

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

        self.transport.notify("textDocument/didOpen", params)
    }

    /// Close a document
    pub fn close_document(&self, uri: &Url) -> Result<()> {
        debug!("Closing document: {}", uri);

        let params = DidCloseTextDocumentParams {
            text_document: TextDocumentIdentifier { uri: uri.clone() },
        };

        self.transport.notify("textDocument/didClose", params)
    }

    /// Get hover information
    pub fn hover(&self, uri: &Url, position: Position) -> Result<Option<Hover>> {
        let params = HoverParams {
            text_document_position_params: TextDocumentPositionParams {
                text_document: TextDocumentIdentifier { uri: uri.clone() },
                position,
            },
            work_done_progress_params: Default::default(),
        };

        self.transport.request("textDocument/hover", params)
    }

    /// Go to definition
    pub fn goto_definition(&self, uri: &Url, position: Position) -> Result<Option<GotoDefinitionResponse>> {
        let params = GotoDefinitionParams {
            text_document_position_params: TextDocumentPositionParams {
                text_document: TextDocumentIdentifier { uri: uri.clone() },
                position,
            },
            work_done_progress_params: Default::default(),
            partial_result_params: Default::default(),
        };

        self.transport.request("textDocument/definition", params)
    }

    /// Find references
    pub fn find_references(
        &self,
        uri: &Url,
        position: Position,
        include_declaration: bool,
    ) -> Result<Option<Vec<Location>>> {
        let params = ReferenceParams {
            text_document_position: TextDocumentPositionParams {
                text_document: TextDocumentIdentifier { uri: uri.clone() },
                position,
            },
            context: ReferenceContext {
                include_declaration,
            },
            work_done_progress_params: Default::default(),
            partial_result_params: Default::default(),
        };

        self.transport.request("textDocument/references", params)
    }

    /// Get completions
    pub fn completion(
        &self,
        uri: &Url,
        position: Position,
        trigger_character: Option<String>,
    ) -> Result<Option<CompletionResponse>> {
        let params = CompletionParams {
            text_document_position: TextDocumentPositionParams {
                text_document: TextDocumentIdentifier { uri: uri.clone() },
                position,
            },
            context: trigger_character.map(|tc| CompletionContext {
                trigger_kind: CompletionTriggerKind::TRIGGER_CHARACTER,
                trigger_character: Some(tc),
            }),
            work_done_progress_params: Default::default(),
            partial_result_params: Default::default(),
        };

        self.transport.request("textDocument/completion", params)
    }

    /// Get document symbols
    pub fn document_symbols(&self, uri: &Url) -> Result<Option<DocumentSymbolResponse>> {
        let params = DocumentSymbolParams {
            text_document: TextDocumentIdentifier { uri: uri.clone() },
            work_done_progress_params: Default::default(),
            partial_result_params: Default::default(),
        };

        self.transport.request("textDocument/documentSymbol", params)
    }

    /// Search workspace symbols
    pub fn workspace_symbols(&self, query: &str) -> Result<Option<Vec<SymbolInformation>>> {
        #[allow(deprecated)]
        let params = WorkspaceSymbolParams {
            query: query.to_string(),
            work_done_progress_params: Default::default(),
            partial_result_params: Default::default(),
        };

        self.transport.request("workspace/symbol", params)
    }

    /// Prepare rename
    pub fn prepare_rename(&self, uri: &Url, position: Position) -> Result<Option<PrepareRenameResponse>> {
        let params = TextDocumentPositionParams {
            text_document: TextDocumentIdentifier { uri: uri.clone() },
            position,
        };

        self.transport.request("textDocument/prepareRename", params)
    }

    /// Rename symbol
    pub fn rename(&self, uri: &Url, position: Position, new_name: &str) -> Result<Option<WorkspaceEdit>> {
        let params = RenameParams {
            text_document_position: TextDocumentPositionParams {
                text_document: TextDocumentIdentifier { uri: uri.clone() },
                position,
            },
            new_name: new_name.to_string(),
            work_done_progress_params: Default::default(),
        };

        self.transport.request("textDocument/rename", params)
    }

    /// Shutdown the language server
    pub fn shutdown(&self) -> Result<()> {
        info!("Shutting down language server for {:?}", self.workspace_root);

        // Send shutdown request
        let _: () = self.transport.request("shutdown", ())?;

        // Send exit notification
        self.transport.notify("exit", ())?;

        Ok(())
    }

    /// Get the language this client handles
    pub fn language(&self) -> Language {
        self.language
    }

    /// Get the workspace root
    pub fn workspace_root(&self) -> &PathBuf {
        &self.workspace_root
    }

    /// Check if the server supports a capability
    pub fn supports_hover(&self) -> bool {
        self.capabilities
            .read()
            .as_ref()
            .map(|c| c.hover_provider.is_some())
            .unwrap_or(false)
    }

    pub fn supports_definition(&self) -> bool {
        self.capabilities
            .read()
            .as_ref()
            .map(|c| c.definition_provider.is_some())
            .unwrap_or(false)
    }

    pub fn supports_references(&self) -> bool {
        self.capabilities
            .read()
            .as_ref()
            .map(|c| c.references_provider.is_some())
            .unwrap_or(false)
    }

    pub fn supports_completion(&self) -> bool {
        self.capabilities
            .read()
            .as_ref()
            .map(|c| c.completion_provider.is_some())
            .unwrap_or(false)
    }

    pub fn supports_document_symbol(&self) -> bool {
        self.capabilities
            .read()
            .as_ref()
            .map(|c| c.document_symbol_provider.is_some())
            .unwrap_or(false)
    }

    pub fn supports_workspace_symbol(&self) -> bool {
        self.capabilities
            .read()
            .as_ref()
            .map(|c| c.workspace_symbol_provider.is_some())
            .unwrap_or(false)
    }

    pub fn supports_rename(&self) -> bool {
        self.capabilities
            .read()
            .as_ref()
            .map(|c| c.rename_provider.is_some())
            .unwrap_or(false)
    }
}