zesdex/crates/zesdex-backend/src/app/lsp/mod.rs

use std::collections::HashMap;
use std::path::Path;
use std::sync::{Arc, Mutex};

mod client;
pub mod provisioner;
pub use client::{path_to_lsp_uri, LspClient};

/// A tracked LSP server entry.
///
/// Holds the spawn metadata and a shared handle to the connected
/// [`LspClient`]. The `Arc<Mutex<...>>` is cloned by callers that need
/// to issue LSP requests from threads or async tasks.
#[derive(Clone)]
pub struct LspServer {
    pub language_id: String,
    pub client: Arc<Mutex<LspClient>>,
}

/// Metadata for a document the manager has announced to an LSP server.
///
/// Used to track the current `version` and `languageId` for files
/// already sent via `textDocument/didOpen`, so subsequent edits can be
/// replayed as `textDocument/didChange` notifications.
#[derive(Clone)]
pub struct OpenDoc {
    pub language: String,
    pub version: i32,
}

/// Central registry of connected LSP servers and per-extension routing.
///
/// Flow: caller calls `connect*` -> client spawned -> entry pushed to
/// `servers` -> `extension_registry` is populated by `register_extensions`.
/// File edits route through `extension_registry` and are dispatched as
/// `didOpen` / `didChange` notifications.
#[derive(Clone)]
pub struct LspManager {
    pub servers: Vec<LspServer>,
    /// Maps file extension (".rs", ".ts", ...) -> language id.
    pub extension_registry: HashMap<String, String>,
    /// Maps document URI -> tracked open document state.
    pub open_files: HashMap<String, OpenDoc>,
}

impl LspManager {
    /// Create an empty manager with no connected servers and empty registries.
    pub fn new() -> Self {
        LspManager {
            servers: Vec::new(),
            extension_registry: HashMap::new(),
            open_files: HashMap::new(),
        }
    }

    /// Spawn an LSP server and register it under `language_id`.
    ///
    /// Fails if a server with the same `language_id` is already connected.
    pub fn connect(
        &mut self,
        command: &str,
        args: &[String],
        language_id: &str,
    ) -> anyhow::Result<()> {
        if self.servers.iter().any(|s| s.language_id == language_id) {
            anyhow::bail!("LSP server for language '{language_id}' is already connected");
        }
        let client = LspClient::spawn(command, args)?;
        self.servers.push(LspServer {
            language_id: language_id.to_string(),
            client: Arc::new(Mutex::new(client)),
        });
        Ok(())
    }

    /// Return a clone of the `Arc<Mutex<LspClient>>` for a connected server.
    ///
    /// Cloning the `Arc` lets callers issue requests without holding a
    /// borrow on the manager.
    pub fn get_client(&self, language_id: &str) -> Option<Arc<Mutex<LspClient>>> {
        self.servers
            .iter()
            .find(|s| s.language_id == language_id)
            .map(|s| s.client.clone())
    }

    /// Shut down and remove a server by language. Returns true if it existed.
    pub fn disconnect(&mut self, language_id: &str) -> bool {
        if let Some(server) = self.servers.iter().find(|s| s.language_id == language_id) {
            if let Ok(mut client) = server.client.lock() {
                client.shutdown();
            }
        }
        let len = self.servers.len();
        self.servers.retain(|s| s.language_id != language_id);
        self.servers.len() < len
    }

    /// Return the language id (e.g. "rust") registered for `language_id`.
    pub fn get_language_id(&self, language_id: &str) -> Option<String> {
        self.servers
            .iter()
            .find(|s| s.language_id == language_id)
            .map(|s| s.language_id.clone())
    }

    /// Register a set of file extensions for an already-connected server.
    ///
    /// Flow: for each `ext`, write `language_id` into `extension_registry`.
    /// Re-registration overwrites the previous target. Unknown language IDs
    /// are accepted at this layer — caller must ensure a server for
    /// `language_id` is connected or will be connected later.
    pub fn register_extensions(&mut self, language_id: &str, extensions: &[&str]) {
        for ext in extensions {
            self.extension_registry
                .insert(ext.to_string(), language_id.to_string());
        }
    }

