leta_daemon/

session.rs

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

use fastrace::trace;
use leta_config::Config;
use leta_fs::{get_language_id, path_to_uri, read_file_content, uri_to_path};
use leta_lsp::LspClient;
use leta_servers::{get_server_env, get_server_for_file, get_server_for_language, ServerConfig};
use serde_json::Value;
use tokio::sync::{Mutex, RwLock};
use tracing::{debug, info};

#[derive(Clone)]
pub struct OpenDocument {
    _uri: String,
    _version: i32,
    pub content: String,
    _language_id: String,
}

pub struct Workspace {
    root: PathBuf,
    server_config: &'static ServerConfig,
    client: Option<Arc<LspClient>>,
    open_documents: HashMap<String, OpenDocument>,
    startup_stats: Option<leta_types::ServerStartupStats>,
}

impl Workspace {
    pub fn new(root: PathBuf, server_config: &'static ServerConfig) -> Self {
        Self {
            root,
            server_config,
            client: None,
            open_documents: HashMap::new(),
            startup_stats: None,
        }
    }

    pub fn startup_stats(&self) -> Option<leta_types::ServerStartupStats> {
        self.startup_stats.clone()
    }

    pub fn client(&self) -> Option<Arc<LspClient>> {
        self.client.clone()
    }

    pub fn server_name(&self) -> &str {
        self.server_config.name
    }

    pub fn open_document_uris(&self) -> Vec<String> {
        self.open_documents.keys().cloned().collect()
    }

    #[trace]
    pub async fn start_server(&mut self) -> Result<leta_types::ServerStartupStats, String> {
        let total_start = std::time::Instant::now();

        if self.client.is_some() {
            return Ok(leta_types::ServerStartupStats {
                server_name: self.server_config.name.to_string(),
                workspace_root: self.root.to_string_lossy().to_string(),
                init_time_ms: 0,
                ready_time_ms: 0,
                total_time_ms: 0,
                functions: Vec::new(),
            });
        }

        info!(
            "Starting {} for {}",
            self.server_config.name,
            self.root.display()
        );

        let env = get_server_env();
        let init_options = self.get_init_options();

        let cmd: Vec<&str> = self.server_config.command.to_vec();

        let start_time = std::time::Instant::now();
        match LspClient::start(&cmd, &self.root, self.server_config.name, env, init_options).await {
            Ok(client) => {
                let init_time = start_time.elapsed();

                if self.server_config.name == "clangd" {
                    client.wait_for_indexing(60).await;
                    self.ensure_workspace_indexed(&client).await;
                }

                self.client = Some(client);
                let total_time = total_start.elapsed();
                let ready_time = std::time::Duration::ZERO;

                info!(
                    "Server {} initialized and ready in {:?}",
                    self.server_config.name, total_time
                );

                let stats = leta_types::ServerStartupStats {
                    server_name: self.server_config.name.to_string(),
                    workspace_root: self.root.to_string_lossy().to_string(),
                    init_time_ms: init_time.as_millis() as u64,
                    ready_time_ms: ready_time.as_millis() as u64,
                    total_time_ms: total_time.as_millis() as u64,
                    functions: Vec::new(),
                };
                self.startup_stats = Some(stats.clone());
                Ok(stats)
            }
            Err(e) => {
                let mut msg = format!(
                    "Language server '{}' for {} failed to start: {}",
                    self.server_config.name,
                    self.server_config.languages.join(", "),
                    e
                );
                if let Some(install_cmd) = self.server_config.install_cmd {
                    msg.push_str(&format!(
                        "\n\nTo install {}, run:\n  {}\n\nIf you just installed it, run `leta daemon restart` to pick up PATH changes.",
                        self.server_config.name, install_cmd
                    ));
                }
                Err(msg)
            }
        }
    }

    fn get_init_options(&self) -> Option<Value> {
        if self.server_config.name == "gopls" {
            Some(serde_json::json!({
                "linksInHover": false,
            }))
        } else {
            None
        }
    }

