kcl_lib/lsp/kcl/

mod.rs

//! Functions for the `kcl` lsp server.
#![allow(dead_code)]

use std::cell::Cell;
use std::collections::HashMap;
use std::sync::Arc;
use std::sync::Mutex;

use anyhow::Result;
#[cfg(feature = "cli")]
use clap::Parser;
use dashmap::DashMap;
use tokio::sync::RwLock;
use tower_lsp::Client;
use tower_lsp::LanguageServer;
use tower_lsp::jsonrpc::Result as RpcResult;
use tower_lsp::lsp_types::CodeAction;
use tower_lsp::lsp_types::CodeActionKind;
use tower_lsp::lsp_types::CodeActionOptions;
use tower_lsp::lsp_types::CodeActionOrCommand;
use tower_lsp::lsp_types::CodeActionParams;
use tower_lsp::lsp_types::CodeActionProviderCapability;
use tower_lsp::lsp_types::CodeActionResponse;
use tower_lsp::lsp_types::ColorInformation;
use tower_lsp::lsp_types::ColorPresentation;
use tower_lsp::lsp_types::ColorPresentationParams;
use tower_lsp::lsp_types::ColorProviderCapability;
use tower_lsp::lsp_types::CompletionItem;
use tower_lsp::lsp_types::CompletionItemKind;
use tower_lsp::lsp_types::CompletionOptions;
use tower_lsp::lsp_types::CompletionParams;
use tower_lsp::lsp_types::CompletionResponse;
use tower_lsp::lsp_types::CreateFilesParams;
use tower_lsp::lsp_types::DeleteFilesParams;
use tower_lsp::lsp_types::Diagnostic;
use tower_lsp::lsp_types::DiagnosticOptions;
use tower_lsp::lsp_types::DiagnosticServerCapabilities;
use tower_lsp::lsp_types::DiagnosticSeverity;
use tower_lsp::lsp_types::DidChangeConfigurationParams;
use tower_lsp::lsp_types::DidChangeTextDocumentParams;
use tower_lsp::lsp_types::DidChangeWatchedFilesParams;
use tower_lsp::lsp_types::DidChangeWorkspaceFoldersParams;
use tower_lsp::lsp_types::DidCloseTextDocumentParams;
use tower_lsp::lsp_types::DidOpenTextDocumentParams;
use tower_lsp::lsp_types::DidSaveTextDocumentParams;
use tower_lsp::lsp_types::DocumentColorParams;
use tower_lsp::lsp_types::DocumentDiagnosticParams;
use tower_lsp::lsp_types::DocumentDiagnosticReport;
use tower_lsp::lsp_types::DocumentDiagnosticReportResult;
use tower_lsp::lsp_types::DocumentFilter;
use tower_lsp::lsp_types::DocumentFormattingParams;
use tower_lsp::lsp_types::DocumentSymbol;
use tower_lsp::lsp_types::DocumentSymbolParams;
use tower_lsp::lsp_types::DocumentSymbolResponse;
use tower_lsp::lsp_types::Documentation;
use tower_lsp::lsp_types::FoldingRange;
use tower_lsp::lsp_types::FoldingRangeParams;
use tower_lsp::lsp_types::FoldingRangeProviderCapability;
use tower_lsp::lsp_types::FullDocumentDiagnosticReport;
use tower_lsp::lsp_types::Hover as LspHover;
use tower_lsp::lsp_types::HoverContents;
use tower_lsp::lsp_types::HoverParams;
use tower_lsp::lsp_types::HoverProviderCapability;
use tower_lsp::lsp_types::InitializeParams;
use tower_lsp::lsp_types::InitializeResult;
use tower_lsp::lsp_types::InitializedParams;
use tower_lsp::lsp_types::InlayHint;
use tower_lsp::lsp_types::InlayHintParams;
use tower_lsp::lsp_types::InsertTextFormat;
use tower_lsp::lsp_types::MarkupContent;
use tower_lsp::lsp_types::MarkupKind;
use tower_lsp::lsp_types::MessageType;
use tower_lsp::lsp_types::OneOf;
use tower_lsp::lsp_types::Position;
use tower_lsp::lsp_types::PrepareRenameResponse;
use tower_lsp::lsp_types::RelatedFullDocumentDiagnosticReport;
use tower_lsp::lsp_types::RenameFilesParams;
use tower_lsp::lsp_types::RenameParams;
use tower_lsp::lsp_types::SemanticToken;
use tower_lsp::lsp_types::SemanticTokenModifier;
use tower_lsp::lsp_types::SemanticTokenType;
use tower_lsp::lsp_types::SemanticTokens;
use tower_lsp::lsp_types::SemanticTokensFullOptions;
use tower_lsp::lsp_types::SemanticTokensLegend;
use tower_lsp::lsp_types::SemanticTokensOptions;
use tower_lsp::lsp_types::SemanticTokensParams;
use tower_lsp::lsp_types::SemanticTokensRegistrationOptions;
use tower_lsp::lsp_types::SemanticTokensResult;
use tower_lsp::lsp_types::SemanticTokensServerCapabilities;
use tower_lsp::lsp_types::ServerCapabilities;
use tower_lsp::lsp_types::SignatureHelp;
use tower_lsp::lsp_types::SignatureHelpOptions;
use tower_lsp::lsp_types::SignatureHelpParams;
use tower_lsp::lsp_types::StaticRegistrationOptions;
use tower_lsp::lsp_types::TextDocumentItem;
use tower_lsp::lsp_types::TextDocumentPositionParams;
use tower_lsp::lsp_types::TextDocumentRegistrationOptions;
use tower_lsp::lsp_types::TextDocumentSyncCapability;
use tower_lsp::lsp_types::TextDocumentSyncKind;
use tower_lsp::lsp_types::TextDocumentSyncOptions;
use tower_lsp::lsp_types::TextEdit;
use tower_lsp::lsp_types::WorkDoneProgressOptions;
use tower_lsp::lsp_types::WorkspaceEdit;
use tower_lsp::lsp_types::WorkspaceFolder;
use tower_lsp::lsp_types::WorkspaceFoldersServerCapabilities;
use tower_lsp::lsp_types::WorkspaceServerCapabilities;

use crate::ModuleId;
use crate::Program;
use crate::SourceRange;
use crate::docs::kcl_doc::ArgData;
use crate::docs::kcl_doc::ModData;
use crate::exec::KclValue;
use crate::execution::cache;
use crate::lsp::LspSuggestion;
use crate::lsp::ToLspRange;
use crate::lsp::backend::Backend as _;
use crate::lsp::kcl::hover::Hover;
use crate::lsp::kcl::hover::HoverOpts;
use crate::lsp::util::IntoDiagnostic;
use crate::parsing::PIPE_OPERATOR;
use crate::parsing::ast::types::Expr;
use crate::parsing::ast::types::VariableKind;
use crate::parsing::token::LexerMode;
use crate::parsing::token::RESERVED_WORDS;
use crate::parsing::token::TokenStream;
use crate::parsing::token::adapter;
use crate::parsing::token::lex;

pub mod custom_notifications;
mod hover;

const SEMANTIC_TOKEN_TYPES: [SemanticTokenType; 10] = [
    SemanticTokenType::NUMBER,
    SemanticTokenType::VARIABLE,
    SemanticTokenType::KEYWORD,
    SemanticTokenType::TYPE,
    SemanticTokenType::STRING,
    SemanticTokenType::OPERATOR,
    SemanticTokenType::COMMENT,
    SemanticTokenType::FUNCTION,
    SemanticTokenType::PARAMETER,
    SemanticTokenType::PROPERTY,
];

const SEMANTIC_TOKEN_MODIFIERS: [SemanticTokenModifier; 5] = [
    SemanticTokenModifier::DECLARATION,
    SemanticTokenModifier::DEFINITION,
    SemanticTokenModifier::DEFAULT_LIBRARY,
    SemanticTokenModifier::READONLY,
    SemanticTokenModifier::STATIC,
];

/// A subcommand for running the server.
#[derive(Clone, Debug)]
#[cfg_attr(feature = "cli", derive(Parser))]
pub struct Server {
    /// Port that the server should listen
    #[cfg_attr(feature = "cli", clap(long, default_value = "8080"))]
    pub socket: i32,

    /// Listen over stdin and stdout instead of a tcp socket.
    #[cfg_attr(feature = "cli", clap(short, long, default_value = "false"))]
    pub stdio: bool,
}

