package-json-lsp 0.1.0

use std::collections::HashMap;
use std::sync::Arc;

use serde_json::json;
use tokio::sync::RwLock;
use tower_lsp::jsonrpc::Result;
use tower_lsp::lsp_types::*;
use tower_lsp::{Client, LanguageServer, LspService, Server};

use crate::color::get_catalog_color;
use crate::document::{Document, position_in_range};
use crate::parser::parse_package_dependencies;
use crate::workspace::WorkspaceManager;

pub async fn run_stdio_server() {
    let stdin = tokio::io::stdin();
    let stdout = tokio::io::stdout();

    let documents = Arc::new(RwLock::new(HashMap::new()));
    let documents_for_service = Arc::clone(&documents);
    let (service, socket) = LspService::new(move |client| Backend {
        client,
        documents: documents_for_service,
        workspace: WorkspaceManager::new(Arc::clone(&documents)),
    });

    Server::new(stdin, stdout, socket).serve(service).await;
}

pub struct Backend {
    client: Client,
    documents: Arc<RwLock<HashMap<Url, Document>>>,
    workspace: WorkspaceManager,
}

#[tower_lsp::async_trait]
impl LanguageServer for Backend {
    async fn initialize(&self, params: InitializeParams) -> Result<InitializeResult> {
        if let Some(folders) = params.workspace_folders {
            self.workspace
                .set_workspace_folders(folders.into_iter().map(|folder| folder.uri).collect())
                .await;
        }

        Ok(InitializeResult {
            capabilities: ServerCapabilities {
                text_document_sync: Some(TextDocumentSyncCapability::Kind(
                    TextDocumentSyncKind::INCREMENTAL,
                )),
                code_action_provider: Some(CodeActionProviderCapability::Simple(true)),
                inlay_hint_provider: Some(OneOf::Left(true)),
                code_lens_provider: Some(CodeLensOptions {
                    resolve_provider: Some(false),
                }),
                hover_provider: Some(HoverProviderCapability::Simple(true)),
                definition_provider: Some(OneOf::Left(true)),
                workspace: Some(WorkspaceServerCapabilities {
                    workspace_folders: Some(WorkspaceFoldersServerCapabilities {
                        supported: Some(true),
                        change_notifications: Some(OneOf::Left(true)),
                    }),
                    file_operations: None,
                }),
                ..ServerCapabilities::default()
            },
            server_info: Some(ServerInfo {
                name: "package-json-lsp".to_string(),
                version: Some(env!("CARGO_PKG_VERSION").to_string()),
            }),
        })
    }

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

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

    async fn did_open(&self, params: DidOpenTextDocumentParams) {
        let doc = Document::new(
            params.text_document.uri.clone(),
            params.text_document.version,
            params.text_document.text,
        );
        self.documents
            .write()
            .await
            .insert(params.text_document.uri.clone(), doc);
        self.send_diagnostics(params.text_document.uri).await;
    }

    async fn did_change(&self, params: DidChangeTextDocumentParams) {
        if let Some(doc) = self
            .documents
            .write()
            .await
            .get_mut(&params.text_document.uri)
        {
            doc.version = params.text_document.version;
            doc.apply_changes(params.content_changes);
        }
        self.workspace
            .clear_document_caches(&params.text_document.uri)
            .await;
    }

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

    async fn did_close(&self, params: DidCloseTextDocumentParams) {
        self.documents
            .write()
            .await
            .remove(&params.text_document.uri);
        self.client
            .publish_diagnostics(params.text_document.uri, Vec::new(), None)
            .await;
    }

    async fn inlay_hint(&self, params: InlayHintParams) -> Result<Option<Vec<InlayHint>>> {
        let Some(document) = self.package_document(&params.text_document.uri).await else {
            return Ok(None);
        };

        let deps = parse_package_dependencies(&document);
        let mut hints = Vec::new();
        for dep in deps {
            let Some(catalog) = dep.catalog else {
                continue;
            };
            let Some(result) = self
                .workspace
                .resolve_catalog(&document.uri, &dep.package_name, &catalog)
                .await
            else {
                continue;
            };
            let color_key = if catalog == "default" {
                "default".to_string()
            } else {
                format!("{catalog}-lens")
            };
            hints.push(InlayHint {
                position: dep.value_range.end,
                label: InlayHintLabel::String(result.version),
                kind: Some(InlayHintKind::TYPE),
                text_edits: None,
                tooltip: None,
                padding_left: None,
                padding_right: None,
                data: Some(json!({
                    "catalog": catalog,
                    "color": get_catalog_color(&color_key),
                })),
            });
        }

        Ok(Some(hints))
    }

