plum

// LSP server for plum. Two diagnostic sources:
//
//   1. Syntax errors — read directly off the tree-sitter parse tree (ERROR /
//      MISSING nodes), which carries real byte/line/column spans.
//
//   2. Type errors — from `plum_checker::checkSource`. `CheckError` carries a
//      `line`/`column` pointing at the enclosing `fn`/`enum`/`test`
//      declaration (not the exact sub-expression/token within it — the
//      checker doesn't track spans that fine yet), so the diagnostic range
//      covers that declaration's line rather than the whole document.
//
// Both are republished on every `did_open`/`did_change`/`did_save`. Type
// checking goes through `plum_core::loadAndMergeWithOverrides`, passing the
// live in-memory buffer for the edited document while every other
// (transitively imported) file still reads from disk — so it reflects
// unsaved edits, not just the last save.
//
// Completion (see `completion.rs`) is a separate, coarser-grained feature:
// keywords + top-level names always; `expr.` member completion and
// not-yet-imported std lib symbols only when the document currently loads
// cleanly enough to build the checker's global tables.

#![allow(non_snake_case)]

use std::collections::HashMap;
use std::path::PathBuf;
use std::sync::OnceLock;

use tokio::sync::Mutex;
use tower_lsp::jsonrpc::Result as RpcResult;
use tower_lsp::lsp_types::*;
use tower_lsp::{Client, LanguageServer, LspService, Server};

mod completion;

struct Backend {
    client: Client,
    docs: Mutex<HashMap<Url, String>>,
    libPath: PathBuf,
    /// Lazily built on first completion request and kept for the server's
    /// lifetime — the std lib doesn't change while `plum-lsp` is running.
    stdIndex: OnceLock<HashMap<String, String>>,
}

impl Backend {
    fn parseTree(source: &str) -> tree_sitter::Tree {
        let mut parser = tree_sitter::Parser::new();
        parser
            .set_language(&tree_sitter_plum::LANGUAGE.into())
            .expect("failed to load plum grammar");
        parser
            .parse(source, None)
            .expect("tree-sitter parse returned None")
    }

    /// Real-span diagnostics straight from the parse tree: any ERROR or
    /// MISSING node becomes one diagnostic at its own range.
    fn syntaxDiagnostics(tree: &tree_sitter::Tree) -> Vec<Diagnostic> {
        let mut out = Vec::new();
        let mut cursor = tree.walk();
        let mut visit = |node: tree_sitter::Node| {
            if node.is_error() || node.is_missing() {
                let start = node.start_position();
                let end = node.end_position();
                out.push(Diagnostic {
                    range: Range {
                        start: Position::new(start.row as u32, start.column as u32),
                        end: Position::new(end.row as u32, end.column as u32),
                    },
                    severity: Some(DiagnosticSeverity::ERROR),
                    source: Some("plum-syntax".into()),
                    message: if node.is_missing() {
                        format!("missing {}", node.kind())
                    } else {
                        "syntax error".into()
                    },
                    ..Default::default()
                });
            }
        };
        // Manual depth-first walk since `TreeCursor` has no built-in visitor.
        loop {
            visit(cursor.node());
            if cursor.goto_first_child() {
                continue;
            }
            loop {
                if cursor.goto_next_sibling() {
                    break;
                }
                if !cursor.goto_parent() {
                    return out;
                }
            }
        }
    }