/// The lsp server backend.
#[derive(Clone)]
pub struct Backend {
    /// The client for the backend.
    pub client: Client,
    /// The file system client to use.
    pub fs: crate::fs::FileSystemHandle,
    /// The workspace folders.
    pub workspace_folders: DashMap<String, WorkspaceFolder>,
    /// The stdlib completions for the language.
    pub stdlib_completions: HashMap<String, CompletionItem>,
    /// The stdlib completions inside a sketch block, where `solver::*`
    /// shadows the regular sketch stdlib.
    pub sketch_block_stdlib_completions: HashMap<String, CompletionItem>,
    /// The stdlib signatures for the language.
    pub stdlib_signatures: HashMap<String, SignatureHelp>,
    /// The stdlib signatures inside a sketch block.
    pub sketch_block_stdlib_signatures: HashMap<String, SignatureHelp>,
    /// For all KwArg functions in std, a map from their arg names to arg help snippets (markdown format).
    pub stdlib_args: HashMap<String, HashMap<String, LspArgData>>,
    /// KwArg docs inside a sketch block.
    pub sketch_block_stdlib_args: HashMap<String, HashMap<String, LspArgData>>,
    /// KCL keywords
    pub kcl_keywords: HashMap<String, CompletionItem>,
    /// Token maps.
    pub(super) token_map: DashMap<String, TokenStream>,
    /// AST maps.
    pub ast_map: DashMap<String, crate::Program>,
    /// Current code.
    pub code_map: DashMap<String, Vec<u8>>,
    /// Diagnostics.
    pub diagnostics_map: DashMap<String, Vec<Diagnostic>>,
    /// Symbols map.
    pub symbols_map: DashMap<String, Vec<DocumentSymbol>>,
    /// Semantic tokens map.
    pub semantic_tokens_map: DashMap<String, Vec<SemanticToken>>,
    /// The Zoo API client.
    pub zoo_client: kittycad::Client,
    /// Optional executor context to use if we want to execute the code.
    pub executor_ctx: Arc<RwLock<Option<crate::execution::ExecutorContext>>>,
    /// If we are currently allowed to execute the ast.
    pub can_execute: Arc<RwLock<bool>>,

    pub is_initialized: Arc<RwLock<bool>>,
}

impl Backend {
    #[cfg(target_arch = "wasm32")]
    pub fn new_wasm(
        client: Client,
        executor_ctx: Option<crate::execution::ExecutorContext>,
        fs: crate::fs::wasm::FileSystemManager,
        zoo_client: kittycad::Client,
    ) -> Result<Self, String> {
        Self::with_file_system(
            client,
            executor_ctx,
            crate::fs::new_file_system_handle(crate::fs::FileManager::new(fs)),
            zoo_client,
        )
    }

    #[cfg(not(target_arch = "wasm32"))]
    pub fn new(
        client: Client,
        executor_ctx: Option<crate::execution::ExecutorContext>,
        zoo_client: kittycad::Client,
    ) -> Result<Self, String> {
        Self::with_file_system(
            client,
            executor_ctx,
            crate::fs::new_file_system_handle(crate::fs::FileManager::new()),
            zoo_client,
        )
    }

    fn with_file_system(
        client: Client,
        executor_ctx: Option<crate::execution::ExecutorContext>,
        fs: crate::fs::FileSystemHandle,
        zoo_client: kittycad::Client,
    ) -> Result<Self, String> {
        let kcl_std = crate::docs::kcl_doc::walk_stdlib();
        let stdlib_completions = get_completions_from_stdlib(&kcl_std).map_err(|e| e.to_string())?;
        let sketch_block_stdlib_completions =
            get_completions_from_stdlib_for_sketch_block(&kcl_std).map_err(|e| e.to_string())?;
        let stdlib_signatures = get_signatures_from_stdlib(&kcl_std);
        let sketch_block_stdlib_signatures = get_signatures_from_stdlib_for_sketch_block(&kcl_std);
        let stdlib_args = get_arg_maps_from_stdlib(&kcl_std);
        let sketch_block_stdlib_args = get_arg_maps_from_stdlib_for_sketch_block(&kcl_std);
        let kcl_keywords = get_keywords();

        Ok(Self {
            client,
            fs,
            stdlib_completions,
            sketch_block_stdlib_completions,
            stdlib_signatures,
            sketch_block_stdlib_signatures,
            stdlib_args,
            sketch_block_stdlib_args,
            kcl_keywords,
            zoo_client,
            can_execute: Arc::new(RwLock::new(executor_ctx.is_some())),
            executor_ctx: Arc::new(RwLock::new(executor_ctx)),
            workspace_folders: Default::default(),
            token_map: Default::default(),
            ast_map: Default::default(),
            code_map: Default::default(),
            diagnostics_map: Default::default(),
            symbols_map: Default::default(),
            semantic_tokens_map: Default::default(),
            is_initialized: Default::default(),
        })
    }

    fn is_in_sketch_block(ast: &crate::Program, position: usize) -> bool {
        let in_sketch_block = Cell::new(false);
        let _ = crate::walk::walk(&ast.ast, |node| {
            if let crate::walk::Node::SketchBlock(sketch_block) = node
                && SourceRange::from(&sketch_block.body).contains(position)
            {
                in_sketch_block.set(true);
                return Ok::<bool, anyhow::Error>(false);
            }

            Ok::<bool, anyhow::Error>(true)
        });

        in_sketch_block.get()
    }

    fn stdlib_completions_for_position<'a>(
        &'a self,
        ast: &crate::Program,
        position: usize,
    ) -> &'a HashMap<String, CompletionItem> {
        if Self::is_in_sketch_block(ast, position) {
            &self.sketch_block_stdlib_completions
        } else {
            &self.stdlib_completions
        }
    }

    fn stdlib_signatures_for_position<'a>(
        &'a self,
        ast: &crate::Program,
        position: usize,
    ) -> &'a HashMap<String, SignatureHelp> {
        if Self::is_in_sketch_block(ast, position) {
            &self.sketch_block_stdlib_signatures
        } else {
            &self.stdlib_signatures
        }
    }

    fn stdlib_args_for_position<'a>(
        &'a self,
        ast: &crate::Program,
        position: usize,
    ) -> &'a HashMap<String, HashMap<String, LspArgData>> {
        if Self::is_in_sketch_block(ast, position) {
            &self.sketch_block_stdlib_args
        } else {
            &self.stdlib_args
        }
    }

    fn remove_from_ast_maps(&self, filename: &str) {
        self.ast_map.remove(filename);
        self.symbols_map.remove(filename);
    }

    fn try_arg_completions(
        &self,
        program: &crate::Program,
        position: usize,
        current_code: &str,
    ) -> Option<impl Iterator<Item = CompletionItem>> {
        let curr_expr = program.ast.get_expr_for_position(position)?;
        let hover =
            curr_expr.get_hover_value_for_position(position, current_code, &HoverOpts::default_for_signature_help())?;

        // Now we can tell if the user's cursor is inside a callable function.
        // If so, get its name (the function name being called.)
        let maybe_callee = match hover {
            Hover::Function { name, range: _ } => Some(name),
            Hover::Signature {
                name,
                parameter_index: _,
                range: _,
            } => Some(name),
            Hover::Comment { .. } => None,
            Hover::Variable { .. } => None,
            Hover::KwArg {
                callee_name,
                name: _,
                range: _,
            } => Some(callee_name),
            Hover::Type { .. } => None,
        };
        let stdlib_args = self.stdlib_args_for_position(program, position);
        let callee_args = maybe_callee.and_then(|fn_name| stdlib_args.get(&fn_name))?;

        let arg_label_completions = callee_args
            .iter()
            // Don't suggest labels for unlabelled args!
            .filter(|(_arg_name, arg_data)| arg_data.props.is_labelled())
            .map(|(arg_name, arg_data)| CompletionItem {
                label: arg_name.to_owned(),
                label_details: None,
                kind: Some(CompletionItemKind::PROPERTY),
                detail: arg_data.props.ty.clone(),
                documentation: arg_data.props.docs.clone().map(|docs| {
                    Documentation::MarkupContent(MarkupContent {
                        kind: MarkupKind::Markdown,
                        value: docs,
                    })
                }),
                deprecated: None,
                preselect: None,
                sort_text: Some(arg_name.to_owned()),
                filter_text: Some(arg_name.to_owned()),
                insert_text: {
                    // let snippet = "${0:x}";
                    if let Some(snippet) = arg_data.props.get_autocomplete_snippet(0).map(|(_i, snippet)| snippet) {
                        Some(snippet)
                    } else {
                        Some(format!("{arg_name} = "))
                    }
                },
                insert_text_format: Some(InsertTextFormat::SNIPPET),
                insert_text_mode: None,
                text_edit: None,
                additional_text_edits: None,
                command: None,
                commit_characters: None,
                data: None,
                tags: None,
            });
        Some(arg_label_completions)
    }
}

// Implement the shared backend trait for the language server.
#[async_trait::async_trait]
impl crate::lsp::backend::Backend for Backend {
    fn client(&self) -> &Client {
        &self.client
    }

    fn fs(&self) -> &crate::fs::FileSystemHandle {
        &self.fs
    }

    async fn is_initialized(&self) -> bool {
        *self.is_initialized.read().await
    }

    async fn set_is_initialized(&self, is_initialized: bool) {
        *self.is_initialized.write().await = is_initialized;
    }

    async fn workspace_folders(&self) -> Vec<WorkspaceFolder> {
        // TODO: fix clone
        self.workspace_folders.iter().map(|i| i.clone()).collect()
    }

    async fn add_workspace_folders(&self, folders: Vec<WorkspaceFolder>) {
        for folder in folders {
            self.workspace_folders.insert(folder.name.to_string(), folder);
        }
    }

    async fn remove_workspace_folders(&self, folders: Vec<WorkspaceFolder>) {
        for folder in folders {
            self.workspace_folders.remove(&folder.name);
        }
    }

    fn code_map(&self) -> &DashMap<String, Vec<u8>> {
        &self.code_map
    }