    /// Open and close all source files to ensure clangd indexes them.
    ///
    /// clangd does lazy indexing - it only indexes files when they're opened.
    /// This means documentSymbol won't work on files that haven't been opened yet.
    /// We work around this by opening all source files during initialization.
    #[trace]
    async fn ensure_workspace_indexed(&mut self, client: &Arc<LspClient>) {
        let walkdir_start = std::time::Instant::now();
        let source_extensions = [".c", ".h", ".cpp", ".hpp", ".cc", ".cxx", ".hxx"];
        let exclude_dirs: std::collections::HashSet<&str> =
            ["build", ".git", "node_modules"].into_iter().collect();

        let mut files_to_index = Vec::new();
        for entry in jwalk::WalkDir::new(&self.root).process_read_dir(
            move |_depth, _path, _state, children| {
                children.retain(|entry| {
                    entry
                        .as_ref()
                        .map(|e| {
                            let name = e.file_name().to_string_lossy();
                            !exclude_dirs.contains(name.as_ref())
                        })
                        .unwrap_or(false)
                });
            },
        ) {
            let Ok(entry) = entry else { continue };
            if entry.file_type().is_file() {
                let path = entry.path();
                if let Some(ext) = path.extension().and_then(|e| e.to_str()) {
                    if source_extensions
                        .iter()
                        .any(|s| s.trim_start_matches('.') == ext)
                    {
                        files_to_index.push(path);
                    }
                }
            }
        }
        let walkdir_elapsed = walkdir_start.elapsed();

        if files_to_index.is_empty() {
            fastrace::local::LocalSpan::add_properties(|| {
                [
                    ("files_found", "0".to_string()),
                    (
                        "walkdir_ms",
                        format!("{:.1}", walkdir_elapsed.as_secs_f64() * 1000.0),
                    ),
                ]
            });
            return;
        }

        info!("Pre-indexing {} files for clangd", files_to_index.len());

        let open_start = std::time::Instant::now();
        for file_path in &files_to_index {
            let _ = self.ensure_document_open(file_path).await;
        }
        let open_elapsed = open_start.elapsed();

        client.wait_for_indexing(30).await;

        info!(
            "Pre-indexing complete, closing {} documents",
            self.open_documents.len()
        );

        let close_start = std::time::Instant::now();
        self.close_all_documents().await;
        let close_elapsed = close_start.elapsed();

        fastrace::local::LocalSpan::add_properties(|| {
            [
                ("files_found", files_to_index.len().to_string()),
                (
                    "walkdir_ms",
                    format!("{:.1}", walkdir_elapsed.as_secs_f64() * 1000.0),
                ),
                (
                    "open_docs_ms",
                    format!("{:.1}", open_elapsed.as_secs_f64() * 1000.0),
                ),
                (
                    "close_docs_ms",
                    format!("{:.1}", close_elapsed.as_secs_f64() * 1000.0),
                ),
            ]
        });
    }

    #[trace]
    pub async fn stop_server(&mut self) {
        if let Some(client) = self.client.take() {
            info!("Stopping {}", self.server_config.name);
            let _ = client.stop().await;
        }
        self.open_documents.clear();
    }

    pub async fn ensure_document_open(&mut self, path: &Path) -> Result<(), String> {
        self.cleanup_deleted_documents().await;

        let uri = path_to_uri(path);

        if let Some(doc) = self.open_documents.get(&uri) {
            let current_content = read_file_content(path).map_err(|e| e.to_string())?;
            if current_content != doc.content {
                self.close_document(path).await;
            } else {
                return Ok(());
            }
        }

        let content = read_file_content(path).map_err(|e| e.to_string())?;
        let language_id = get_language_id(path).to_string();

        let doc = OpenDocument {
            _uri: uri.clone(),
            _version: 1,
            content: content.clone(),
            _language_id: language_id.clone(),
        };

        self.open_documents.insert(uri.clone(), doc);

        if let Some(client) = &self.client {
            let params = serde_json::json!({
                "textDocument": {
                    "uri": uri,
                    "languageId": language_id,
                    "version": 1,
                    "text": content,
                }
            });
            let _ = client
                .send_notification("textDocument/didOpen", params)
                .await;

            // ruby-lsp processes messages asynchronously in a queue, so we need to ensure
            // the didOpen is fully processed before subsequent operations can succeed.
            // We do this by sending a simple request and waiting for its response.
            if client.server_name() == "ruby-lsp" {
                let symbol_params = serde_json::json!({
                    "textDocument": {"uri": uri}
                });
                let _ = client
                    .send_request_raw("textDocument/documentSymbol", symbol_params)
                    .await;
            }
        }

        Ok(())
    }

    #[trace]
    pub async fn close_document(&mut self, path: &Path) {
        let uri = path_to_uri(path);
        if self.open_documents.remove(&uri).is_none() {
            return;
        }

        if let Some(client) = &self.client {
            let params = serde_json::json!({
                "textDocument": {"uri": uri}
            });
            let _ = client
                .send_notification("textDocument/didClose", params)
                .await;
        }
    }