    /// Diagnostics from `plum_checker::checkSource`, run against the live
    /// in-memory buffer (via `loadAndMergeWithOverrides`) rather than
    /// whatever's on disk, so an unsaved edit is reflected immediately.
    /// Every other (transitively imported) file still resolves from disk.
    /// Returns `None` (rather than an error) if the document has no on-disk
    /// path or fails to load — those aren't type errors, so we don't want to
    /// flash a bogus diagnostic.
    fn typeDiagnostics(&self, uri: &Url, text: &str) -> Option<Vec<Diagnostic>> {
        let path = uri.to_file_path().ok()?;
        let canon = std::fs::canonicalize(&path).unwrap_or(path.clone());
        let mut overrides = HashMap::new();
        overrides.insert(canon, text.to_string());
        let ast = plum_core::loadAndMergeWithOverrides(&path, &self.libPath, &overrides).ok()?;
        let lastLine = text.lines().count().max(1) as u32 - 1;
        match plum_checker::checkSource(&ast) {
            Ok(()) => Some(Vec::new()),
            Err(errors) => Some(
                errors
                    .into_iter()
                    .map(|e| {
                        // `CheckError`'s position is 1-based and points at the
                        // enclosing declaration, not a sub-expression — clamp
                        // to the document and cover the whole line so an
                        // editor still shows something sensible even if the
                        // line itself is short/blank.
                        let line = e.line.saturating_sub(1).min(lastLine as usize) as u32;
                        Diagnostic {
                            range: Range {
                                start: Position::new(line, 0),
                                end: Position::new(line, u32::MAX),
                            },
                            severity: Some(DiagnosticSeverity::ERROR),
                            source: Some("plum-check".into()),
                            message: e.message,
                            ..Default::default()
                        }
                    })
                    .collect(),
            ),
        }
    }

    /// Builds `plum_checker`'s global tables for the live buffer, the same
    /// way `typeDiagnostics` does — used by member/std-lib completion, which
    /// both need method/field/import-target lookups, not just the document's
    /// own AST.
    fn mergedSource(&self, uri: &Url, text: &str) -> Option<plum_core::ast::Source> {
        let path = uri.to_file_path().ok()?;
        let canon = std::fs::canonicalize(&path).unwrap_or(path.clone());
        let mut overrides = HashMap::new();
        overrides.insert(canon, text.to_string());
        plum_core::loadAndMergeWithOverrides(&path, &self.libPath, &overrides).ok()
    }

    async fn completionItems(&self, uri: &Url, pos: Position) -> Vec<CompletionItem> {
        let text = match self.docs.lock().await.get(uri) {
            Some(t) => t.clone(),
            None => return Vec::new(),
        };
        let line_text = text.lines().nth(pos.line as usize).unwrap_or("");
        let sanitized = completion::sanitizeForParsing(&text, pos);
        let tree = Self::parseTree(&sanitized);
        let own_source = plum_core::AstParser::new(&sanitized).parseSource(tree.root_node());

        if let Some(receiver) = completion::receiverBeforeDot(line_text, pos.character as usize) {
            let Some(merged) = self.mergedSource(uri, &sanitized) else {
                return Vec::new();
            };
            let (global_env, methods, enum_variants, enum_params) =
                plum_checker::buildGlobalTables(&merged);
            let min_required = plum_checker::buildMinRequiredArgs(&merged);
            let ctx = plum_checker::CheckCtx {
                methods: &methods,
                enum_variants: &enum_variants,
                enum_params: &enum_params,
                min_required: &min_required,
            };
            let cursor_fn = completion::enclosingFn(&own_source, pos.line as usize + 1);
            return completion::memberItems(cursor_fn, &receiver, &global_env, &ctx);
        }

        let mut items = completion::keywordItems();
        let mut seen: std::collections::HashSet<String> = std::collections::HashSet::new();
        let pushUnique = |item: CompletionItem,
                          items: &mut Vec<CompletionItem>,
                          seen: &mut std::collections::HashSet<String>| {
            if seen.insert(item.label.clone()) {
                items.push(item);
            }
        };

        // The merged source (own file + imports) if it loads cleanly, else
        // just the document's own declarations — still useful mid-edit.
        let base_source = self.mergedSource(uri, &sanitized);
        for item in completion::topLevelItems(base_source.as_ref().unwrap_or(&own_source)) {
            pushUnique(item, &mut items, &mut seen);
        }

        let stdIndex = self
            .stdIndex
            .get_or_init(|| completion::stdLibIndex(&self.libPath));
        for (name, import_path) in stdIndex {
            if seen.contains(name) {
                continue;
            }
            let already_imported = own_source.imports.iter().any(|i| &i.path == import_path);
            if already_imported {
                continue;
            }
            items.push(CompletionItem {
                label: name.clone(),
                kind: Some(CompletionItemKind::CLASS),
                detail: Some(format!("{import_path} (not imported)")),
                additional_text_edits: Some(vec![completion::importInsertionEdit(
                    &text,
                    import_path,
                )]),
                ..Default::default()
            });
        }

        items
    }