    async fn insert_code_map(&self, uri: String, text: Vec<u8>) {
        self.code_map.insert(uri, text);
    }

    async fn remove_from_code_map(&self, uri: String) -> Option<Vec<u8>> {
        self.code_map.remove(&uri).map(|x| x.1)
    }

    async fn clear_code_state(&self) {
        self.code_map.clear();
        self.token_map.clear();
        self.ast_map.clear();
        self.diagnostics_map.clear();
        self.symbols_map.clear();
        self.semantic_tokens_map.clear();
    }

    fn current_diagnostics_map(&self) -> &DashMap<String, Vec<Diagnostic>> {
        &self.diagnostics_map
    }

    async fn inner_on_change(&self, params: TextDocumentItem, force: bool) {
        if force {
            crate::bust_cache().await;
        }

        let filename = params.uri.to_string();
        // We already updated the code map in the shared backend.

        // Lets update the tokens.
        let module_id = ModuleId::default();
        // Mode-aware lex. The old lexer keeps its exact bail-on-error behavior.
        // The new lexer always yields a token stream so semantic highlighting
        // survives a lexical error; any lexical error is carried in `lex_error`
        // and reported below, after semantic tokens are computed.
        let (tokens, lex_error) = match LexerMode::resolve() {
            LexerMode::Old => match lex(&params.text, module_id) {
                Ok(tokens) => (tokens, None),
                Err(err) => {
                    self.add_to_diagnostics(&params, &[err], true).await;
                    self.token_map.remove(&filename);
                    self.remove_from_ast_maps(&filename);
                    self.semantic_tokens_map.remove(&filename);
                    return;
                }
            },
            LexerMode::New => {
                let result = adapter::lex_with_diagnostics(&params.text, module_id);
                let lex_error = result.to_lexical_error();
                (result.tokens, lex_error)
            }
        };

        // Get the previous tokens.
        let tokens_changed = match self.token_map.get(&filename) {
            Some(previous_tokens) => *previous_tokens != tokens,
            _ => true,
        };

        let had_diagnostics = self.has_diagnostics(params.uri.as_ref()).await;

        // Check if the tokens are the same.
        if !tokens_changed && !force && !had_diagnostics {
            // We return early here because the tokens are the same.
            return;
        }

        if tokens_changed {
            // Update our token map.
            self.token_map.insert(params.uri.to_string(), tokens.clone());
            // Update our semantic tokens.
            self.update_semantic_tokens(&tokens, &params).await;
        }

        // With the new lexer a lexical error is surfaced as a diagnostic, but the
        // token stream and semantic tokens (computed above) are retained so the
        // editor keeps highlighting. No AST is produced (mirrors the parse-error
        // path below).
        if let Some(err) = lex_error {
            self.add_to_diagnostics(&params, &[err], true).await;
            self.remove_from_ast_maps(&filename);
            return;
        }

        // Lets update the ast.

        let (ast, errs) = match crate::parsing::parse_tokens(tokens.clone()).0 {
            Ok(result) => result,
            Err(err) => {
                self.add_to_diagnostics(&params, &[err], true).await;
                self.remove_from_ast_maps(&filename);
                return;
            }
        };

        self.add_to_diagnostics(&params, &errs, true).await;

        if errs.iter().any(|e| e.severity == crate::errors::Severity::Fatal) {
            self.remove_from_ast_maps(&filename);
            return;
        }

        let Some(mut ast) = ast else {
            self.remove_from_ast_maps(&filename);
            return;
        };

        // Here we will want to store the digest and compare, but for now
        // we're doing this in a non-load-bearing capacity so we can remove
        // this if it backfires and only hork the LSP.
        ast.compute_digest();

        // Save it as a program.
        let ast = crate::Program {
            ast,
            original_file_contents: params.text.clone(),
        };

        // Check if the ast changed.
        let ast_changed = match self.ast_map.get(&filename) {
            Some(old_ast) => {
                // Check if the ast changed.
                *old_ast.ast != *ast.ast
            }
            None => true,
        };

        if !ast_changed && !force && !had_diagnostics {
            // Return early if the ast did not change and we don't need to force.
            return;
        }

        if ast_changed {
            self.ast_map.insert(params.uri.to_string(), ast.clone());
            // Update the symbols map.
            self.symbols_map.insert(
                params.uri.to_string(),
                ast.ast.get_lsp_symbols(&params.text).unwrap_or_default(),
            );

            // Update our semantic tokens.
            self.update_semantic_tokens(&tokens, &params).await;

            let mut discovered_findings: Vec<_> = ast.lint_all().into_iter().flatten().collect();
            // Filter out Z0005 (old sketch syntax) from LSP diagnostics
            // TODO: Remove this filter once the transpiler is complete and all tests are updated
            discovered_findings.retain(|finding| finding.finding.code != "Z0005");
            self.add_to_diagnostics(&params, &discovered_findings, false).await;
        }

        // Send the notification to the client that the ast was updated.
        if self.can_execute().await || self.executor_ctx().await.is_none() {
            // Only send the notification if we can execute.
            // Otherwise it confuses the client.
            self.client
                .send_notification::<custom_notifications::AstUpdated>(ast.ast.clone())
                .await;
        }

        // Execute the code if we have an executor context.
        // This function automatically executes if we should & updates the diagnostics if we got
        // errors.
        if self.execute(&params, &ast).await.is_err() {
            return;
        }

        // If we made it here we can clear the diagnostics.
        self.clear_diagnostics_map(&params.uri, Some(DiagnosticSeverity::ERROR))
            .await;
    }
}

impl Backend {
    pub async fn can_execute(&self) -> bool {
        *self.can_execute.read().await
    }

