catenary-mcp 1.4.0

// SPDX-License-Identifier: GPL-3.0-or-later
// Copyright (C) 2026 Mark Wells <contact@markwells.dev>

use anyhow::{Result, anyhow};
use lsp_types::{
    DidChangeTextDocumentParams, DidCloseTextDocumentParams, DidOpenTextDocumentParams,
    TextDocumentContentChangeEvent, TextDocumentIdentifier, TextDocumentItem, Uri,
    VersionedTextDocumentIdentifier,
};
use std::collections::HashMap;
use std::path::{Path, PathBuf};
use std::time::{Instant, SystemTime};
use tokio::fs;
use tracing::{debug, trace};

/// Tracks the state of an open document.
struct OpenDocument {
    version: i32,
    content: String,
    mtime: SystemTime,
    last_accessed: Instant,
}

/// Manages document lifecycle for the LSP server.
///
/// The LSP protocol requires documents to be explicitly opened before
/// most operations. This manager handles opening documents on first
/// access, tracking their versions, and detecting changes on disk.
pub struct DocumentManager {
    documents: HashMap<PathBuf, OpenDocument>,
}

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

impl DocumentManager {
    /// Creates a new, empty `DocumentManager`.
    #[must_use]
    pub fn new() -> Self {
        Self {
            documents: HashMap::new(),
        }
    }

    /// Ensures a document is open and returns the notification to send if needed.
    ///
    /// If the document is already open but the file has changed on disk,
    /// returns a `didChange` notification instead.
    ///
    /// # Errors
    ///
    /// Returns an error if:
    /// - The path cannot be canonicalized.
    /// - File metadata cannot be read.
    /// - The file cannot be read from disk.
    /// - The path cannot be converted to a valid URI.
    pub async fn ensure_open(&mut self, path: &Path) -> Result<Option<DocumentNotification>> {
        let path = path.canonicalize()?;
        let metadata = fs::metadata(&path).await?;
        let mtime = metadata.modified()?;

        if let Some(doc) = self.documents.get_mut(&path) {
            // Document already open - check if it changed on disk
            if mtime > doc.mtime {
                let content = fs::read_to_string(&path).await?;
                if content != doc.content {
                    doc.version += 1;
                    doc.content.clone_from(&content);
                    doc.mtime = mtime;
                    doc.last_accessed = Instant::now();

                    debug!("Document changed on disk: {}", path.display());

                    return Ok(Some(DocumentNotification::Change(
                        DidChangeTextDocumentParams {
                            text_document: VersionedTextDocumentIdentifier {
                                uri: path_to_uri(&path)?,
                                version: doc.version,
                            },
                            content_changes: vec![TextDocumentContentChangeEvent {
                                range: None,
                                range_length: None,
                                text: content,
                            }],
                        },
                    )));
                }
            }

            doc.last_accessed = Instant::now();
            trace!("Document already open: {}", path.display());
            return Ok(None);
        }

        // Document not open - read and open it
        let content = fs::read_to_string(&path).await?;
        let uri = path_to_uri(&path)?;

        // Detect language ID from extension
        let language_id = detect_language_id(&path);

        let doc = OpenDocument {
            version: 1,
            content: content.clone(),
            mtime,
            last_accessed: Instant::now(),
        };

        self.documents.insert(path.clone(), doc);
        debug!("Opening document: {} ({})", path.display(), language_id);

        Ok(Some(DocumentNotification::Open(
            DidOpenTextDocumentParams {
                text_document: TextDocumentItem {
                    uri,
                    language_id: language_id.to_string(),
                    version: 1,
                    text: content,
                },
            },
        )))
    }