    /// Notify the relevant LSP server that a file's contents have changed.
    ///
    /// Flow: resolve language by extension from the registry -> read file contents ->
    /// either send `didOpen` (first time) or `didChange` (already tracked)
    /// -> update `open_files` with the new version.
    ///
    /// Non-critical failures (file missing, server unreachable, send
    /// error) are logged with `tracing::warn!` rather than propagated,
    /// so a stale notification cannot abort the calling flow.
    pub fn did_change_file(&mut self, path: &Path) {
        let Some(ext) = path
            .extension()
            .and_then(|e| e.to_str())
            .map(|s| format!(".{s}"))
        else {
            tracing::warn!("did_change_file: path has no extension: {:?}", path);
            return;
        };

        let Some(language_id) = self.extension_registry.get(&ext).cloned() else {
            tracing::warn!(
                "did_change_file: no LSP server registered for extension '{}'",
                ext
            );
            return;
        };

        let uri = path_to_lsp_uri(&path.to_string_lossy());

        let text = match std::fs::read_to_string(path) {
            Ok(t) => t,
            Err(e) => {
                tracing::warn!("did_change_file: failed to read {:?}: {}", path, e);
                return;
            }
        };

        let Some(client) = self.get_client(&language_id) else {
            tracing::warn!("did_change_file: no client for language '{}'", language_id);
            return;
        };

        let next_version = match self.open_files.get(&uri) {
            Some(existing) => existing.version + 1,
            None => 1,
        };

        let send_result = {
            let mut client = match client.lock() {
                Ok(c) => c,
                Err(e) => {
                    tracing::warn!(
                        "did_change_file: client mutex poisoned for '{}': {}",
                        language_id,
                        e
                    );
                    return;
                }
            };
            if self.open_files.contains_key(&uri) {
                client.did_change(&uri, next_version, &text)
            } else {
                client.did_open(&uri, &language_id, next_version, &text)
            }
        };

        if let Err(e) = send_result {
            tracing::warn!(
                "did_change_file: failed to notify '{}' for {}: {}",
                language_id,
                uri,
                e
            );
            return;
        }

        self.open_files.insert(
            uri.clone(),
            OpenDoc {
                language: language_id,
                version: next_version,
            },
        );
    }

    /// Shut down every connected server and clear the server list.
    ///
    /// Flow: iterate `servers` -> call `client.shutdown()` on each ->
    /// drop the vec. Failures from individual shutdowns are swallowed
    /// because the goal is best-effort termination during teardown.
    pub fn shutdown_all(&mut self) {
        for server in &self.servers {
            if let Ok(mut client) = server.client.lock() {
                client.shutdown();
            }
        }
        self.servers.clear();
    }

    /// Snapshot the connected servers as `(language_id, has_open_docs)` pairs.
    ///
    /// `has_open_docs` is true if any tracked `OpenDoc` was registered
    /// against this server's clients. Useful for status displays.
    pub fn list_servers(&self) -> Vec<(String, bool)> {
        self.servers
            .iter()
            .map(|s| {
                let lang = s.language_id.clone();
                let has_open = self
                    .open_files
                    .values()
                    .any(|d| d.language == s.language_id);
                (lang, has_open)
            })
            .collect()
    }

    /// Connect an LSP server and register its default extensions in one call.
    ///
    /// Flow: invoke `connect` -> on success, register `extensions` against
    /// `language_id` in `extension_registry`. If `connect` fails, the registries
    /// are left untouched and the error is propagated.
    pub fn connect_with_extensions(
        &mut self,
        command: &str,
        args: &[String],
        language_id: &str,
        extensions: &[&str],
    ) -> anyhow::Result<()> {
        self.connect(command, args, language_id)?;
        self.register_extensions(language_id, extensions);
        Ok(())
    }
}

impl Default for LspManager {
    fn default() -> Self {
        Self::new()
    }
}