    pub async fn executor_ctx(&self) -> tokio::sync::RwLockReadGuard<'_, Option<crate::execution::ExecutorContext>> {
        self.executor_ctx.read().await
    }

    async fn update_semantic_tokens(&self, tokens: &TokenStream, params: &TextDocumentItem) {
        // Update the semantic tokens map.
        let mut semantic_tokens = vec![];
        let mut last_position = Position::new(0, 0);
        for token in tokens.as_slice() {
            let Ok(token_type) = SemanticTokenType::try_from(token.token_type) else {
                // We continue here because not all tokens can be converted this way, we will get
                // the rest from the ast.
                continue;
            };

            let mut token_type_index = match self.get_semantic_token_type_index(&token_type) {
                Some(index) => index,
                // This is actually bad this should not fail.
                // The test for listing all semantic token types should make this never happen.
                None => {
                    self.client
                        .log_message(
                            MessageType::ERROR,
                            format!("token type `{token_type:?}` not accounted for"),
                        )
                        .await;
                    continue;
                }
            };

            let source_range: SourceRange = token.into();
            let position = source_range.start_to_lsp_position(&params.text);

            // Calculate the token modifiers.
            // Get the value at the current position.
            let token_modifiers_bitset = match self.ast_map.get(params.uri.as_str()) {
                Some(ast) => {
                    let token_index = Arc::new(Mutex::new(token_type_index));
                    let modifier_index: Arc<Mutex<u32>> = Arc::new(Mutex::new(0));
                    crate::walk::walk(&ast.ast, |node: crate::walk::Node| {
                        let Ok(node_range): Result<SourceRange, _> = (&node).try_into() else {
                            return Ok(true);
                        };

                        if !node_range.contains(source_range.start()) {
                            return Ok(true);
                        }

                        let get_modifier = |modifier: Vec<SemanticTokenModifier>| -> Result<bool> {
                            let mut mods = modifier_index.lock().map_err(|_| anyhow::anyhow!("mutex"))?;
                            let Some(token_modifier_index) = self.get_semantic_token_modifier_index(modifier) else {
                                return Ok(true);
                            };
                            if *mods == 0 {
                                *mods = token_modifier_index;
                            } else {
                                *mods |= token_modifier_index;
                            }
                            Ok(false)
                        };

                        match node {
                            crate::walk::Node::TagDeclarator(_) => {
                                return get_modifier(vec![
                                    SemanticTokenModifier::DEFINITION,
                                    SemanticTokenModifier::STATIC,
                                ]);
                            }
                            crate::walk::Node::VariableDeclarator(variable) => {
                                let sr: SourceRange = (&variable.id).into();
                                if sr.contains(source_range.start()) {
                                    if let Expr::FunctionExpression(_) = &variable.init {
                                        let mut ti = token_index.lock().map_err(|_| anyhow::anyhow!("mutex"))?;
                                        *ti = match self.get_semantic_token_type_index(&SemanticTokenType::FUNCTION) {
                                            Some(index) => index,
                                            None => token_type_index,
                                        };
                                    }

                                    return get_modifier(vec![
                                        SemanticTokenModifier::DECLARATION,
                                        SemanticTokenModifier::READONLY,
                                    ]);
                                }
                            }
                            crate::walk::Node::Parameter(_) => {
                                let mut ti = token_index.lock().map_err(|_| anyhow::anyhow!("mutex"))?;
                                *ti = match self.get_semantic_token_type_index(&SemanticTokenType::PARAMETER) {
                                    Some(index) => index,
                                    None => token_type_index,
                                };
                                return Ok(false);
                            }
                            crate::walk::Node::MemberExpression(member_expression) => {
                                let sr: SourceRange = (&member_expression.property).into();
                                if sr.contains(source_range.start()) {
                                    let mut ti = token_index.lock().map_err(|_| anyhow::anyhow!("mutex"))?;
                                    *ti = match self.get_semantic_token_type_index(&SemanticTokenType::PROPERTY) {
                                        Some(index) => index,
                                        None => token_type_index,
                                    };
                                    return Ok(false);
                                }
                            }
                            crate::walk::Node::ObjectProperty(object_property) => {
                                let sr: SourceRange = (&object_property.key).into();
                                if sr.contains(source_range.start()) {
                                    let mut ti = token_index.lock().map_err(|_| anyhow::anyhow!("mutex"))?;
                                    *ti = match self.get_semantic_token_type_index(&SemanticTokenType::PROPERTY) {
                                        Some(index) => index,
                                        None => token_type_index,
                                    };
                                }
                                return get_modifier(vec![SemanticTokenModifier::DECLARATION]);
                            }
                            crate::walk::Node::CallExpressionKw(call_expr) => {
                                let sr: SourceRange = (&call_expr.callee).into();
                                if sr.contains(source_range.start()) {
                                    let mut ti = token_index.lock().map_err(|_| anyhow::anyhow!("mutex"))?;
                                    *ti = match self.get_semantic_token_type_index(&SemanticTokenType::FUNCTION) {
                                        Some(index) => index,
                                        None => token_type_index,
                                    };

                                    if self.stdlib_completions.contains_key(&call_expr.callee.name.name) {
                                        // This is a stdlib function.
                                        return get_modifier(vec![SemanticTokenModifier::DEFAULT_LIBRARY]);
                                    }

                                    return Ok(false);
                                }
                            }
                            _ => {}
                        }
                        Ok(true)
                    })
                    .unwrap_or_default();

                    let t = match token_index.lock() {
                        Ok(guard) => *guard,
                        _ => 0,
                    };
                    token_type_index = t;

                    match modifier_index.lock() {
                        Ok(guard) => *guard,
                        _ => 0,
                    }
                }
                _ => 0,
            };

            // We need to check if we are on the last token of the line.
            // If we are starting from the end of the last line just add 1 to the line.
            // Check if we are on the last token of the line.
            if let Some(line) = params.text.lines().nth(position.line as usize)
                && line.len() == position.character as usize
            {
                // We are on the last token of the line.
                // We need to add a new line.
                let semantic_token = SemanticToken {
                    delta_line: position.line - last_position.line + 1,
                    delta_start: 0,
                    length: (token.end - token.start) as u32,
                    token_type: token_type_index,
                    token_modifiers_bitset,
                };

                semantic_tokens.push(semantic_token);

                last_position = Position::new(position.line + 1, 0);
                continue;
            }

            let semantic_token = SemanticToken {
                delta_line: position.line - last_position.line,
                delta_start: if position.line != last_position.line {
                    position.character
                } else {
                    position.character - last_position.character
                },
                length: (token.end - token.start) as u32,
                token_type: token_type_index,
                token_modifiers_bitset,
            };

            semantic_tokens.push(semantic_token);

            last_position = position;
        }
        self.semantic_tokens_map.insert(params.uri.to_string(), semantic_tokens);
    }

    async fn clear_diagnostics_map(&self, uri: &url::Url, severity: Option<DiagnosticSeverity>) {
        let Some(mut items) = self.diagnostics_map.get_mut(uri.as_str()) else {
            return;
        };

        // If we only want to clear a specific severity, do that.
        if let Some(severity) = severity {
            items.retain(|x| x.severity != Some(severity));
        } else {
            items.clear();
        }

        if items.is_empty() {
            #[cfg(not(target_arch = "wasm32"))]
            {
                self.client.publish_diagnostics(uri.clone(), items.clone(), None).await;
            }

            // We need to drop the items here.
            drop(items);

            self.diagnostics_map.remove(uri.as_str());
        } else {
            // We don't need to update the map since we used get_mut.

            #[cfg(not(target_arch = "wasm32"))]
            {
                self.client.publish_diagnostics(uri.clone(), items.clone(), None).await;
            }
        }
    }

    async fn add_to_diagnostics<DiagT: IntoDiagnostic + std::fmt::Debug>(
        &self,
        params: &TextDocumentItem,
        diagnostics: &[DiagT],
        clear_all_before_add: bool,
    ) {
        if diagnostics.is_empty() {
            return;
        }

        if clear_all_before_add {
            self.clear_diagnostics_map(&params.uri, None).await;
        } else if diagnostics.iter().all(|x| x.severity() == DiagnosticSeverity::ERROR) {
            // If the diagnostic is an error, it will be the only error we get since that halts
            // execution.
            // Clear the diagnostics before we add a new one.
            self.clear_diagnostics_map(&params.uri, Some(DiagnosticSeverity::ERROR))
                .await;
        } else if diagnostics
            .iter()
            .all(|x| x.severity() == DiagnosticSeverity::INFORMATION)
        {
            // If the diagnostic is a lint, we will pass them all to add at once so we need to
            // clear the old ones.
            self.clear_diagnostics_map(&params.uri, Some(DiagnosticSeverity::INFORMATION))
                .await;
        }

        let mut items = match self.diagnostics_map.get(params.uri.as_str()) {
            Some(items) => {
                // TODO: Would be awesome to fix the clone here.
                items.clone()
            }
            _ => {
                vec![]
            }
        };

        for diagnostic in diagnostics {
            let lsp_d = diagnostic.to_lsp_diagnostics(&params.text);
            // Make sure we don't duplicate diagnostics.
            for d in lsp_d {
                if !items.iter().any(|x| x == &d) {
                    items.push(d);
                }
            }
        }

        self.diagnostics_map.insert(params.uri.to_string(), items.clone());

        self.client.publish_diagnostics(params.uri.clone(), items, None).await;
    }

    async fn execute(&self, params: &TextDocumentItem, ast: &Program) -> Result<()> {
        // Check if we can execute.
        if !self.can_execute().await {
            return Ok(());
        }

        // Execute the code if we have an executor context.
        let ctx = self.executor_ctx().await;
        let Some(ref executor_ctx) = *ctx else {
            return Ok(());
        };

        if !self.is_initialized().await {
            // We are not initialized yet.
            return Ok(());
        }

        // Use run_mock for mock contexts, run_with_caching for live contexts
        let result = if executor_ctx.is_mock() {
            executor_ctx
                .run_mock(ast, &crate::execution::MockConfig::default())
                .await
        } else {
            executor_ctx.run_with_caching(ast.clone()).await
        };

        match result {
            Err(err) => {
                self.add_to_diagnostics(params, &[err], false).await;

                // Since we already published the diagnostics we don't really care about the error
                // string.
                Err(anyhow::anyhow!("failed to execute code"))
            }
            Ok(_) => Ok(()),
        }
    }

    pub fn get_semantic_token_type_index(&self, token_type: &SemanticTokenType) -> Option<u32> {
        SEMANTIC_TOKEN_TYPES
            .iter()
            .position(|x| *x == *token_type)
            .map(|y| y as u32)
    }

    pub fn get_semantic_token_modifier_index(&self, token_types: Vec<SemanticTokenModifier>) -> Option<u32> {
        if token_types.is_empty() {
            return None;
        }

        let mut modifier = None;
        for token_type in token_types {
            if let Some(index) = SEMANTIC_TOKEN_MODIFIERS
                .iter()
                .position(|x| *x == token_type)
                .map(|y| y as u32)
            {
                modifier = match modifier {
                    Some(modifier) => Some(modifier | index),
                    None => Some(index),
                };
            }
        }
        modifier
    }

    pub async fn update_can_execute(
        &self,
        params: custom_notifications::UpdateCanExecuteParams,
    ) -> RpcResult<custom_notifications::UpdateCanExecuteResponse> {
        let mut can_execute = self.can_execute.write().await;

        if *can_execute == params.can_execute {
            return Ok(custom_notifications::UpdateCanExecuteResponse {});
        }

        *can_execute = params.can_execute;

        Ok(custom_notifications::UpdateCanExecuteResponse {})
    }

    /// Returns the new string for the code after rename.
    pub fn inner_prepare_rename(
        &self,
        params: &TextDocumentPositionParams,
        new_name: &str,
    ) -> RpcResult<Option<(String, String)>> {
        let filename = params.text_document.uri.to_string();

        let Some(current_code) = self.code_map.get(&filename) else {
            return Ok(None);
        };
        let Ok(current_code) = std::str::from_utf8(&current_code) else {
            return Ok(None);
        };

        // Parse the ast.
        // I don't know if we need to do this again since it should be updated in the context.
        // But I figure better safe than sorry since this will write back out to the file.
        let module_id = ModuleId::default();
        let Ok(mut ast) = crate::parsing::parse_str(current_code, module_id).parse_errs_as_err() else {
            return Ok(None);
        };

        // Let's convert the position to a character index.
        let pos = position_to_char_index(params.position, current_code);
        // Now let's perform the rename on the ast.
        ast.rename_symbol(new_name, pos);
        // Now recast it.
        let recast = ast.recast_top(&Default::default(), 0);

        Ok(Some((current_code.to_string(), recast)))
    }
}