    /// Marks a document as closed and returns the notification to send.
    ///
    /// # Errors
    ///
    /// Returns an error if the path cannot be canonicalized or converted to a URI.
    pub fn close(&mut self, path: &Path) -> Result<Option<DidCloseTextDocumentParams>> {
        let path = path.canonicalize()?;

        if self.documents.remove(&path).is_some() {
            debug!("Closing document: {}", path.display());
            Ok(Some(DidCloseTextDocumentParams {
                text_document: TextDocumentIdentifier {
                    uri: path_to_uri(&path)?,
                },
            }))
        } else {
            Ok(None)
        }
    }

    /// Returns paths of documents that haven't been accessed within the timeout.
    #[must_use]
    pub fn stale_documents(&self, timeout_secs: u64) -> Vec<PathBuf> {
        let now = Instant::now();
        let timeout = std::time::Duration::from_secs(timeout_secs);
        self.documents
            .iter()
            .filter_map(|(path, doc)| {
                if now.duration_since(doc.last_accessed) >= timeout {
                    Some(path.clone())
                } else {
                    None
                }
            })
            .collect()
    }

    /// Returns the URI for an open document.
    ///
    /// # Errors
    ///
    /// Returns an error if the path cannot be canonicalized or converted to a URI.
    pub fn uri_for_path(&self, path: &Path) -> Result<Uri> {
        path_to_uri(&path.canonicalize()?)
    }

    /// Returns the language ID for a given path.
    #[must_use]
    pub fn language_id_for_path(&self, path: &Path) -> &'static str {
        detect_language_id(path)
    }

    /// Checks if there are any open documents for the given language ID.
    #[must_use]
    pub fn has_open_documents(&self, language_id: &str) -> bool {
        self.documents
            .keys()
            .any(|path| detect_language_id(path) == language_id)
    }

    /// Notifies the manager that a file was written externally (by Catenary itself).
    ///
    /// Updates internal state with the new content and returns the appropriate
    /// LSP notification to send, without re-reading from disk.
    ///
    /// # Errors
    ///
    /// Returns an error if the path cannot be canonicalized or converted to a URI.
    pub fn notify_external_write(
        &mut self,
        path: &Path,
        content: &str,
        mtime: SystemTime,
    ) -> Result<DocumentNotification> {
        let path = path.canonicalize()?;
        let uri = path_to_uri(&path)?;

        if let Some(doc) = self.documents.get_mut(&path) {
            // Already open — send didChange
            doc.version += 1;
            doc.content = content.to_string();
            doc.mtime = mtime;
            doc.last_accessed = Instant::now();

            debug!("External write (change): {}", path.display());

            Ok(DocumentNotification::Change(DidChangeTextDocumentParams {
                text_document: VersionedTextDocumentIdentifier {
                    uri,
                    version: doc.version,
                },
                content_changes: vec![TextDocumentContentChangeEvent {
                    range: None,
                    range_length: None,
                    text: content.to_string(),
                }],
            }))
        } else {
            // Not open — send didOpen
            let language_id = detect_language_id(&path);

            let doc = OpenDocument {
                version: 1,
                content: content.to_string(),
                mtime,
                last_accessed: Instant::now(),
            };

            self.documents.insert(path.clone(), doc);
            debug!(
                "External write (open): {} ({})",
                path.display(),
                language_id
            );

            Ok(DocumentNotification::Open(DidOpenTextDocumentParams {
                text_document: TextDocumentItem {
                    uri,
                    language_id: language_id.to_string(),
                    version: 1,
                    text: content.to_string(),
                },
            }))
        }
    }
}

/// Notification to send to the LSP server.
pub enum DocumentNotification {
    /// A `textDocument/didOpen` notification.
    Open(DidOpenTextDocumentParams),
    /// A `textDocument/didChange` notification.
    Change(DidChangeTextDocumentParams),
}

fn path_to_uri(path: &Path) -> Result<Uri> {
    let uri_str = format!("file://{}", path.display());
    uri_str
        .parse()
        .map_err(|e| anyhow!("Invalid path for URI: {}: {}", path.display(), e))
}

