zesdex/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 {
    #[allow(dead_code)]
    pub name: String,
    #[allow(dead_code)]
    pub command: String,
    #[allow(dead_code)]
    pub args: Vec<String>,
    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 `find_server_for_path` / `find_server_for_extension`
/// and are dispatched as `didOpen` / `didChange` notifications.
#[derive(Clone)]
pub struct LspManager {
    pub servers: Vec<LspServer>,
    /// Maps file extension (".rs", ".ts", ...) -> server name.
    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 `name`.
    ///
    /// Fails if a server with the same name is already connected.
    pub fn connect(
        &mut self,
        name: &str,
        command: &str,
        args: &[String],
        language_id: &str,
    ) -> anyhow::Result<()> {
        if self.servers.iter().any(|s| s.name == name) {
            anyhow::bail!("LSP server '{name}' is already connected");
        }
        let client = LspClient::spawn(command, args)?;
        self.servers.push(LspServer {
            name: name.to_string(),
            command: command.to_string(),
            args: args.to_vec(),
            language_id: language_id.to_string(),
            client: Arc::new(Mutex::new(client)),
        });
        Ok(())
    }

    /// Look up a connected server by name and return a reference to its entry.
    #[allow(dead_code)]
    pub fn find_server(&self, name: &str) -> Option<&LspServer> {
        self.servers.iter().find(|s| s.name == name)
    }

    /// 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, name: &str) -> Option<Arc<Mutex<LspClient>>> {
        self.servers.iter().find(|s| s.name == name).map(|s| s.client.clone())
    }

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

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

    /// Resolve an extension (".rs", ".ts", ...) to its server's client.
    ///
    /// Flow: lookup `extension_registry` -> resolve server name -> clone client.
    /// Returns `None` if no server has been registered for `ext`.
    #[allow(dead_code)]
    pub fn find_server_for_extension(&self, ext: &str) -> Option<Arc<Mutex<LspClient>>> {
        self.extension_registry
            .get(ext)
            .and_then(|name| self.get_client(name))
    }

    /// Resolve a file path to its server's client by extension.
    ///
    /// Flow: extract the extension from `path` -> delegate to
    /// `find_server_for_extension`. Files without an extension or with
    /// an unmapped extension return `None`.
    #[allow(dead_code)]
    pub fn find_server_for_path(&self, path: &Path) -> Option<Arc<Mutex<LspClient>>> {
        path.extension()
            .and_then(|e| e.to_str())
            .map(|s| format!(".{s}"))
            .and_then(|ext| self.find_server_for_extension(&ext))
    }

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

    /// Return the registered server name for a given language id.
    ///
    /// Flow: scan `servers` for the first entry whose `language_id` matches.
    /// Used when callers have a language hint rather than a file path.
    #[allow(dead_code)]
    pub fn get_server_name(&self, language: &str) -> Option<String> {
        self.servers
            .iter()
            .find(|s| s.language_id == language)
            .map(|s| s.name.clone())
    }

    /// Notify the relevant LSP server that a file's contents have changed.
    ///
    /// Flow: resolve server by extension -> 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 server_name = if let Some(name) = self.extension_registry.get(&ext) { name.clone() } 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 language_id = self
            .get_language_id(&server_name)
            .unwrap_or_else(|| "plaintext".to_string());

        let Some(client) = self.get_client(&server_name) else {
            tracing::warn!("did_change_file: server '{}' has no client", server_name);
            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 '{}': {}", server_name, 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 {}: {}",
                server_name,
                uri,
                e
            );
            return;
        }

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

    }

    /// Record that `server_name` has an open document at `uri`.
    ///
    /// Flow: insert/overwrite the `OpenDoc` entry in `open_files`.
    /// Does not contact the LSP server — pure local bookkeeping.
    #[allow(dead_code)]
    pub fn track_open_doc(&mut self, server_name: &str, uri: &str, language: &str, version: i32) {
        // server_name retained for future routing extensions; not stored today.
        let _ = server_name;
        self.open_files.insert(
            uri.to_string(),
            OpenDoc {
                language: language.to_string(),
                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 `(name, language_id, has_open_docs)` triples.
    ///
    /// `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, String, bool)> {
        self.servers
            .iter()
            .map(|s| {
                let name = s.name.clone();
                let lang = s.language_id.clone();
                let has_open = self
                    .open_files
                    .values()
                    .any(|d| d.language == s.language_id);
                (name, lang, has_open)
            })
            .collect()
    }

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

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