#[tower_lsp::async_trait]
impl LanguageServer for Backend {
    async fn initialize(&self, params: InitializeParams) -> RpcResult<InitializeResult> {
        self.client
            .log_message(MessageType::INFO, format!("initialize: {params:?}"))
            .await;

        Ok(InitializeResult {
            capabilities: ServerCapabilities {
                color_provider: Some(ColorProviderCapability::Simple(true)),
                code_action_provider: Some(CodeActionProviderCapability::Options(CodeActionOptions {
                    code_action_kinds: Some(vec![CodeActionKind::QUICKFIX]),
                    resolve_provider: Some(false),
                    work_done_progress_options: WorkDoneProgressOptions::default(),
                })),
                completion_provider: Some(CompletionOptions {
                    resolve_provider: Some(false),
                    trigger_characters: Some(vec![".".to_string()]),
                    work_done_progress_options: Default::default(),
                    all_commit_characters: None,
                    ..Default::default()
                }),
                diagnostic_provider: Some(DiagnosticServerCapabilities::Options(DiagnosticOptions {
                    ..Default::default()
                })),
                document_formatting_provider: Some(OneOf::Left(true)),
                folding_range_provider: Some(FoldingRangeProviderCapability::Simple(true)),
                hover_provider: Some(HoverProviderCapability::Simple(true)),
                inlay_hint_provider: Some(OneOf::Left(true)),
                rename_provider: Some(OneOf::Left(true)),
                semantic_tokens_provider: Some(SemanticTokensServerCapabilities::SemanticTokensRegistrationOptions(
                    SemanticTokensRegistrationOptions {
                        text_document_registration_options: {
                            TextDocumentRegistrationOptions {
                                document_selector: Some(vec![DocumentFilter {
                                    language: Some("kcl".to_string()),
                                    scheme: Some("file".to_string()),
                                    pattern: None,
                                }]),
                            }
                        },
                        semantic_tokens_options: SemanticTokensOptions {
                            work_done_progress_options: WorkDoneProgressOptions::default(),
                            legend: SemanticTokensLegend {
                                token_types: SEMANTIC_TOKEN_TYPES.to_vec(),
                                token_modifiers: SEMANTIC_TOKEN_MODIFIERS.to_vec(),
                            },
                            range: Some(false),
                            full: Some(SemanticTokensFullOptions::Bool(true)),
                        },
                        static_registration_options: StaticRegistrationOptions::default(),
                    },
                )),
                signature_help_provider: Some(SignatureHelpOptions {
                    trigger_characters: None,
                    retrigger_characters: None,
                    ..Default::default()
                }),
                text_document_sync: Some(TextDocumentSyncCapability::Options(TextDocumentSyncOptions {
                    open_close: Some(true),
                    change: Some(TextDocumentSyncKind::FULL),
                    ..Default::default()
                })),
                workspace: Some(WorkspaceServerCapabilities {
                    workspace_folders: Some(WorkspaceFoldersServerCapabilities {
                        supported: Some(true),
                        change_notifications: Some(OneOf::Left(true)),
                    }),
                    file_operations: None,
                }),
                ..Default::default()
            },
            ..Default::default()
        })
    }

    async fn initialized(&self, params: InitializedParams) {
        self.do_initialized(params).await
    }

    async fn shutdown(&self) -> RpcResult<()> {
        self.do_shutdown().await
    }

    async fn did_change_workspace_folders(&self, params: DidChangeWorkspaceFoldersParams) {
        self.do_did_change_workspace_folders(params).await
    }

    async fn did_change_configuration(&self, params: DidChangeConfigurationParams) {
        self.do_did_change_configuration(params).await
    }

    async fn did_change_watched_files(&self, params: DidChangeWatchedFilesParams) {
        self.do_did_change_watched_files(params).await
    }

    async fn did_create_files(&self, params: CreateFilesParams) {
        self.do_did_create_files(params).await
    }

    async fn did_rename_files(&self, params: RenameFilesParams) {
        self.do_did_rename_files(params).await
    }

    async fn did_delete_files(&self, params: DeleteFilesParams) {
        self.do_did_delete_files(params).await
    }

    async fn did_open(&self, params: DidOpenTextDocumentParams) {
        self.do_did_open(params).await
    }

    async fn did_change(&self, params: DidChangeTextDocumentParams) {
        self.do_did_change(params).await;
    }

    async fn did_save(&self, params: DidSaveTextDocumentParams) {
        self.do_did_save(params).await
    }

    async fn did_close(&self, params: DidCloseTextDocumentParams) {
        self.do_did_close(params).await;
    }

    async fn hover(&self, params: HoverParams) -> RpcResult<Option<LspHover>> {
        let filename = params.text_document_position_params.text_document.uri.to_string();

        let Some(current_code) = self.code_map.get(&filename) else {
            return Ok(None);
        };
        let Ok(current_code) = std::str::from_utf8(&current_code) else {
            return Ok(None);
        };

        let pos = position_to_char_index(params.text_document_position_params.position, current_code);

        // Let's iterate over the AST and find the node that contains the cursor.
        let Some(ast) = self.ast_map.get(&filename) else {
            return Ok(None);
        };

        let Some(hover) = ast
            .ast
            .get_hover_value_for_position(pos, current_code, &HoverOpts::default_for_hover())
        else {
            return Ok(None);
        };
        let stdlib_completions = self.stdlib_completions_for_position(&ast, pos);
        let stdlib_args = self.stdlib_args_for_position(&ast, pos);

        match hover {
            Hover::Function { name, range } => {
                let (sig, docs) = if let Some(Some(result)) = with_cached_var(&name, |value| {
                    match value {
                        // User-defined function
                        KclValue::Function { value, .. } if !value.is_std() => {
                            // TODO get docs from comments
                            Some((value.ast.signature(), ""))
                        }
                        _ => None,
                    }
                })
                .await
                {
                    result
                } else {
                    // Get the docs for this function.
                    let Some(completion) = stdlib_completions.get(&name) else {
                        return Ok(None);
                    };
                    let Some(docs) = &completion.documentation else {
                        return Ok(None);
                    };

                    let docs = match docs {
                        Documentation::String(docs) => docs,
                        Documentation::MarkupContent(MarkupContent { value, .. }) => value,
                    };

                    let docs = if docs.len() > 320 {
                        let end = docs.find("\n\n").or_else(|| docs.find("\n\r\n")).unwrap_or(320);
                        &docs[..end]
                    } else {
                        &**docs
                    };

                    let Some(label_details) = &completion.label_details else {
                        return Ok(None);
                    };

                    let sig = if let Some(detail) = &label_details.detail {
                        detail.clone()
                    } else {
                        String::new()
                    };

                    (sig, docs)
                };

                Ok(Some(LspHover {
                    contents: HoverContents::Markup(MarkupContent {
                        kind: MarkupKind::Markdown,
                        value: format!("```\n{name}{sig}\n```\n\n{docs}"),
                    }),
                    range: Some(range),
                }))
            }
            Hover::Type { name, range } => {
                let Some(completion) = stdlib_completions.get(&name) else {
                    return Ok(None);
                };
                let Some(docs) = &completion.documentation else {
                    return Ok(None);
                };

                let docs = match docs {
                    Documentation::String(docs) => docs,
                    Documentation::MarkupContent(MarkupContent { value, .. }) => value,
                };

                let docs = if docs.len() > 320 {
                    let end = docs.find("\n\n").or_else(|| docs.find("\n\r\n")).unwrap_or(320);
                    &docs[..end]
                } else {
                    &**docs
                };

                Ok(Some(LspHover {
                    contents: HoverContents::Markup(MarkupContent {
                        kind: MarkupKind::Markdown,
                        value: format!("```\n{name}\n```\n\n{docs}"),
                    }),
                    range: Some(range),
                }))
            }
            Hover::KwArg {
                name,
                callee_name,
                range,
            } => {
                // TODO handle user-defined functions too

                let Some(arg_map) = stdlib_args.get(&callee_name) else {
                    return Ok(None);
                };

                let Some(arg_entry) = arg_map.get(&name) else {
                    return Ok(None);
                };

                Ok(Some(LspHover {
                    contents: HoverContents::Markup(MarkupContent {
                        kind: MarkupKind::Markdown,
                        value: arg_entry.tip.clone(),
                    }),
                    range: Some(range),
                }))
            }
            Hover::Variable {
                name,
                ty: Some(ty),
                range,
            } => Ok(Some(LspHover {
                contents: HoverContents::Markup(MarkupContent {
                    kind: MarkupKind::Markdown,
                    value: format!("```\n{name}: {ty}\n```"),
                }),
                range: Some(range),
            })),
            Hover::Variable { name, ty: None, range } => Ok(with_cached_var(&name, |value| {
                let mut text: String = format!("```\n{name}");
                if let Some(ty) = value.principal_type() {
                    text.push_str(&format!(": {}", ty.human_friendly_type()));
                }
                if let Some(v) = value.value_str() {
                    text.push_str(&format!(" = {v}"));
                }
                text.push_str("\n```");

                LspHover {
                    contents: HoverContents::Markup(MarkupContent {
                        kind: MarkupKind::Markdown,
                        value: text,
                    }),
                    range: Some(range),
                }
            })
            .await),
            Hover::Signature { .. } => Ok(None),
            Hover::Comment { value, range } => Ok(Some(LspHover {
                contents: HoverContents::Markup(MarkupContent {
                    kind: MarkupKind::Markdown,
                    value,
                }),
                range: Some(range),
            })),
        }
    }