fn detect_language_id(path: &Path) -> &'static str {
    if let Some(file_name) = path.file_name().and_then(|n| n.to_str()) {
        match file_name {
            "Dockerfile" => return "dockerfile",
            "Makefile" => return "makefile",
            "CMakeLists.txt" => return "cmake",
            "Cargo.toml" | "Cargo.lock" => return "toml",
            _ => {}
        }
    }

    match path.extension().and_then(|e| e.to_str()) {
        Some("rs") => "rust",
        Some("go") => "go",
        Some("py") => "python",
        Some("js") => "javascript",
        Some("ts") => "typescript",
        Some("tsx") => "typescriptreact",
        Some("jsx") => "javascriptreact",
        Some("c") => "c",
        Some("cpp" | "cc" | "cxx" | "h" | "hpp") => "cpp",
        Some("cs") => "csharp",
        Some("java") => "java",
        Some("kt" | "kts") => "kotlin",
        Some("swift") => "swift",
        Some("rb") => "ruby",
        Some("php") => "php",
        Some("sh" | "bash" | "zsh") => "shellscript",
        Some("json") => "json",
        Some("yaml" | "yml") => "yaml",
        Some("toml") => "toml",
        Some("md") => "markdown",
        Some("html") => "html",
        Some("css") => "css",
        Some("scss") => "scss",
        Some("lua") => "lua",
        Some("sql") => "sql",
        Some("zig") => "zig",
        Some("mojo") => "mojo",
        Some("dart") => "dart",
        Some("m" | "mm") => "objective-c",
        Some("nix") => "nix",
        Some("proto") => "proto",
        Some("graphql" | "gql") => "graphql",
        Some("r" | "R") => "r",
        Some("jl") => "julia",
        Some("scala" | "sc") => "scala",
        Some("hs") => "haskell",
        Some("ex" | "exs") => "elixir",
        Some("erl" | "hrl") => "erlang",
        Some("cmake") => "cmake",
        _ => "plaintext",
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use std::io::Write;
    use tempfile::NamedTempFile;

    #[tokio::test]
    async fn test_open_document() -> Result<()> {
        let mut file = NamedTempFile::with_suffix(".rs")?;
        writeln!(file, "fn main() {{}}")?;

        let mut manager = DocumentManager::new();
        let notification = manager.ensure_open(file.path()).await?;

        assert!(notification.is_some());
        if let Some(DocumentNotification::Open(params)) = notification {
            assert_eq!(params.text_document.language_id, "rust");
            assert_eq!(params.text_document.version, 1);
            assert!(params.text_document.text.contains("fn main()"));
        } else {
            anyhow::bail!("Expected Open notification");
        }
        Ok(())
    }

    #[tokio::test]
    async fn test_already_open_no_change() -> Result<()> {
        let mut file = NamedTempFile::with_suffix(".py")?;
        writeln!(file, "print('hello')")?;

        let mut manager = DocumentManager::new();

        // First open
        let notification1 = manager.ensure_open(file.path()).await?;
        assert!(notification1.is_some());

        // Second access - no notification since file unchanged
        let notification2 = manager.ensure_open(file.path()).await?;
        assert!(notification2.is_none());
        Ok(())
    }

    #[tokio::test]
    async fn test_document_changed_on_disk() -> Result<()> {
        let file = NamedTempFile::with_suffix(".js")?;
        let path = file.path().to_path_buf();
        std::fs::write(&path, "const x = 1;")?;

        let mut manager = DocumentManager::new();

        // First open
        let notification1 = manager.ensure_open(&path).await?;
        assert!(matches!(notification1, Some(DocumentNotification::Open(_))));

        // Modify file (need delay for mtime to differ on some filesystems)
        tokio::time::sleep(tokio::time::Duration::from_millis(50)).await;
        std::fs::write(&path, "const x = 2;")?;

        // Re-access - should get Change notification since content differs
        let notification2 = manager.ensure_open(&path).await?;
        assert!(
            matches!(notification2, Some(DocumentNotification::Change(_))),
            "Expected Change notification after file modification"
        );
        Ok(())
    }

    #[tokio::test]
    async fn test_close_document() -> Result<()> {
        let mut file = NamedTempFile::with_suffix(".go")?;
        writeln!(file, "package main")?;

        let mut manager = DocumentManager::new();
        manager.ensure_open(file.path()).await?;

        let close_params = manager.close(file.path())?;
        assert!(close_params.is_some());

        // Closing again should return None
        let close_params2 = manager.close(file.path())?;
        assert!(close_params2.is_none());
        Ok(())
    }

    #[tokio::test]
    async fn test_stale_documents() -> Result<()> {
        let mut file = NamedTempFile::with_suffix(".txt")?;
        writeln!(file, "test")?;

        let mut manager = DocumentManager::new();
        manager.ensure_open(file.path()).await?;

        // Wait a moment so the document becomes stale
        tokio::time::sleep(tokio::time::Duration::from_millis(100)).await;

        // With 0 second timeout, document should be stale (100ms > 0s)
        let stale = manager.stale_documents(0);
        assert_eq!(stale.len(), 1);

        // With large timeout, nothing should be stale
        let stale = manager.stale_documents(3600);
        assert!(stale.is_empty());
        Ok(())
    }

    #[test]
    fn test_language_detection() {
        assert_eq!(detect_language_id(Path::new("test.rs")), "rust");
        assert_eq!(detect_language_id(Path::new("test.py")), "python");
        assert_eq!(detect_language_id(Path::new("test.js")), "javascript");
        assert_eq!(detect_language_id(Path::new("test.ts")), "typescript");
        assert_eq!(detect_language_id(Path::new("test.tsx")), "typescriptreact");
        assert_eq!(detect_language_id(Path::new("test.go")), "go");
        assert_eq!(detect_language_id(Path::new("test.php")), "php");
        assert_eq!(detect_language_id(Path::new("test.sh")), "shellscript");
        assert_eq!(detect_language_id(Path::new("test.bash")), "shellscript");
        assert_eq!(detect_language_id(Path::new("test.cs")), "csharp");
        assert_eq!(detect_language_id(Path::new("test.kt")), "kotlin");
        assert_eq!(detect_language_id(Path::new("test.swift")), "swift");
        assert_eq!(detect_language_id(Path::new("test.html")), "html");
        assert_eq!(detect_language_id(Path::new("test.css")), "css");
        assert_eq!(detect_language_id(Path::new("test.scss")), "scss");
        assert_eq!(detect_language_id(Path::new("Dockerfile")), "dockerfile");
        assert_eq!(detect_language_id(Path::new("Makefile")), "makefile");
        assert_eq!(detect_language_id(Path::new("CMakeLists.txt")), "cmake");
        assert_eq!(detect_language_id(Path::new("test.zig")), "zig");
        assert_eq!(detect_language_id(Path::new("test.nix")), "nix");
        assert_eq!(detect_language_id(Path::new("test.proto")), "proto");
        assert_eq!(detect_language_id(Path::new("test.graphql")), "graphql");
        assert_eq!(detect_language_id(Path::new("test.r")), "r");
        assert_eq!(detect_language_id(Path::new("test.jl")), "julia");
        assert_eq!(detect_language_id(Path::new("test.ex")), "elixir");
        assert_eq!(detect_language_id(Path::new("Cargo.toml")), "toml");
        assert_eq!(detect_language_id(Path::new("test.unknown")), "plaintext");
        assert_eq!(detect_language_id(Path::new("noextension")), "plaintext");
    }

    #[test]
    fn test_path_to_uri() -> Result<()> {
        let uri = path_to_uri(Path::new("/home/user/test.rs"))?;
        assert!(uri.as_str().starts_with("file:///home/user/test.rs"));
        Ok(())
    }
}