    async fn hover(&self, params: HoverParams) -> Result<Option<Hover>> {
        let Some(document) = self
            .package_document(&params.text_document_position_params.text_document.uri)
            .await
        else {
            return Ok(None);
        };
        let position = params.text_document_position_params.position;
        let deps = parse_package_dependencies(&document);

        for dep in &deps {
            let Some(catalog) = &dep.catalog else {
                continue;
            };
            if position_in_range(position, dep.value_range)
                && let Some(result) = self
                    .workspace
                    .resolve_catalog(&document.uri, &dep.package_name, catalog)
                    .await
            {
                return Ok(Some(Hover {
                    contents: HoverContents::Markup(MarkupContent {
                        kind: MarkupKind::Markdown,
                        value: format!(
                            "- {} Catalog: `{}`\n- Version: `{}`",
                            result.manager.as_str(),
                            catalog,
                            result.version
                        ),
                    }),
                    range: Some(dep.value_range),
                }));
            }
        }

        for dep in deps {
            if position_in_range(position, dep.property_range) {
                let Some(outdated) = self
                    .workspace
                    .resolve_version(&document.uri, &dep.package_name)
                    .await
                else {
                    return Ok(Some(Hover {
                        contents: HoverContents::Markup(MarkupContent {
                            kind: MarkupKind::Markdown,
                            value: "latest version".to_string(),
                        }),
                        range: Some(dep.property_range),
                    }));
                };

                return Ok(Some(Hover {
                    contents: HoverContents::Markup(MarkupContent {
                        kind: MarkupKind::Markdown,
                        value: format!(
                            "- Wanted: `{}`\n- Latest: `{}`",
                            outdated.wanted, outdated.latest
                        ),
                    }),
                    range: Some(dep.property_range),
                }));
            }
        }

        Ok(None)
    }

    async fn goto_definition(
        &self,
        params: GotoDefinitionParams,
    ) -> Result<Option<GotoDefinitionResponse>> {
        let Some(document) = self
            .package_document(&params.text_document_position_params.text_document.uri)
            .await
        else {
            return Ok(None);
        };
        let position = params.text_document_position_params.position;

        for dep in parse_package_dependencies(&document) {
            if !position_in_range(position, dep.value_range) {
                continue;
            }
            if let Some(catalog) = dep.catalog
                && let Some(result) = self
                    .workspace
                    .resolve_catalog(&document.uri, &dep.package_name, &catalog)
                    .await
                && let Some(definition) = result.definition
            {
                return Ok(Some(GotoDefinitionResponse::Scalar(definition)));
            }
            if dep.is_workspace_ref
                && let Some(result) = self
                    .workspace
                    .resolve_workspace_package(&document.uri, &dep.package_name)
                    .await
            {
                return Ok(Some(GotoDefinitionResponse::Scalar(result.definition)));
            }
        }

        Ok(None)
    }

    async fn code_lens(&self, params: CodeLensParams) -> Result<Option<Vec<CodeLens>>> {
        let Some(document) = self.package_document(&params.text_document.uri).await else {
            return Ok(None);
        };

        let mut lenses = Vec::new();
        for dep in parse_package_dependencies(&document) {
            let Some(outdated) = self
                .workspace
                .resolve_version(&document.uri, &dep.package_name)
                .await
            else {
                continue;
            };

            let version = if let Some(catalog) = &dep.catalog {
                self.workspace
                    .resolve_catalog(&document.uri, &dep.package_name, catalog)
                    .await
                    .map(|result| result.version)
                    .unwrap_or_else(|| dep.version_string.clone())
            } else {
                dep.version_string.clone()
            };
            let prefix = version_prefix(&version);
            let title = if outdated.current == outdated.wanted {
                format!("Latest: {prefix}{}", outdated.latest)
            } else {
                format!(
                    "Wanted: {prefix}{} | Latest: {prefix}{}",
                    outdated.wanted, outdated.latest
                )
            };

            lenses.push(CodeLens {
                range: dep.value_range,
                command: Some(Command {
                    title,
                    command: "package-json-lsp:update".to_string(),
                    arguments: None,
                }),
                data: None,
            });
        }

        Ok(Some(lenses))
    }