    async fn completion(&self, params: CompletionParams) -> RpcResult<Option<CompletionResponse>> {
        let mut completions = vec![CompletionItem {
            label: PIPE_OPERATOR.to_string(),
            label_details: None,
            kind: Some(CompletionItemKind::OPERATOR),
            detail: Some("A pipe operator.".to_string()),
            documentation: Some(Documentation::MarkupContent(MarkupContent {
                kind: MarkupKind::Markdown,
                value: "A pipe operator.".to_string(),
            })),
            deprecated: Some(false),
            preselect: None,
            sort_text: None,
            filter_text: None,
            insert_text: Some("|> ".to_string()),
            insert_text_format: Some(InsertTextFormat::PLAIN_TEXT),
            insert_text_mode: None,
            text_edit: None,
            additional_text_edits: None,
            command: None,
            commit_characters: None,
            data: None,
            tags: None,
        }];

        // Get the current line up to cursor
        let Some(current_code) = self
            .code_map
            .get(params.text_document_position.text_document.uri.as_ref())
        else {
            return Ok(Some(CompletionResponse::Array(completions)));
        };
        let Ok(current_code) = std::str::from_utf8(&current_code) else {
            return Ok(Some(CompletionResponse::Array(completions)));
        };

        // Get the current line up to cursor, with bounds checking
        if let Some(line) = current_code
            .lines()
            .nth(params.text_document_position.position.line as usize)
        {
            let char_pos = params.text_document_position.position.character as usize;
            if char_pos <= line.len() {
                let line_prefix = &line[..char_pos];
                // Get last word
                let last_word = line_prefix
                    .split(|c: char| c.is_whitespace() || c.is_ascii_punctuation())
                    .next_back()
                    .unwrap_or("");

                // If the last word starts with a digit, return no completions
                if !last_word.is_empty() && last_word.chars().next().unwrap().is_ascii_digit() {
                    return Ok(None);
                }
            }
        }

        // Add more to the completions if we have more.
        let Some(ast) = self
            .ast_map
            .get(params.text_document_position.text_document.uri.as_ref())
        else {
            completions.extend(self.stdlib_completions.values().cloned());
            completions.extend(self.kcl_keywords.values().cloned());
            return Ok(Some(CompletionResponse::Array(completions)));
        };

        let position = position_to_char_index(params.text_document_position.position, current_code);
        if ast.ast.in_comment(position) {
            // If we are in a code comment we don't want to show completions.
            return Ok(None);
        }

        completions.extend(self.stdlib_completions_for_position(&ast, position).values().cloned());
        completions.extend(self.kcl_keywords.values().cloned());

        // If we're inside a CallExpression or something where a function parameter label could be completed,
        // then complete it.
        // Let's find the AST node that the user's cursor is in.
        if let Some(arg_label_completions) = self.try_arg_completions(&ast, position, current_code) {
            completions.extend(arg_label_completions);
        }

        // Get the completion items for the ast.
        let Ok(variables) = ast.ast.completion_items(position) else {
            return Ok(Some(CompletionResponse::Array(completions)));
        };

        // Get our variables from our AST to include in our completions.
        completions.extend(variables);

        Ok(Some(CompletionResponse::Array(completions)))
    }

    async fn diagnostic(&self, params: DocumentDiagnosticParams) -> RpcResult<DocumentDiagnosticReportResult> {
        let filename = params.text_document.uri.to_string();

        // Get the current diagnostics for this file.
        let Some(items) = self.diagnostics_map.get(&filename) else {
            // Send an empty report.
            return Ok(DocumentDiagnosticReportResult::Report(DocumentDiagnosticReport::Full(
                RelatedFullDocumentDiagnosticReport {
                    related_documents: None,
                    full_document_diagnostic_report: FullDocumentDiagnosticReport {
                        result_id: None,
                        items: vec![],
                    },
                },
            )));
        };

        Ok(DocumentDiagnosticReportResult::Report(DocumentDiagnosticReport::Full(
            RelatedFullDocumentDiagnosticReport {
                related_documents: None,
                full_document_diagnostic_report: FullDocumentDiagnosticReport {
                    result_id: None,
                    items: items.clone(),
                },
            },
        )))
    }

    async fn signature_help(&self, params: SignatureHelpParams) -> RpcResult<Option<SignatureHelp>> {
        let filename = params.text_document_position_params.text_document.uri.to_string();

        let Some(current_code) = self.code_map.get(&filename) else {
            return Ok(None);
        };
        let Ok(current_code) = std::str::from_utf8(&current_code) else {
            return Ok(None);
        };

        let pos = position_to_char_index(params.text_document_position_params.position, current_code);

        let ast = self.ast_map.get(&filename);
        let stdlib_signatures = ast
            .as_ref()
            .map(|ast| self.stdlib_signatures_for_position(ast, pos))
            .unwrap_or(&self.stdlib_signatures);

        // Get the character at the position.
        let Some(ch) = current_code.chars().nth(pos) else {
            return Ok(None);
        };

        let check_char = |ch: char| {
            // If  we are on a (, then get the string in front of the (
            // and try to get the signature.
            // We do these before the ast check because we might not have a valid ast.
            if ch == '(' {
                // If the current character is not a " " then get the next space after
                // our position so we can split on that.
                // Find the next space after the current position.
                let next_space = if ch != ' ' {
                    if let Some(next_space) = current_code[pos..].find(' ') {
                        pos + next_space
                    } else if let Some(next_space) = current_code[pos..].find('(') {
                        pos + next_space
                    } else {
                        pos
                    }
                } else {
                    pos
                };
                let p2 = std::cmp::max(pos, next_space);

                let last_word = current_code[..p2].split_whitespace().last()?;

                // Get the function name.
                return stdlib_signatures.get(last_word);
            } else if ch == ',' {
                // If we have a comma, then get the string in front of
                // the closest ( and try to get the signature.

                // Find the last ( before the comma.
                let last_paren = current_code[..pos].rfind('(')?;
                // Get the string in front of the (.
                let last_word = current_code[..last_paren].split_whitespace().last()?;
                // Get the function name.
                return stdlib_signatures.get(last_word);
            }

            None
        };

        if let Some(signature) = check_char(ch) {
            return Ok(Some(signature.clone()));
        }

        // Check if we have context.
        if let Some(context) = params.context
            && let Some(character) = context.trigger_character
        {
            for character in character.chars() {
                // Check if we are on a ( or a ,.
                if (character == '(' || character == ',')
                    && let Some(signature) = check_char(character)
                {
                    return Ok(Some(signature.clone()));
                }
            }
        }

        let Some(ast) = ast else {
            return Ok(None);
        };

        let Some(value) = ast.ast.get_expr_for_position(pos) else {
            return Ok(None);
        };

        let Some(hover) =
            value.get_hover_value_for_position(pos, current_code, &HoverOpts::default_for_signature_help())
        else {
            return Ok(None);
        };

        match hover {
            Hover::Function { name, range: _ } => {
                // Get the docs for this function.
                let Some(signature) = stdlib_signatures.get(&name) else {
                    return Ok(None);
                };

                Ok(Some(signature.clone()))
            }
            Hover::Signature {
                name,
                parameter_index,
                range: _,
            } => {
                let Some(signature) = stdlib_signatures.get(&name) else {
                    return Ok(None);
                };

                let mut signature = signature.clone();

                signature.active_parameter = Some(parameter_index);

                Ok(Some(signature))
            }
            _ => {
                return Ok(None);
            }
        }
    }

    async fn inlay_hint(&self, _params: InlayHintParams) -> RpcResult<Option<Vec<InlayHint>>> {
        // TODO: do this

        Ok(None)
    }

    async fn semantic_tokens_full(&self, params: SemanticTokensParams) -> RpcResult<Option<SemanticTokensResult>> {
        let filename = params.text_document.uri.to_string();

        let Some(semantic_tokens) = self.semantic_tokens_map.get(&filename) else {
            return Ok(None);
        };

        Ok(Some(SemanticTokensResult::Tokens(SemanticTokens {
            result_id: None,
            data: semantic_tokens.clone(),
        })))
    }

    async fn document_symbol(&self, params: DocumentSymbolParams) -> RpcResult<Option<DocumentSymbolResponse>> {
        let filename = params.text_document.uri.to_string();

        let Some(symbols) = self.symbols_map.get(&filename) else {
            return Ok(None);
        };

        Ok(Some(DocumentSymbolResponse::Nested(symbols.clone())))
    }