    #[trace]
    pub async fn close_all_documents(&mut self) {
        if let Some(client) = &self.client {
            for uri in self.open_documents.keys() {
                let params = serde_json::json!({
                    "textDocument": {"uri": uri}
                });
                let _ = client
                    .send_notification("textDocument/didClose", params)
                    .await;
            }
        }
        self.open_documents.clear();
    }

    pub async fn cleanup_deleted_documents(&mut self) {
        let deleted_uris: Vec<String> = self
            .open_documents
            .keys()
            .filter(|uri| !uri_to_path(uri).exists())
            .cloned()
            .collect();

        for uri in deleted_uris {
            self.open_documents.remove(&uri);
            if let Some(client) = &self.client {
                let params = serde_json::json!({
                    "textDocument": {"uri": uri}
                });
                let _ = client
                    .send_notification("textDocument/didClose", params)
                    .await;
            }
        }
    }
}

type StartupLockKey = (PathBuf, String);
type StartupLocks = Mutex<HashMap<StartupLockKey, Arc<Mutex<()>>>>;

pub struct Session {
    workspaces: RwLock<HashMap<PathBuf, HashMap<String, Workspace>>>,
    config: RwLock<Config>,
    workspace_profiling: RwLock<Vec<leta_types::WorkspaceProfilingData>>,
    startup_locks: StartupLocks,
}

impl Session {
    pub fn new(config: Config) -> Self {
        Self {
            workspaces: RwLock::new(HashMap::new()),
            config: RwLock::new(config),
            workspace_profiling: RwLock::new(Vec::new()),
            startup_locks: Mutex::new(HashMap::new()),
        }
    }

    pub async fn add_workspace_profiling(&self, data: leta_types::WorkspaceProfilingData) {
        let mut profiling = self.workspace_profiling.write().await;
        profiling.retain(|p| p.workspace_root != data.workspace_root);
        profiling.push(data);
    }

    pub async fn get_workspace_profiling(&self) -> Vec<leta_types::WorkspaceProfilingData> {
        self.workspace_profiling.read().await.clone()
    }

    #[trace]
    pub async fn config(&self) -> Config {
        self.config.read().await.clone()
    }

    #[trace]
    pub async fn get_or_create_workspace(
        &self,
        file_path: &Path,
        workspace_root: &Path,
    ) -> Result<WorkspaceHandle<'_>, String> {
        let server_config = {
            let config = self.config.read().await;
            get_server_for_file(file_path, Some(&config))
                .ok_or_else(|| format!("No language server found for {}", file_path.display()))?
        };
        // config lock dropped here before acquiring workspace lock