    async fn code_action(&self, params: CodeActionParams) -> Result<Option<CodeActionResponse>> {
        let Some(document) = self.package_document(&params.text_document.uri).await else {
            return Ok(None);
        };

        for dep in parse_package_dependencies(&document) {
            if !position_in_range(params.range.start, dep.value_range)
                && !position_in_range(params.range.end, dep.value_range)
            {
                continue;
            }
            let Some(outdated) = self
                .workspace
                .resolve_version(&document.uri, &dep.package_name)
                .await
            else {
                return Ok(None);
            };

            let (uri, range, version) = if let Some(catalog) = &dep.catalog {
                let Some(catalog_info) = self
                    .workspace
                    .resolve_catalog(&document.uri, &dep.package_name, catalog)
                    .await
                else {
                    return Ok(None);
                };
                let Some(definition) = catalog_info.definition else {
                    return Ok(None);
                };
                (definition.uri, definition.range, catalog_info.version)
            } else {
                (
                    document.uri.clone(),
                    dep.value_range,
                    dep.version_string.clone(),
                )
            };

            let prefix = version_prefix(&version);
            let new_text = format!("{prefix}{}", outdated.latest);
            let edit = WorkspaceEdit {
                changes: Some(HashMap::from([(
                    uri,
                    vec![TextEdit {
                        range,
                        new_text: new_text.clone(),
                    }],
                )])),
                document_changes: None,
                change_annotations: None,
            };

            return Ok(Some(vec![CodeActionOrCommand::CodeAction(CodeAction {
                title: format!(
                    "Update {} to latest version ({})",
                    dep.package_name, new_text
                ),
                kind: None,
                diagnostics: Some(
                    params
                        .context
                        .diagnostics
                        .into_iter()
                        .filter(|diagnostic| {
                            diagnostic.source.as_deref() == Some("pnpm")
                                && diagnostic.code.as_ref().is_some_and(|code| match code {
                                    NumberOrString::String(value) => value == "outdated",
                                    NumberOrString::Number(_) => false,
                                })
                        })
                        .collect(),
                ),
                edit: Some(edit),
                command: None,
                is_preferred: None,
                disabled: None,
                data: None,
            })]));
        }

        Ok(None)
    }
}

impl Backend {
    async fn package_document(&self, uri: &Url) -> Option<Document> {
        if !uri.path().ends_with("package.json") {
            return None;
        }
        self.documents.read().await.get(uri).cloned()
    }

    async fn send_diagnostics(&self, uri: Url) {
        let Some(document) = self.package_document(&uri).await else {
            return;
        };
        let mut diagnostics = Vec::new();

        for dep in parse_package_dependencies(&document) {
            let Some(outdated) = self
                .workspace
                .resolve_version(&document.uri, &dep.package_name)
                .await
            else {
                continue;
            };

            diagnostics.push(Diagnostic {
                range: dep.value_range,
                severity: Some(DiagnosticSeverity::INFORMATION),
                code: Some(NumberOrString::String("outdated".to_string())),
                code_description: None,
                source: Some("pnpm".to_string()),
                message: format!(
                    "- Wanted: `{}`\n- Latest: `{}`",
                    outdated.wanted, outdated.latest
                ),
                related_information: None,
                tags: outdated
                    .is_deprecated
                    .then_some(vec![DiagnosticTag::DEPRECATED]),
                data: None,
            });
        }

        self.client
            .publish_diagnostics(uri, diagnostics, Some(document.version))
            .await;
    }
}

fn version_prefix(version: &str) -> &str {
    if version.starts_with(">=") {
        ">="
    } else if version.starts_with(['~', '^', '>']) {
        &version[..1]
    } else {
        ""
    }
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn preserves_version_prefixes() {
        assert_eq!(version_prefix("^1.0.0"), "^");
        assert_eq!(version_prefix("~1.0.0"), "~");
        assert_eq!(version_prefix(">1.0.0"), ">");
        assert_eq!(version_prefix(">=1.0.0"), ">=");
        assert_eq!(version_prefix("1.0.0"), "");
    }
}