    async fn formatting(&self, params: DocumentFormattingParams) -> RpcResult<Option<Vec<TextEdit>>> {
        let filename = params.text_document.uri.to_string();

        let Some(current_code) = self.code_map.get(&filename) else {
            return Ok(None);
        };
        let Ok(current_code) = std::str::from_utf8(&current_code) else {
            return Ok(None);
        };

        // Parse the ast.
        // I don't know if we need to do this again since it should be updated in the context.
        // But I figure better safe than sorry since this will write back out to the file.
        let module_id = ModuleId::default();
        let Ok(ast) = crate::parsing::parse_str(current_code, module_id).parse_errs_as_err() else {
            return Ok(None);
        };
        // Now recast it.
        let recast = ast.recast_top(
            &crate::parsing::ast::types::FormatOptions {
                tab_size: params.options.tab_size as usize,
                insert_final_newline: params.options.insert_final_newline.unwrap_or(false),
                use_tabs: !params.options.insert_spaces,
            },
            0,
        );
        let source_range = SourceRange::new(0, current_code.len(), module_id);
        let range = source_range.to_lsp_range(current_code);
        Ok(Some(vec![TextEdit {
            new_text: recast,
            range,
        }]))
    }

    async fn rename(&self, params: RenameParams) -> RpcResult<Option<WorkspaceEdit>> {
        let Some((current_code, new_code)) =
            self.inner_prepare_rename(&params.text_document_position, &params.new_name)?
        else {
            return Ok(None);
        };

        let source_range = SourceRange::new(0, current_code.len(), ModuleId::default());
        let range = source_range.to_lsp_range(&current_code);
        Ok(Some(WorkspaceEdit {
            changes: Some(HashMap::from([(
                params.text_document_position.text_document.uri,
                vec![TextEdit {
                    new_text: new_code,
                    range,
                }],
            )])),
            document_changes: None,
            change_annotations: None,
        }))
    }

    async fn prepare_rename(&self, params: TextDocumentPositionParams) -> RpcResult<Option<PrepareRenameResponse>> {
        if self
            .inner_prepare_rename(&params, "someNameNoOneInTheirRightMindWouldEverUseForTesting")?
            .is_none()
        {
            return Ok(None);
        }

        // Return back to the client, that it is safe to use the rename behavior.
        Ok(Some(PrepareRenameResponse::DefaultBehavior { default_behavior: true }))
    }

    async fn folding_range(&self, params: FoldingRangeParams) -> RpcResult<Option<Vec<FoldingRange>>> {
        let filename = params.text_document.uri.to_string();

        // Get the ast.
        let Some(ast) = self.ast_map.get(&filename) else {
            return Ok(None);
        };

        // Get the folding ranges.
        let folding_ranges = ast.ast.get_lsp_folding_ranges();

        if folding_ranges.is_empty() {
            return Ok(None);
        }

        Ok(Some(folding_ranges))
    }

    async fn code_action(&self, params: CodeActionParams) -> RpcResult<Option<CodeActionResponse>> {
        // Get the actual diagnostics from our map to ensure we have the full diagnostic info
        let filename = params.text_document.uri.to_string();
        let stored_diagnostics = self
            .diagnostics_map
            .get(&filename)
            .map(|d| d.clone())
            .unwrap_or_default();

        // Use stored diagnostics if available, otherwise fall back to params.context.diagnostics
        let diagnostics_to_check: Vec<_> = if stored_diagnostics.is_empty() {
            params.context.diagnostics.clone()
        } else {
            // Match params.context.diagnostics with stored diagnostics by range
            params
                .context
                .diagnostics
                .iter()
                .filter_map(|param_diag| {
                    stored_diagnostics
                        .iter()
                        .find(|stored_diag| stored_diag.range == param_diag.range)
                        .cloned()
                })
                .collect()
        };

        let actions = diagnostics_to_check
            .into_iter()
            .filter_map(|diagnostic| {
                // Handle regular suggestions (like camelCase)
                let (suggestion, range) = diagnostic
                    .data
                    .as_ref()
                    .and_then(|data| serde_json::from_value::<LspSuggestion>(data.clone()).ok())?;
                let edit = TextEdit {
                    range,
                    new_text: suggestion.insert,
                };
                let changes = HashMap::from([(params.text_document.uri.clone(), vec![edit])]);

                // If you add more code action kinds, make sure you add it to the server
                // capabilities on initialization!
                Some(CodeActionOrCommand::CodeAction(CodeAction {
                    title: suggestion.title,
                    kind: Some(CodeActionKind::QUICKFIX),
                    diagnostics: Some(vec![diagnostic]),
                    edit: Some(WorkspaceEdit {
                        changes: Some(changes),
                        document_changes: None,
                        change_annotations: None,
                    }),
                    command: None,
                    is_preferred: Some(true),
                    disabled: None,
                    data: None,
                }))
            })
            .collect();

        Ok(Some(actions))
    }

    async fn document_color(&self, params: DocumentColorParams) -> RpcResult<Vec<ColorInformation>> {
        let filename = params.text_document.uri.to_string();

        let Some(current_code) = self.code_map.get(&filename) else {
            return Ok(vec![]);
        };
        let Ok(current_code) = std::str::from_utf8(&current_code) else {
            return Ok(vec![]);
        };

        // Get the ast from our map.
        let Some(ast) = self.ast_map.get(&filename) else {
            return Ok(vec![]);
        };

        // Get the colors from the ast.
        let Ok(colors) = ast.ast.document_color(current_code) else {
            return Ok(vec![]);
        };

        Ok(colors)
    }

    async fn color_presentation(&self, params: ColorPresentationParams) -> RpcResult<Vec<ColorPresentation>> {
        let filename = params.text_document.uri.to_string();

        let Some(current_code) = self.code_map.get(&filename) else {
            return Ok(vec![]);
        };
        let Ok(current_code) = std::str::from_utf8(&current_code) else {
            return Ok(vec![]);
        };

        // Get the ast from our map.
        let Some(ast) = self.ast_map.get(&filename) else {
            return Ok(vec![]);
        };

        let pos_start = position_to_char_index(params.range.start, current_code);
        let pos_end = position_to_char_index(params.range.end, current_code);

        // Get the colors from the ast.
        let Ok(Some(presentation)) = ast.ast.color_presentation(&params.color, pos_start, pos_end) else {
            return Ok(vec![]);
        };

        Ok(vec![presentation])
    }
}

#[derive(Clone, Copy, Debug, Eq, PartialEq)]
enum StdlibCompletionContext {
    Default,
    SketchBlock,
}

fn is_sketch2_doc(doc: &crate::docs::kcl_doc::DocData) -> bool {
    doc.qual_name().starts_with("std::solver::")
}

fn strip_sketch2_prefix(value: &str) -> String {
    value.strip_prefix("solver::").unwrap_or(value).to_owned()
}

fn rewrite_completion_for_sketch_block(mut completion: CompletionItem) -> CompletionItem {
    completion.label = strip_sketch2_prefix(&completion.label);
    completion.insert_text = completion.insert_text.map(|text| strip_sketch2_prefix(&text));
    completion
}

fn rewrite_signature_for_sketch_block(mut signature: SignatureHelp) -> SignatureHelp {
    for info in &mut signature.signatures {
        info.label = strip_sketch2_prefix(&info.label);
    }
    signature
}

fn should_skip_stdlib_doc(
    doc: &crate::docs::kcl_doc::DocData,
    has_existing: bool,
    context: StdlibCompletionContext,
) -> bool {
    match context {
        StdlibCompletionContext::Default => is_sketch2_doc(doc) || (has_existing && doc.is_experimental()),
        StdlibCompletionContext::SketchBlock => has_existing && doc.is_experimental() && !is_sketch2_doc(doc),
    }
}

/// Get completions from our stdlib.
pub fn get_completions_from_stdlib(kcl_std: &ModData) -> Result<HashMap<String, CompletionItem>> {
    get_completions_from_stdlib_in_context(kcl_std, StdlibCompletionContext::Default)
}

pub fn get_completions_from_stdlib_for_sketch_block(kcl_std: &ModData) -> Result<HashMap<String, CompletionItem>> {
    get_completions_from_stdlib_in_context(kcl_std, StdlibCompletionContext::SketchBlock)
}

fn get_completions_from_stdlib_in_context(
    kcl_std: &ModData,
    context: StdlibCompletionContext,
) -> Result<HashMap<String, CompletionItem>> {
    let mut completions = HashMap::new();

    for d in kcl_std.all_docs() {
        if let Some(mut ci) = d.to_completion_item() {
            let preferred = d.preferred_name();
            if should_skip_stdlib_doc(d, completions.contains_key(preferred), context) {
                continue;
            }
            let key = if context == StdlibCompletionContext::SketchBlock && is_sketch2_doc(d) {
                ci = rewrite_completion_for_sketch_block(ci);
                strip_sketch2_prefix(preferred)
            } else {
                preferred.to_owned()
            };
            completions.insert(key, ci);
        }
    }

    let variable_kinds = VariableKind::to_completion_items();
    for variable_kind in variable_kinds {
        completions.insert(variable_kind.label.clone(), variable_kind);
    }

    Ok(completions)
}

/// Get signatures from our stdlib.
pub fn get_signatures_from_stdlib(kcl_std: &ModData) -> HashMap<String, SignatureHelp> {
    get_signatures_from_stdlib_in_context(kcl_std, StdlibCompletionContext::Default)
}

pub fn get_signatures_from_stdlib_for_sketch_block(kcl_std: &ModData) -> HashMap<String, SignatureHelp> {
    get_signatures_from_stdlib_in_context(kcl_std, StdlibCompletionContext::SketchBlock)
}