        self.get_or_create_workspace_for_server(workspace_root, server_config)
            .await
    }

    #[trace]
    pub async fn get_or_create_workspace_for_language(
        &self,
        language_id: &str,
        workspace_root: &Path,
    ) -> Result<WorkspaceHandle<'_>, String> {
        let server_config = {
            let config = self.config.read().await;
            get_server_for_language(language_id, Some(&config))
                .ok_or_else(|| format!("No language server found for language {}", language_id))?
        };
        // config lock dropped here before acquiring workspace lock

        self.get_or_create_workspace_for_server(workspace_root, server_config)
            .await
    }

    #[trace]
    async fn get_or_create_workspace_for_server(
        &self,
        workspace_root: &Path,
        server_config: &'static ServerConfig,
    ) -> Result<WorkspaceHandle<'_>, String> {
        let workspace_root = workspace_root
            .canonicalize()
            .unwrap_or_else(|_| workspace_root.to_path_buf());

        // Get or create a per-workspace/server lock to prevent concurrent starts
        let startup_lock = {
            let mut locks = self.startup_locks.lock().await;
            let key = (workspace_root.clone(), server_config.name.to_string());
            locks
                .entry(key)
                .or_insert_with(|| Arc::new(Mutex::new(())))
                .clone()
        };

        // Hold the startup lock while we check and potentially start the server
        let _startup_guard = startup_lock.lock().await;

        // Check if workspace exists (read lock only)
        let needs_create = {
            let workspaces = self.workspaces.read().await;
            if let Some(servers) = workspaces.get(&workspace_root) {
                if let Some(ws) = servers.get(server_config.name) {
                    ws.client.is_none() // needs restart
                } else {
                    true // needs create
                }
            } else {
                true // needs create
            }
        };

        if needs_create {
            debug!(
                "Starting {} for {}",
                server_config.name,
                workspace_root.display()
            );
            let mut new_workspace = Workspace::new(workspace_root.clone(), server_config);
            new_workspace.start_server().await?;

            // Insert with write lock (quick operation)
            let mut workspaces = self.workspaces.write().await;
            let servers = workspaces
                .entry(workspace_root.clone())
                .or_insert_with(HashMap::new);
            servers.insert(server_config.name.to_string(), new_workspace);
        }

        Ok(WorkspaceHandle {
            session: self,
            workspace_root,
            server_name: server_config.name.to_string(),
        })
    }

    #[allow(dead_code)]
    #[trace]
    pub async fn get_workspace_for_file(&self, file_path: &Path) -> Option<WorkspaceHandle<'_>> {
        let file_path = file_path
            .canonicalize()
            .unwrap_or_else(|_| file_path.to_path_buf());
        let config = self.config.read().await;
        let server_config = get_server_for_file(&file_path, Some(&config))?;

        let workspaces = self.workspaces.read().await;
        for (root, servers) in workspaces.iter() {
            if file_path.starts_with(root) && servers.contains_key(server_config.name) {
                return Some(WorkspaceHandle {
                    session: self,
                    workspace_root: root.clone(),
                    server_name: server_config.name.to_string(),
                });
            }
        }
        None
    }

    #[trace]
    pub async fn list_workspaces(&self) -> Vec<(String, String, Option<u32>, Vec<String>)> {
        let workspaces = self.workspaces.read().await;
        let mut result = Vec::new();

        for (root, servers) in workspaces.iter() {
            for (_, ws) in servers.iter() {
                let server_pid = ws.client.as_ref().and_then(|c| c.pid());
                result.push((
                    root.to_string_lossy().to_string(),
                    ws.server_name().to_string(),
                    server_pid,
                    ws.open_document_uris(),
                ));
            }
        }

        result
    }

    #[trace]
    pub async fn restart_workspace(&self, root: &Path) -> Result<Vec<String>, String> {
        let root = root.canonicalize().unwrap_or_else(|_| root.to_path_buf());
        let mut workspaces = self.workspaces.write().await;

        let mut restarted = Vec::new();
        if let Some(servers) = workspaces.get_mut(&root) {
            for (name, workspace) in servers.iter_mut() {
                workspace.stop_server().await;
                workspace.start_server().await?;
                restarted.push(name.clone());
            }
        }
        Ok(restarted)
    }

    #[trace]
    pub async fn remove_workspace(&self, root: &Path) -> Result<Vec<String>, String> {
        let root = root.canonicalize().unwrap_or_else(|_| root.to_path_buf());
        let mut workspaces = self.workspaces.write().await;

        let mut stopped = Vec::new();
        if let Some(mut servers) = workspaces.remove(&root) {
            for (name, mut workspace) in servers.drain() {
                workspace.stop_server().await;
                stopped.push(name);
            }
        }
        Ok(stopped)
    }

    #[trace]
    pub async fn close_all(&self) {
        let mut workspaces = self.workspaces.write().await;
        for (_, mut servers) in workspaces.drain() {
            for (_, mut workspace) in servers.drain() {
                workspace.stop_server().await;
            }
        }
    }
}

pub struct WorkspaceHandle<'a> {
    session: &'a Session,
    workspace_root: PathBuf,
    server_name: String,
}

impl<'a> WorkspaceHandle<'a> {
    #[trace]
    pub async fn client(&self) -> Option<Arc<LspClient>> {
        let workspaces = self.session.workspaces.read().await;
        workspaces
            .get(&self.workspace_root)
            .and_then(|servers| servers.get(&self.server_name))
            .and_then(|ws| ws.client())
    }

    pub fn server_name(&self) -> &str {
        &self.server_name
    }

    pub async fn get_startup_stats(&self) -> Option<leta_types::ServerStartupStats> {
        let workspaces = self.session.workspaces.read().await;
        workspaces
            .get(&self.workspace_root)
            .and_then(|servers| servers.get(&self.server_name))
            .and_then(|ws| ws.startup_stats())
    }

    #[trace]
    pub async fn wait_for_ready(&self, timeout_secs: u64) -> bool {
        if let Some(client) = self.client().await {
            client.wait_for_indexing(timeout_secs).await
        } else {
            false
        }
    }