    async fn publishAll(&self, uri: Url) {
        let text = match self.docs.lock().await.get(&uri) {
            Some(t) => t.clone(),
            None => return,
        };
        let tree = Self::parseTree(&text);
        let mut diagnostics = Self::syntaxDiagnostics(&tree);
        if diagnostics.is_empty() {
            if let Some(mut typeErrors) = self.typeDiagnostics(&uri, &text) {
                diagnostics.append(&mut typeErrors);
            }
        }
        self.client
            .publish_diagnostics(uri, diagnostics, None)
            .await;
    }
}

#[tower_lsp::async_trait]
impl LanguageServer for Backend {
    async fn initialize(&self, _: InitializeParams) -> RpcResult<InitializeResult> {
        Ok(InitializeResult {
            capabilities: ServerCapabilities {
                text_document_sync: Some(TextDocumentSyncCapability::Kind(
                    TextDocumentSyncKind::FULL,
                )),
                completion_provider: Some(CompletionOptions {
                    trigger_characters: Some(vec![".".to_string()]),
                    ..Default::default()
                }),
                ..Default::default()
            },
            server_info: Some(ServerInfo {
                name: "plum-lsp".into(),
                version: Some(env!("CARGO_PKG_VERSION").into()),
            }),
        })
    }

    async fn initialized(&self, _: InitializedParams) {
        self.client
            .log_message(MessageType::INFO, "plum-lsp initialized")
            .await;
    }

    async fn shutdown(&self) -> RpcResult<()> {
        Ok(())
    }

    async fn did_open(&self, params: DidOpenTextDocumentParams) {
        let uri = params.text_document.uri;
        self.docs
            .lock()
            .await
            .insert(uri.clone(), params.text_document.text);
        self.publishAll(uri).await;
    }

    async fn did_change(&self, mut params: DidChangeTextDocumentParams) {
        let uri = params.text_document.uri;
        // FULL sync: the last content change carries the whole new text.
        if let Some(change) = params.content_changes.pop() {
            self.docs.lock().await.insert(uri.clone(), change.text);
        }
        self.publishAll(uri).await;
    }

    async fn did_save(&self, params: DidSaveTextDocumentParams) {
        self.publishAll(params.text_document.uri).await;
    }

    async fn did_close(&self, params: DidCloseTextDocumentParams) {
        self.docs.lock().await.remove(&params.text_document.uri);
    }

    async fn completion(&self, params: CompletionParams) -> RpcResult<Option<CompletionResponse>> {
        let uri = params.text_document_position.text_document.uri;
        let pos = params.text_document_position.position;
        Ok(Some(CompletionResponse::Array(
            self.completionItems(&uri, pos).await,
        )))
    }
}

#[tokio::main]
async fn main() {
    // `--lib-path` mirrors `plum-cli`'s own required flag (see
    // plum-cli/src/main.rs's `Command::Run`/`Compile`) — the checker has no
    // built-in default location for the std lib. It's the directory
    // CONTAINING `plum-std/` (an `import std/Bool` resolves to
    // `<lib_path>/plum-std/Bool.plum` — see `plum_core::loader::loadImport`),
    // not `plum-std` itself; the repo root is `.` when run from a checkout.
    let libPath = std::env::args()
        .collect::<Vec<_>>()
        .windows(2)
        .find(|w| w[0] == "--lib-path")
        .map(|w| PathBuf::from(&w[1]))
        .unwrap_or_else(|| PathBuf::from("."));

    let stdin = tokio::io::stdin();
    let stdout = tokio::io::stdout();

    let (service, socket) = LspService::new(|client| Backend {
        client,
        docs: Mutex::new(HashMap::new()),
        libPath,
        stdIndex: OnceLock::new(),
    });
    Server::new(stdin, stdout, socket).serve(service).await;
}