fn get_signatures_from_stdlib_in_context(
    kcl_std: &ModData,
    context: StdlibCompletionContext,
) -> HashMap<String, SignatureHelp> {
    let mut signatures = HashMap::new();

    for d in kcl_std.all_docs() {
        if let Some(mut sig) = d.to_signature_help() {
            let preferred = d.preferred_name();
            if should_skip_stdlib_doc(d, signatures.contains_key(preferred), context) {
                continue;
            }
            let key = if context == StdlibCompletionContext::SketchBlock && is_sketch2_doc(d) {
                sig = rewrite_signature_for_sketch_block(sig);
                strip_sketch2_prefix(preferred)
            } else {
                preferred.to_owned()
            };
            signatures.insert(key, sig);
        }
    }

    signatures
}

/// Get KCL keywords
pub fn get_keywords() -> HashMap<String, CompletionItem> {
    RESERVED_WORDS
        .keys()
        .map(|k| (k.to_string(), keyword_to_completion(k.to_string())))
        .collect()
}

fn keyword_to_completion(kw: String) -> CompletionItem {
    CompletionItem {
        label: kw,
        label_details: None,
        kind: Some(CompletionItemKind::KEYWORD),
        detail: None,
        documentation: None,
        deprecated: Some(false),
        preselect: None,
        sort_text: None,
        filter_text: None,
        insert_text: None,
        insert_text_format: None,
        insert_text_mode: None,
        text_edit: None,
        additional_text_edits: None,
        command: None,
        commit_characters: None,
        data: None,
        tags: None,
    }
}

#[derive(Clone, Debug)]
pub struct LspArgData {
    pub tip: String,
    pub props: ArgData,
}

/// Get signatures from our stdlib.
pub fn get_arg_maps_from_stdlib(kcl_std: &ModData) -> HashMap<String, HashMap<String, LspArgData>> {
    get_arg_maps_from_stdlib_in_context(kcl_std, StdlibCompletionContext::Default)
}

pub fn get_arg_maps_from_stdlib_for_sketch_block(kcl_std: &ModData) -> HashMap<String, HashMap<String, LspArgData>> {
    get_arg_maps_from_stdlib_in_context(kcl_std, StdlibCompletionContext::SketchBlock)
}

fn get_arg_maps_from_stdlib_in_context(
    kcl_std: &ModData,
    context: StdlibCompletionContext,
) -> HashMap<String, HashMap<String, LspArgData>> {
    let mut result = HashMap::new();

    for d in kcl_std.all_docs() {
        let crate::docs::kcl_doc::DocData::Fn(f) = d else {
            continue;
        };
        let arg_map: HashMap<String, _> = f
            .args
            .iter()
            .map(|data| {
                let mut tip = "```\n".to_owned();
                tip.push_str(&data.to_string());
                tip.push_str("\n```");
                if let Some(docs) = &data.docs {
                    tip.push_str("\n\n");
                    tip.push_str(docs);
                }
                let arg_data = LspArgData {
                    tip,
                    props: data.clone(),
                };
                (data.name.clone(), arg_data)
            })
            .collect();
        if !arg_map.is_empty() {
            let key = if context == StdlibCompletionContext::SketchBlock && is_sketch2_doc(d) {
                strip_sketch2_prefix(&f.preferred_name)
            } else {
                f.preferred_name.clone()
            };
            if should_skip_stdlib_doc(d, result.contains_key(&key), context) {
                continue;
            }
            result.insert(key, arg_map);
        }
    }

    result
}

/// Convert a position to a character index from the start of the file.
fn position_to_char_index(position: Position, code: &str) -> usize {
    // Get the character position from the start of the file.
    let mut char_position = 0;
    for (index, line) in code.lines().enumerate() {
        if index == position.line as usize {
            char_position += position.character as usize;
            break;
        } else {
            char_position += line.len() + 1;
        }
    }

    let end_of_file = if code.is_empty() { 0 } else { code.len() - 1 };
    std::cmp::min(char_position, end_of_file)
}

async fn with_cached_var<T>(name: &str, f: impl Fn(&KclValue) -> T) -> Option<T> {
    let mem = cache::read_old_memory().await?;
    let value = mem.stack.get(name, SourceRange::default()).ok()?;

    Some(f(&value))
}

#[cfg(test)]
mod tests {
    use pretty_assertions::assert_eq;

    use super::*;

    #[test]
    fn test_position_to_char_index_first_line() {
        let code = r#"def foo():
return 42"#;
        let position = Position::new(0, 3);
        let index = position_to_char_index(position, code);
        assert_eq!(index, 3);
    }

    #[test]
    fn test_position_to_char_index() {
        let code = r#"def foo():
return 42"#;
        let position = Position::new(1, 4);
        let index = position_to_char_index(position, code);
        assert_eq!(index, 15);
    }

    #[test]
    fn test_position_to_char_index_with_newline() {
        let code = r#"def foo():

return 42"#;
        let position = Position::new(2, 0);
        let index = position_to_char_index(position, code);
        assert_eq!(index, 12);
    }

    #[test]
    fn test_position_to_char_at_end() {
        let code = r#"def foo():
return 42"#;

        let position = Position::new(1, 8);
        let index = position_to_char_index(position, code);
        assert_eq!(index, 19);
    }

    #[test]
    fn test_position_to_char_empty() {
        let code = r#""#;
        let position = Position::new(0, 0);
        let index = position_to_char_index(position, code);
        assert_eq!(index, 0);
    }

    #[test]
    fn stdlib_completions_keyed_by_preferred_name() {
        let stdlib = crate::docs::kcl_doc::walk_stdlib();
        let completions = get_completions_from_stdlib(&stdlib).unwrap();

        assert!(
            completions.contains_key("clone"),
            "prelude fn should be keyed by short name"
        );
        assert!(
            completions.contains_key("units::toRadians"),
            "module fn should be keyed by preferred name with module prefix"
        );
        assert!(
            !completions.contains_key("toRadians"),
            "module fn should NOT be keyed by bare short name"
        );
    }

    #[test]
    fn stdlib_signatures_keyed_by_preferred_name() {
        let stdlib = crate::docs::kcl_doc::walk_stdlib();
        let signatures = get_signatures_from_stdlib(&stdlib);

        assert!(
            signatures.contains_key("clone"),
            "prelude fn should be keyed by short name"
        );
        assert!(
            signatures.contains_key("units::toRadians"),
            "module fn should be keyed by preferred name with module prefix"
        );
        assert!(
            !signatures.contains_key("toRadians"),
            "module fn should NOT be keyed by bare short name"
        );
    }

    #[test]
    fn stdlib_args_keyed_by_preferred_name() {
        let stdlib = crate::docs::kcl_doc::walk_stdlib();
        let args = get_arg_maps_from_stdlib(&stdlib);

        assert!(args.contains_key("clone"), "prelude fn should be keyed by short name");
        assert!(
            args.contains_key("units::toRadians"),
            "module fn should be keyed by preferred name with module prefix"
        );
        assert!(
            !args.contains_key("toRadians"),
            "module fn should NOT be keyed by bare short name"
        );
    }

    #[test]
    fn default_completions_exclude_solver() {
        let stdlib = crate::docs::kcl_doc::walk_stdlib();
        let completions = get_completions_from_stdlib(&stdlib).unwrap();

        assert!(
            completions.contains_key("line"),
            "prelude line should be present in default context"
        );
        assert!(
            !completions.contains_key("solver::line"),
            "solver::line should NOT appear in default context"
        );
        assert!(
            !completions.contains_key("solver::coincident"),
            "solver-only functions should NOT appear in default context"
        );
    }

    #[test]
    fn sketch1_block_completions_exclude_solver() {
        // Sketch v1 (startSketchOn()) blocks use the Default completions map (is_in_sketch_block
        // returns false for sketch v1 syntax);  solver functions must be absent.
        let stdlib = crate::docs::kcl_doc::walk_stdlib();
        let completions = get_completions_from_stdlib(&stdlib).unwrap();

        assert!(
            completions.contains_key("line"),
            "prelude line should be present for old sketch blocks"
        );
        assert!(
            !completions.contains_key("solver::line"),
            "solver::line should NOT appear for old sketch blocks"
        );
        assert!(
            !completions.contains_key("solver::coincident"),
            "solver-only functions should NOT appear for old sketch blocks"
        );
    }

    #[test]
    fn sketch2_completions_include_solver() {
        let stdlib = crate::docs::kcl_doc::walk_stdlib();
        let completions = get_completions_from_stdlib_for_sketch_block(&stdlib).unwrap();

        // Solver functions should shadow prelude functions under their bare name
        assert!(
            completions.contains_key("line"),
            "solver line should be present under bare name in sketch2 context"
        );
        let line = &completions["line"];
        assert!(
            line.insert_text
                .as_ref()
                .is_some_and(|t| t.starts_with("line(start = ")),
            "sketch2 line should use solver signature (start =), not prelude signature"
        );

        // Solver-only functions should also be present under bare name
        assert!(
            completions.contains_key("coincident"),
            "solver-only functions should be present under bare name in sketch2 context"
        );
        assert!(
            !completions.contains_key("solver::coincident"),
            "solver functions should NOT keep solver:: prefix in sketch2 context"
        );
    }
}