    #[trace]
    pub async fn ensure_document_open(&self, path: &Path) -> Result<(), String> {
        self.cleanup_deleted_documents().await;

        let uri = path_to_uri(path);

        // First check if document needs updating (read lock only)
        let (needs_open, needs_reopen, client) = {
            let workspaces = self.session.workspaces.read().await;
            let workspace = workspaces
                .get(&self.workspace_root)
                .and_then(|servers| servers.get(&self.server_name))
                .ok_or_else(|| "Workspace not found".to_string())?;

            let client = workspace.client();

            if let Some(doc) = workspace.open_documents.get(&uri) {
                let current_content = read_file_content(path).map_err(|e| e.to_string())?;
                if current_content != doc.content {
                    (false, true, client) // needs reopen (close then open)
                } else {
                    (false, false, client) // already open with same content
                }
            } else {
                (true, false, client) // needs open
            }
        };

        if !needs_open && !needs_reopen {
            return Ok(());
        }

        // Close first if needed
        if needs_reopen {
            self.close_document(path).await;
        }

        // Read file content
        let content = read_file_content(path).map_err(|e| e.to_string())?;
        let language_id = get_language_id(path).to_string();

        // Insert document record (write lock, but no LSP call)
        {
            let mut workspaces = self.session.workspaces.write().await;
            let workspace = workspaces
                .get_mut(&self.workspace_root)
                .and_then(|servers| servers.get_mut(&self.server_name))
                .ok_or_else(|| "Workspace not found".to_string())?;

            let doc = OpenDocument {
                _uri: uri.clone(),
                _version: 1,
                content: content.clone(),
                _language_id: language_id.clone(),
            };
            workspace.open_documents.insert(uri.clone(), doc);
        }

        // Send LSP notification OUTSIDE the lock
        if let Some(client) = client {
            let params = serde_json::json!({
                "textDocument": {
                    "uri": uri,
                    "languageId": language_id,
                    "version": 1,
                    "text": content,
                }
            });
            let _ = client
                .send_notification("textDocument/didOpen", params)
                .await;

            // ruby-lsp processes messages asynchronously in a queue
            if client.server_name() == "ruby-lsp" {
                let symbol_params = serde_json::json!({
                    "textDocument": {"uri": uri}
                });
                let _ = client
                    .send_request_raw("textDocument/documentSymbol", symbol_params)
                    .await;
            }
        }

        Ok(())
    }

    #[trace]
    pub async fn close_document(&self, path: &Path) {
        tracing::trace!("WorkspaceHandle::close_document acquiring write lock");
        let mut workspaces = self.session.workspaces.write().await;
        if let Some(servers) = workspaces.get_mut(&self.workspace_root) {
            if let Some(workspace) = servers.get_mut(&self.server_name) {
                workspace.close_document(path).await;
            }
        }
        tracing::trace!("WorkspaceHandle::close_document releasing write lock");
    }

    #[trace]
    pub async fn is_document_open(&self, path: &Path) -> bool {
        let uri = path_to_uri(path);
        let workspaces = self.session.workspaces.read().await;
        if let Some(servers) = workspaces.get(&self.workspace_root) {
            if let Some(workspace) = servers.get(&self.server_name) {
                return workspace.open_documents.contains_key(&uri);
            }
        }
        false
    }

    async fn cleanup_deleted_documents(&self) {
        let mut workspaces = self.session.workspaces.write().await;
        if let Some(servers) = workspaces.get_mut(&self.workspace_root) {
            if let Some(workspace) = servers.get_mut(&self.server_name) {
                workspace.cleanup_deleted_documents().await;
            }
        }
    }

    #[trace]
    pub async fn notify_files_changed(
        &self,
        changes: &[(PathBuf, leta_lsp::lsp_types::FileChangeType)],
    ) {
        let client = match self.client().await {
            Some(c) => c,
            None => return,
        };

        let file_events: Vec<leta_lsp::lsp_types::FileEvent> = changes
            .iter()
            .map(|(path, change_type)| leta_lsp::lsp_types::FileEvent {
                uri: path_to_uri(path).parse().unwrap(),
                typ: *change_type,
            })
            .collect();

        let params = leta_lsp::lsp_types::DidChangeWatchedFilesParams {
            changes: file_events,
        };

        let _ = client
            .send_notification("workspace/didChangeWatchedFiles", params)
            .await;
    }
}