use std::collections::HashMap;

use ion_core::ast::{Param, StmtKind, UseImports};
use ion_core::error::IonError;
use ion_core::lexer::Lexer;
use ion_core::parser::Parser;
use ion_core::token::Token;

use lsp_server::{
    Connection, ErrorCode, Message, Notification as LspNotification, Request, RequestId,
    Response, ResponseError,
};
use lsp_types::notification::{
    self, DidChangeTextDocument, DidOpenTextDocument, DidSaveTextDocument,
    Notification as NotificationTrait,
};
use lsp_types::request::{
    Completion, DocumentSymbolRequest, GotoDefinition, HoverRequest, References, Rename,
    Request as RequestTrait,
};
use lsp_types::{
    CompletionItem, CompletionItemKind, CompletionList, CompletionOptions, CompletionParams,
    CompletionResponse, Diagnostic, DiagnosticSeverity, DidChangeTextDocumentParams,
    DidOpenTextDocumentParams, DocumentSymbol, DocumentSymbolParams, DocumentSymbolResponse,
    GotoDefinitionParams, GotoDefinitionResponse, Hover, HoverContents, HoverParams,
    HoverProviderCapability, InitializeResult, Location, MarkupContent, MarkupKind,
    OneOf, Position, Range, ReferenceParams, RenameParams, ServerCapabilities, SymbolKind,
    TextDocumentSyncCapability, TextDocumentSyncKind, TextEdit, Url, WorkspaceEdit,
};

// ---- Definition tracking ----

#[derive(Debug, Clone)]
struct Definition {
    name: String,
    kind: DefKind,
    line: u32, // 0-based
    col: u32,
    detail: String,
}

#[derive(Debug, Clone)]
enum DefKind {
    Function,
    Variable,
}

fn collect_definitions(source: &str) -> Vec<Definition> {
    let mut lexer = Lexer::new(source);
    let tokens = match lexer.tokenize() {
        Ok(t) => t,
        Err(_) => return vec![],
    };
    let mut parser = Parser::new(tokens);
    // Use recovering parse so definitions are available even with errors
    let output = parser.parse_program_recovering();
    let mut defs = Vec::new();
    collect_defs_from_stmts(&output.program.stmts, &mut defs);
    defs
}

fn collect_defs_from_stmts(stmts: &[ion_core::ast::Stmt], defs: &mut Vec<Definition>) {
    for stmt in stmts {
        let line = if stmt.span.line > 0 {
            stmt.span.line as u32 - 1
        } else {
            0
        };
        let col = if stmt.span.col > 0 {
            stmt.span.col as u32 - 1
        } else {
            0
        };
        match &stmt.kind {
            StmtKind::FnDecl { name, params, body } => {
                defs.push(Definition {
                    name: name.clone(),
                    kind: DefKind::Function,
                    line,
                    col,
                    detail: format_fn_signature(name, params),
                });
                collect_defs_from_stmts(body, defs);
            }
            StmtKind::Let {
                pattern: ion_core::ast::Pattern::Ident(name),
                ..
            } => {
                defs.push(Definition {
                    name: name.clone(),
                    kind: DefKind::Variable,
                    line,
                    col,
                    detail: format!("let {}", name),
                });
            }
            StmtKind::For { body, .. } => collect_defs_from_stmts(body, defs),
            StmtKind::While { body, .. } => collect_defs_from_stmts(body, defs),
            StmtKind::WhileLet { body, .. } => collect_defs_from_stmts(body, defs),
            StmtKind::Loop { body, .. } => collect_defs_from_stmts(body, defs),
            StmtKind::ExprStmt { expr, .. } => collect_defs_from_expr(expr, defs),
            StmtKind::Use { path, imports } => {
                let module_path = path.join("::");
                let names: Vec<String> = match imports {
                    UseImports::Single(name) => vec![name.clone()],
                    UseImports::Names(names) => names.clone(),
                    UseImports::Glob => vec![format!("{}::*", module_path)],
                };
                for name in names {
                    defs.push(Definition {
                        name: name.clone(),
                        kind: DefKind::Variable,
                        line,
                        col,
                        detail: format!("use {}::{}", module_path, name),
                    });
                }
            }
            _ => {}
        }
    }
}

fn collect_defs_from_expr(expr: &ion_core::ast::Expr, defs: &mut Vec<Definition>) {
    use ion_core::ast::ExprKind;
    match &expr.kind {
        ExprKind::If {
            then_body,
            else_body,
            ..
        } => {
            collect_defs_from_stmts(then_body, defs);
            if let Some(eb) = else_body {
                collect_defs_from_stmts(eb, defs);
            }
        }
        ExprKind::IfLet {
            then_body,
            else_body,
            ..
        } => {
            collect_defs_from_stmts(then_body, defs);
            if let Some(eb) = else_body {
                collect_defs_from_stmts(eb, defs);
            }
        }
        ExprKind::Match { arms, .. } => {
            for arm in arms {
                collect_defs_from_expr(&arm.body, defs);
            }
        }
        ExprKind::Block(stmts) => collect_defs_from_stmts(stmts, defs),
        ExprKind::TryCatch { body, handler, .. } => {
            collect_defs_from_stmts(body, defs);
            collect_defs_from_stmts(handler, defs);
        }
        _ => {}
    }
}

fn format_fn_signature(name: &str, params: &[Param]) -> String {
    let params_str: Vec<String> = params
        .iter()
        .map(|p| {
            if let Some(default) = &p.default {
                format!("{} = {:?}", p.name, default)
            } else {
                p.name.clone()
            }
        })
        .collect();
    format!("fn {}({})", name, params_str.join(", "))
}

// ---- Word at position ----

fn word_at_position(source: &str, line: u32, col: u32) -> Option<String> {
    let target_line = source.lines().nth(line as usize)?;
    let col = col as usize;
    if col > target_line.len() {
        return None;
    }
    let bytes = target_line.as_bytes();
    let mut start = col;
    while start > 0 && is_ident_char(bytes[start - 1]) {
        start -= 1;
    }
    let mut end = col;
    while end < bytes.len() && is_ident_char(bytes[end]) {
        end += 1;
    }
    if start == end {
        return None;
    }
    Some(target_line[start..end].to_string())
}

fn is_ident_char(b: u8) -> bool {
    b.is_ascii_alphanumeric() || b == b'_'
}

// ---- Builtins ----

struct BuiltinInfo {
    name: &'static str,
    signature: &'static str,
    description: &'static str,
}

const BUILTINS: &[BuiltinInfo] = &[
    BuiltinInfo {
        name: "len",
        signature: "len(x)",
        description: "Length of list, string, dict, or bytes",
    },
    BuiltinInfo {
        name: "range",
        signature: "range(n) / range(start, end)",
        description: "Create a range [0..n) or [start..end)",
    },
    BuiltinInfo {
        name: "enumerate",
        signature: "enumerate(list)",
        description: "List of (index, value) tuples",
    },
    BuiltinInfo {
        name: "type_of",
        signature: "type_of(x)",
        description: "Returns type name as string",
    },
    BuiltinInfo {
        name: "str",
        signature: "str(x)",
        description: "Convert to string",
    },
    BuiltinInfo {
        name: "int",
        signature: "int(x)",
        description: "Convert to int",
    },
    BuiltinInfo {
        name: "float",
        signature: "float(x)",
        description: "Convert to float",
    },
    BuiltinInfo {
        name: "bytes",
        signature: "bytes() / bytes(list) / bytes(string) / bytes(n)",
        description: "Create bytes",
    },
    BuiltinInfo {
        name: "bytes_from_hex",
        signature: "bytes_from_hex(string)",
        description: "Bytes from hex string",
    },
    BuiltinInfo {
        name: "assert",
        signature: "assert(cond) / assert(cond, msg)",
        description: "Error if condition is false",
    },
    BuiltinInfo {
        name: "assert_eq",
        signature: "assert_eq(a, b) / assert_eq(a, b, msg)",
        description: "Error if values are not equal",
    },
    BuiltinInfo {
        name: "channel",
        signature: "channel(buffer_size)",
        description: "Create a buffered channel (tx, rx)",
    },
    BuiltinInfo {
        name: "set",
        signature: "set() / set(list)",
        description: "Create a set from a list (deduplicates elements)",
    },
    BuiltinInfo {
        name: "cell",
        signature: "cell(value)",
        description: "Create a mutable reference cell for shared closure state",
    },
    BuiltinInfo {
        name: "sleep",
        signature: "sleep(ms)",
        description: "Sleep for given milliseconds (requires concurrency feature)",
    },
    BuiltinInfo {
        name: "timeout",
        signature: "timeout(ms, fn)",
        description: "Run function with time limit, returns Option (Some or None on timeout)",
    },
];

const KEYWORDS: &[&str] = &[
    "let", "mut", "fn", "if", "else", "while", "for", "loop", "break", "continue", "return",
    "match", "in", "true", "false", "None", "Some", "Ok", "Err", "async", "spawn", "select", "try",
    "catch", "use",
];

// ---- Main ----

fn main() {
    let (connection, io_threads) = Connection::stdio();

    let capabilities = ServerCapabilities {
        text_document_sync: Some(TextDocumentSyncCapability::Kind(TextDocumentSyncKind::FULL)),
        document_symbol_provider: Some(OneOf::Left(true)),
        definition_provider: Some(OneOf::Left(true)),
        references_provider: Some(OneOf::Left(true)),
        rename_provider: Some(OneOf::Left(true)),
        hover_provider: Some(HoverProviderCapability::Simple(true)),
        completion_provider: Some(CompletionOptions {
            trigger_characters: Some(vec![".".to_string()]),
            ..Default::default()
        }),
        ..Default::default()
    };

    let init_result = InitializeResult {
        capabilities,
        server_info: Some(lsp_types::ServerInfo {
            name: "ion-lsp".to_string(),
            version: Some("0.1.0".to_string()),
        }),
    };

    let init_json = serde_json::to_value(init_result).unwrap();
    connection.initialize(init_json).unwrap();

    let mut documents: HashMap<Url, String> = HashMap::new();

    for msg in &connection.receiver {
        match msg {
            Message::Request(req) => {
                if connection.handle_shutdown(&req).unwrap() {
                    break;
                }
                handle_request(&connection, &documents, req);
            }
            Message::Notification(not) => {
                handle_notification(&connection, &mut documents, not);
            }
            Message::Response(_) => {}
        }
    }

    io_threads.join().unwrap();
}

// ---- Request handling ----

fn handle_request(conn: &Connection, documents: &HashMap<Url, String>, req: Request) {
    if req.method == DocumentSymbolRequest::METHOD {
        let (id, params): (RequestId, DocumentSymbolParams) =
            req.extract(DocumentSymbolRequest::METHOD).unwrap();
        let uri = &params.text_document.uri;
        let symbols = if let Some(source) = documents.get(uri) {
            compute_symbols(source)
        } else {
            vec![]
        };
        let result = DocumentSymbolResponse::Nested(symbols);
        let resp = Response::new_ok(id, serde_json::to_value(result).unwrap());
        conn.sender.send(Message::Response(resp)).unwrap();
    } else if req.method == GotoDefinition::METHOD {
        let (id, params): (RequestId, GotoDefinitionParams) =
            req.extract(GotoDefinition::METHOD).unwrap();
        let uri = params
            .text_document_position_params
            .text_document
            .uri
            .clone();
        let pos = params.text_document_position_params.position;
        let result = if let Some(source) = documents.get(&uri) {
            handle_goto_definition(source, &uri, pos)
        } else {
            None
        };
        let resp = Response::new_ok(id, serde_json::to_value(result).unwrap());
        conn.sender.send(Message::Response(resp)).unwrap();
    } else if req.method == HoverRequest::METHOD {
        let (id, params): (RequestId, HoverParams) = req.extract(HoverRequest::METHOD).unwrap();
        let uri = &params.text_document_position_params.text_document.uri;
        let pos = params.text_document_position_params.position;
        let result = if let Some(source) = documents.get(uri) {
            handle_hover(source, pos)
        } else {
            None
        };
        let resp = Response::new_ok(id, serde_json::to_value(result).unwrap());
        conn.sender.send(Message::Response(resp)).unwrap();
    } else if req.method == Completion::METHOD {
        let (id, params): (RequestId, CompletionParams) = req.extract(Completion::METHOD).unwrap();
        let uri = &params.text_document_position.text_document.uri;
        let pos = params.text_document_position.position;
        let result = if let Some(source) = documents.get(uri) {
            handle_completion(source, pos)
        } else {
            CompletionResponse::List(CompletionList {
                is_incomplete: false,
                items: vec![],
            })
        };
        let resp = Response::new_ok(id, serde_json::to_value(result).unwrap());
        conn.sender.send(Message::Response(resp)).unwrap();
    } else if req.method == References::METHOD {
        let (id, params): (RequestId, ReferenceParams) =
            req.extract(References::METHOD).unwrap();
        let uri = params.text_document_position.text_document.uri.clone();
        let pos = params.text_document_position.position;
        let include_decl = params.context.include_declaration;
        let result = documents
            .get(&uri)
            .map(|source| handle_references(source, &uri, pos, include_decl))
            .unwrap_or_default();
        let resp = Response::new_ok(id, serde_json::to_value(result).unwrap());
        conn.sender.send(Message::Response(resp)).unwrap();
    } else if req.method == Rename::METHOD {
        let (id, params): (RequestId, RenameParams) = req.extract(Rename::METHOD).unwrap();
        let uri = params.text_document_position.text_document.uri.clone();
        let pos = params.text_document_position.position;
        let new_name = params.new_name;
        let resp = match documents.get(&uri) {
            Some(source) => match handle_rename(source, &uri, pos, &new_name) {
                Ok(edit) => Response::new_ok(id, serde_json::to_value(edit).unwrap()),
                Err(msg) => Response {
                    id,
                    result: None,
                    error: Some(ResponseError {
                        code: ErrorCode::InvalidParams as i32,
                        message: msg,
                        data: None,
                    }),
                },
            },
            None => Response::new_ok(id, serde_json::Value::Null),
        };
        conn.sender.send(Message::Response(resp)).unwrap();
    }
}

// ---- Go to Definition ----

fn handle_goto_definition(
    source: &str,
    uri: &Url,
    pos: Position,
) -> Option<GotoDefinitionResponse> {
    let word = word_at_position(source, pos.line, pos.character)?;
    let defs = collect_definitions(source);

    for def in &defs {
        if def.name == word {
            let loc = Location {
                uri: uri.clone(),
                range: Range {
                    start: Position::new(def.line, def.col),
                    end: Position::new(def.line, def.col + def.name.len() as u32),
                },
            };
            return Some(GotoDefinitionResponse::Scalar(loc));
        }
    }
    None
}

// ---- Hover ----

fn handle_hover(source: &str, pos: Position) -> Option<Hover> {
    let word = word_at_position(source, pos.line, pos.character)?;

    // Check builtins first
    for bi in BUILTINS {
        if bi.name == word {
            let content = format!("```ion\n{}\n```\n{}", bi.signature, bi.description);
            return Some(Hover {
                contents: HoverContents::Markup(MarkupContent {
                    kind: MarkupKind::Markdown,
                    value: content,
                }),
                range: None,
            });
        }
    }

    // Check user-defined functions/variables
    let defs = collect_definitions(source);
    for def in &defs {
        if def.name == word {
            let content = match def.kind {
                DefKind::Function => format!("```ion\n{}\n```", def.detail),
                DefKind::Variable => format!("```ion\n{}\n```", def.detail),
            };
            return Some(Hover {
                contents: HoverContents::Markup(MarkupContent {
                    kind: MarkupKind::Markdown,
                    value: content,
                }),
                range: None,
            });
        }
    }

    // Check keywords
    let keyword_doc = match word.as_str() {
        "let" => Some("Declare a variable. Use `mut` for mutable bindings.\n\n```ion\nlet x = 10;\nlet mut y = 0;\n```"),
        "fn" => Some("Declare a function.\n\n```ion\nfn add(a, b) { a + b }\n```"),
        "if" => Some("Conditional expression.\n\n```ion\nif x > 0 { \"positive\" } else { \"non-positive\" }\n```"),
        "match" => Some("Pattern matching expression.\n\n```ion\nmatch value {\n    Some(x) => x,\n    None => 0,\n}\n```"),
        "for" => Some("Iterate over a collection.\n\n```ion\nfor x in [1, 2, 3] { println(x); }\n```"),
        "while" => Some("Loop while condition is true.\n\n```ion\nwhile x < 10 { x += 1; }\n```"),
        "loop" => Some("Infinite loop. Use `break` to exit.\n\n```ion\nlet result = loop { if done { break 42; } };\n```"),
        "spawn" => Some("Spawn a concurrent task (requires `concurrency` feature).\n\n```ion\nlet t = spawn compute(100);\nlet result = t.await;\n```"),
        "async" => Some("Structured concurrency scope.\n\n```ion\nlet result = async {\n    let t = spawn work();\n    t.await\n};\n```"),
        "Some" => Some("`Option` variant containing a value.\n\n```ion\nSome(42)\n```"),
        "None" => Some("`Option` variant representing no value."),
        "Ok" => Some("`Result` variant representing success.\n\n```ion\nOk(42)\n```"),
        "Err" => Some("`Result` variant representing failure.\n\n```ion\nErr(\"something failed\")\n```"),
        "try" => Some("Begin a try/catch block.\n\n```ion\nlet result = try { risky() } catch e { fallback(e) };\n```"),
        "catch" => Some("Handle errors from a try block.\n\n```ion\ntry { risky() } catch e { println(e); }\n```"),
        "break" => Some("Exit a loop. Optionally return a value.\n\n```ion\nlet x = loop { break 42; };\n```"),
        "continue" => Some("Skip to the next iteration of a loop."),
        "return" => Some("Return a value from a function early.\n\n```ion\nfn check(x) { if x < 0 { return Err(\"negative\"); } Ok(x) }\n```"),
        "mut" => Some("Mark a binding as mutable.\n\n```ion\nlet mut count = 0;\ncount += 1;\n```"),
        "in" => Some("Used in `for` loops and membership tests.\n\n```ion\nfor x in [1, 2, 3] { println(x); }\n```"),
        "use" => Some("Import names from a module.\n\n```ion\nuse math::add;          // single import\nuse math::{add, PI};    // multiple imports\nuse math::*;            // glob import\n```"),
        _ => None,
    };

    keyword_doc.map(|doc| Hover {
        contents: HoverContents::Markup(MarkupContent {
            kind: MarkupKind::Markdown,
            value: doc.to_string(),
        }),
        range: None,
    })
}

// ---- Completion ----

fn handle_completion(source: &str, pos: Position) -> CompletionResponse {
    let mut items = Vec::new();

    // Check if we're after a dot (method completion)
    let line_text = source.lines().nth(pos.line as usize).unwrap_or("");
    let col = pos.character as usize;

    let is_dot_completion = col > 0 && line_text.as_bytes().get(col - 1) == Some(&b'.');

    // Check if we're after a `::` (module member completion)
    let module_prefix = {
        let before_cursor = &line_text[..col.min(line_text.len())];
        if let Some(idx) = before_cursor.rfind("::") {
            let prefix = before_cursor[..idx].trim();
            // Extract the last word as the module name
            let mod_name = prefix
                .rsplit(|c: char| c.is_whitespace() || c == '(' || c == ',' || c == '{')
                .next()
                .unwrap_or(prefix);
            if !mod_name.is_empty() {
                Some(mod_name.to_string())
            } else {
                None
            }
        } else {
            None
        }
    };

    // Check if we're in a type annotation position (after `:` in a let binding)
    let is_type_position = {
        let before_cursor = &line_text[..col.min(line_text.len())];
        let trimmed = before_cursor.trim_start();
        (trimmed.starts_with("let ") || trimmed.starts_with("let mut "))
            && before_cursor.contains(':')
            && !before_cursor.contains('=')
    };

    if let Some(ref mod_name) = module_prefix {
        let members: &[(&str, &str, CompletionItemKind)] = match mod_name.as_str() {
            "math" => &[
                (
                    "PI",
                    "Pi constant (3.14159...)",
                    CompletionItemKind::CONSTANT,
                ),
                (
                    "E",
                    "Euler's number (2.71828...)",
                    CompletionItemKind::CONSTANT,
                ),
                ("TAU", "Tau (2π)", CompletionItemKind::CONSTANT),
                ("INF", "Positive infinity", CompletionItemKind::CONSTANT),
                ("NAN", "Not-a-number", CompletionItemKind::CONSTANT),
                ("abs", "Absolute value", CompletionItemKind::FUNCTION),
                ("min", "Minimum of arguments", CompletionItemKind::FUNCTION),
                ("max", "Maximum of arguments", CompletionItemKind::FUNCTION),
                ("floor", "Floor (round down)", CompletionItemKind::FUNCTION),
                ("ceil", "Ceiling (round up)", CompletionItemKind::FUNCTION),
                ("round", "Round to nearest", CompletionItemKind::FUNCTION),
                ("sqrt", "Square root", CompletionItemKind::FUNCTION),
                ("pow", "Exponentiation", CompletionItemKind::FUNCTION),
                (
                    "clamp",
                    "Clamp value to range",
                    CompletionItemKind::FUNCTION,
                ),
                ("sin", "Sine", CompletionItemKind::FUNCTION),
                ("cos", "Cosine", CompletionItemKind::FUNCTION),
                ("tan", "Tangent", CompletionItemKind::FUNCTION),
                (
                    "atan2",
                    "Two-argument arctangent",
                    CompletionItemKind::FUNCTION,
                ),
                ("log", "Natural logarithm", CompletionItemKind::FUNCTION),
                ("log2", "Base-2 logarithm", CompletionItemKind::FUNCTION),
                ("log10", "Base-10 logarithm", CompletionItemKind::FUNCTION),
                ("is_nan", "Check if NaN", CompletionItemKind::FUNCTION),
                ("is_inf", "Check if infinite", CompletionItemKind::FUNCTION),
            ],
            "json" => &[
                (
                    "encode",
                    "Value to JSON string",
                    CompletionItemKind::FUNCTION,
                ),
                (
                    "decode",
                    "JSON string to value",
                    CompletionItemKind::FUNCTION,
                ),
                (
                    "pretty",
                    "Pretty-printed JSON string",
                    CompletionItemKind::FUNCTION,
                ),
            ],
            "io" => &[
                (
                    "print",
                    "Print without newline",
                    CompletionItemKind::FUNCTION,
                ),
                (
                    "println",
                    "Print with newline",
                    CompletionItemKind::FUNCTION,
                ),
                (
                    "eprintln",
                    "Print to stderr with newline",
                    CompletionItemKind::FUNCTION,
                ),
            ],
            "string" => &[(
                "join",
                "Join list elements into string with optional separator",
                CompletionItemKind::FUNCTION,
            )],
            _ => &[],
        };
        for (label, doc, kind) in members {
            items.push(CompletionItem {
                label: label.to_string(),
                kind: Some(*kind),
                documentation: Some(lsp_types::Documentation::String(doc.to_string())),
                ..Default::default()
            });
        }
        return CompletionResponse::List(CompletionList {
            is_incomplete: false,
            items,
        });
    } else if is_type_position {
        let type_names = [
            ("int", "Integer type"),
            ("float", "Floating-point type"),
            ("bool", "Boolean type"),
            ("string", "String type"),
            ("bytes", "Byte string type"),
            ("list", "List type (e.g. list<int>)"),
            ("dict", "Dictionary type (e.g. dict<string, int>)"),
            ("tuple", "Tuple type"),
            ("set", "Set type"),
            ("fn", "Function type"),
            ("cell", "Mutable reference cell type"),
            ("any", "Any type (accepts all values)"),
            ("Option", "Option type (e.g. Option<int>)"),
            ("Result", "Result type (e.g. Result<int, string>)"),
        ];
        for (name, doc) in type_names {
            items.push(CompletionItem {
                label: name.to_string(),
                kind: Some(CompletionItemKind::TYPE_PARAMETER),
                documentation: Some(lsp_types::Documentation::String(doc.to_string())),
                ..Default::default()
            });
        }
    } else if is_dot_completion {
        // Provide method completions
        let methods = [
            // String methods
            ("len", "len()", "Length", CompletionItemKind::METHOD),
            (
                "is_empty",
                "is_empty()",
                "True if empty",
                CompletionItemKind::METHOD,
            ),
            (
                "contains",
                "contains(sub)",
                "Contains substring/element",
                CompletionItemKind::METHOD,
            ),
            (
                "starts_with",
                "starts_with(prefix)",
                "Starts with prefix",
                CompletionItemKind::METHOD,
            ),
            (
                "ends_with",
                "ends_with(suffix)",
                "Ends with suffix",
                CompletionItemKind::METHOD,
            ),
            (
                "trim",
                "trim()",
                "Strip whitespace",
                CompletionItemKind::METHOD,
            ),
            (
                "split",
                "split(delim)",
                "Split by delimiter",
                CompletionItemKind::METHOD,
            ),
            (
                "replace",
                "replace(from, to)",
                "Replace occurrences",
                CompletionItemKind::METHOD,
            ),
            (
                "to_upper",
                "to_upper()",
                "Uppercase",
                CompletionItemKind::METHOD,
            ),
            (
                "to_lower",
                "to_lower()",
                "Lowercase",
                CompletionItemKind::METHOD,
            ),
            (
                "chars",
                "chars()",
                "List of characters",
                CompletionItemKind::METHOD,
            ),
            (
                "reverse",
                "reverse()",
                "Reversed copy",
                CompletionItemKind::METHOD,
            ),
            (
                "find",
                "find(sub)",
                "Index of first occurrence",
                CompletionItemKind::METHOD,
            ),
            (
                "repeat",
                "repeat(n)",
                "Repeat n times",
                CompletionItemKind::METHOD,
            ),
            (
                "to_int",
                "to_int()",
                "Parse as integer",
                CompletionItemKind::METHOD,
            ),
            (
                "to_float",
                "to_float()",
                "Parse as float",
                CompletionItemKind::METHOD,
            ),
            // List methods
            (
                "push",
                "push(val)",
                "Append value",
                CompletionItemKind::METHOD,
            ),
            (
                "pop",
                "pop()",
                "Remove last element",
                CompletionItemKind::METHOD,
            ),
            (
                "first",
                "first()",
                "First element (Option)",
                CompletionItemKind::METHOD,
            ),
            (
                "last",
                "last()",
                "Last element (Option)",
                CompletionItemKind::METHOD,
            ),
            ("sort", "sort()", "Sorted copy", CompletionItemKind::METHOD),
            (
                "flatten",
                "flatten()",
                "Flatten one level",
                CompletionItemKind::METHOD,
            ),
            (
                "join",
                "join(sep)",
                "Join with separator",
                CompletionItemKind::METHOD,
            ),
            (
                "enumerate",
                "enumerate()",
                "Index-value tuples",
                CompletionItemKind::METHOD,
            ),
            (
                "zip",
                "zip(other)",
                "Zip with another list",
                CompletionItemKind::METHOD,
            ),
            (
                "map",
                "map(fn)",
                "Apply function",
                CompletionItemKind::METHOD,
            ),
            (
                "filter",
                "filter(fn)",
                "Keep matching elements",
                CompletionItemKind::METHOD,
            ),
            (
                "fold",
                "fold(init, fn)",
                "Reduce with accumulator",
                CompletionItemKind::METHOD,
            ),
            (
                "flat_map",
                "flat_map(fn)",
                "Map then flatten",
                CompletionItemKind::METHOD,
            ),
            (
                "any",
                "any(fn)",
                "True if any match",
                CompletionItemKind::METHOD,
            ),
            (
                "all",
                "all(fn)",
                "True if all match",
                CompletionItemKind::METHOD,
            ),
            // Dict methods
            ("keys", "keys()", "List of keys", CompletionItemKind::METHOD),
            (
                "values",
                "values()",
                "List of values",
                CompletionItemKind::METHOD,
            ),
            (
                "entries",
                "entries()",
                "List of (key, value) tuples",
                CompletionItemKind::METHOD,
            ),
            (
                "contains_key",
                "contains_key(key)",
                "Key exists",
                CompletionItemKind::METHOD,
            ),
            (
                "get",
                "get(key)",
                "Get value by key",
                CompletionItemKind::METHOD,
            ),
            (
                "insert",
                "insert(key, val)",
                "Insert entry",
                CompletionItemKind::METHOD,
            ),
            (
                "remove",
                "remove(key)",
                "Remove entry",
                CompletionItemKind::METHOD,
            ),
            (
                "merge",
                "merge(other)",
                "Merge dicts",
                CompletionItemKind::METHOD,
            ),
            // New list/string/dict methods
            (
                "index",
                "index(val)",
                "Index of first occurrence",
                CompletionItemKind::METHOD,
            ),
            (
                "count",
                "count(val)",
                "Count occurrences",
                CompletionItemKind::METHOD,
            ),
            (
                "dedup",
                "dedup()",
                "Remove consecutive duplicates",
                CompletionItemKind::METHOD,
            ),
            (
                "unique",
                "unique()",
                "Remove all duplicates",
                CompletionItemKind::METHOD,
            ),
            (
                "sum",
                "sum()",
                "Sum of elements",
                CompletionItemKind::METHOD,
            ),
            (
                "window",
                "window(n)",
                "Sliding windows of size n",
                CompletionItemKind::METHOD,
            ),
            (
                "sort_by",
                "sort_by(fn)",
                "Sort by key function",
                CompletionItemKind::METHOD,
            ),
            (
                "to_string",
                "to_string()",
                "Convert to string",
                CompletionItemKind::METHOD,
            ),
            (
                "pad_start",
                "pad_start(width, char)",
                "Pad start to width",
                CompletionItemKind::METHOD,
            ),
            (
                "pad_end",
                "pad_end(width, char)",
                "Pad end to width",
                CompletionItemKind::METHOD,
            ),
            (
                "strip_prefix",
                "strip_prefix(prefix)",
                "Remove prefix if present",
                CompletionItemKind::METHOD,
            ),
            (
                "strip_suffix",
                "strip_suffix(suffix)",
                "Remove suffix if present",
                CompletionItemKind::METHOD,
            ),
            (
                "char_len",
                "char_len()",
                "Character count (Unicode-aware)",
                CompletionItemKind::METHOD,
            ),
            (
                "bytes",
                "bytes()",
                "Byte representation",
                CompletionItemKind::METHOD,
            ),
            (
                "update",
                "update(other)",
                "Merge dict (mutating)",
                CompletionItemKind::METHOD,
            ),
            (
                "keys_of",
                "keys_of(val)",
                "Keys with matching value",
                CompletionItemKind::METHOD,
            ),
            // Option/Result methods
            (
                "unwrap",
                "unwrap()",
                "Extract value or error",
                CompletionItemKind::METHOD,
            ),
            (
                "unwrap_or",
                "unwrap_or(default)",
                "Unwrap or return default",
                CompletionItemKind::METHOD,
            ),
            (
                "expect",
                "expect(msg)",
                "Unwrap or error",
                CompletionItemKind::METHOD,
            ),
            (
                "is_some",
                "is_some()",
                "True if Some",
                CompletionItemKind::METHOD,
            ),
            (
                "is_none",
                "is_none()",
                "True if None",
                CompletionItemKind::METHOD,
            ),
            ("is_ok", "is_ok()", "True if Ok", CompletionItemKind::METHOD),
            (
                "is_err",
                "is_err()",
                "True if Err",
                CompletionItemKind::METHOD,
            ),
            (
                "map_err",
                "map_err(fn)",
                "Transform error",
                CompletionItemKind::METHOD,
            ),
            (
                "and_then",
                "and_then(fn)",
                "Flat-map on Ok/Some",
                CompletionItemKind::METHOD,
            ),
            (
                "or_else",
                "or_else(fn)",
                "Call fn on Err/None",
                CompletionItemKind::METHOD,
            ),
            // Bytes methods
            (
                "to_list",
                "to_list()",
                "Convert to list",
                CompletionItemKind::METHOD,
            ),
            (
                "to_str",
                "to_str()",
                "Decode as UTF-8",
                CompletionItemKind::METHOD,
            ),
            (
                "to_hex",
                "to_hex()",
                "Hex-encoded string",
                CompletionItemKind::METHOD,
            ),
            (
                "slice",
                "slice(start, end)",
                "Sub-slice",
                CompletionItemKind::METHOD,
            ),
            // Task methods
            (
                "await",
                "await",
                "Wait for task",
                CompletionItemKind::METHOD,
            ),
            (
                "is_finished",
                "is_finished()",
                "Check if done",
                CompletionItemKind::METHOD,
            ),
            (
                "cancel",
                "cancel()",
                "Cancel task",
                CompletionItemKind::METHOD,
            ),
            (
                "is_cancelled",
                "is_cancelled()",
                "Check if cancelled",
                CompletionItemKind::METHOD,
            ),
            // Channel methods
            (
                "send",
                "send(val)",
                "Send value",
                CompletionItemKind::METHOD,
            ),
            (
                "recv",
                "recv()",
                "Receive value",
                CompletionItemKind::METHOD,
            ),
            (
                "try_recv",
                "try_recv()",
                "Non-blocking receive",
                CompletionItemKind::METHOD,
            ),
            (
                "close",
                "close()",
                "Close channel",
                CompletionItemKind::METHOD,
            ),
            // List methods: chunk, reduce, min, max
            (
                "chunk",
                "chunk(n)",
                "Split into chunks of size n",
                CompletionItemKind::METHOD,
            ),
            (
                "reduce",
                "reduce(fn)",
                "Reduce list with fn (no initial value)",
                CompletionItemKind::METHOD,
            ),
            (
                "min",
                "min()",
                "Minimum element",
                CompletionItemKind::METHOD,
            ),
            (
                "max",
                "max()",
                "Maximum element",
                CompletionItemKind::METHOD,
            ),
            // String methods: trim_start, trim_end
            (
                "trim_start",
                "trim_start()",
                "Trim leading whitespace",
                CompletionItemKind::METHOD,
            ),
            (
                "trim_end",
                "trim_end()",
                "Trim trailing whitespace",
                CompletionItemKind::METHOD,
            ),
            // Set methods
            (
                "add",
                "add(val)",
                "Add element to set",
                CompletionItemKind::METHOD,
            ),
            (
                "remove",
                "remove(val)",
                "Remove element from set",
                CompletionItemKind::METHOD,
            ),
            (
                "union",
                "union(other)",
                "Union of two sets",
                CompletionItemKind::METHOD,
            ),
            (
                "intersection",
                "intersection(other)",
                "Intersection of two sets",
                CompletionItemKind::METHOD,
            ),
            (
                "difference",
                "difference(other)",
                "Difference of two sets",
                CompletionItemKind::METHOD,
            ),
            // Cell methods
            (
                "set",
                "set(val)",
                "Set cell value",
                CompletionItemKind::METHOD,
            ),
        ];

        for (label, detail, doc, kind) in methods {
            items.push(CompletionItem {
                label: label.to_string(),
                kind: Some(kind),
                detail: Some(detail.to_string()),
                documentation: Some(lsp_types::Documentation::String(doc.to_string())),
                ..Default::default()
            });
        }
    } else {
        // Keywords
        for kw in KEYWORDS {
            items.push(CompletionItem {
                label: kw.to_string(),
                kind: Some(CompletionItemKind::KEYWORD),
                ..Default::default()
            });
        }

        // Builtins
        for bi in BUILTINS {
            items.push(CompletionItem {
                label: bi.name.to_string(),
                kind: Some(CompletionItemKind::FUNCTION),
                detail: Some(bi.signature.to_string()),
                documentation: Some(lsp_types::Documentation::String(bi.description.to_string())),
                ..Default::default()
            });
        }

        // User definitions
        let defs = collect_definitions(source);
        for def in &defs {
            let kind = match def.kind {
                DefKind::Function => CompletionItemKind::FUNCTION,
                DefKind::Variable => CompletionItemKind::VARIABLE,
            };
            items.push(CompletionItem {
                label: def.name.clone(),
                kind: Some(kind),
                detail: Some(def.detail.clone()),
                ..Default::default()
            });
        }
    }

    CompletionResponse::List(CompletionList {
        is_incomplete: false,
        items,
    })
}

// ---- References & Rename ----

/// Tokenize the source and return the location of every `Token::Ident` whose
/// text equals `name`. Range columns and lines are converted to LSP 0-based
/// coordinates. Tokens inside strings, comments, and other non-identifier
/// constructs are naturally excluded by the lexer.
fn find_identifier_occurrences(source: &str, name: &str) -> Vec<Range> {
    let tokens = match Lexer::new(source).tokenize() {
        Ok(t) => t,
        Err(_) => return Vec::new(),
    };
    let mut out = Vec::new();
    for tok in tokens {
        if let Token::Ident(text) = &tok.token {
            if text == name {
                let line = tok.line.saturating_sub(1) as u32;
                let col = tok.col.saturating_sub(1) as u32;
                out.push(Range {
                    start: Position::new(line, col),
                    end: Position::new(line, col + name.len() as u32),
                });
            }
        }
    }
    out
}

fn is_valid_identifier(name: &str) -> bool {
    let mut chars = name.chars();
    match chars.next() {
        Some(c) if c.is_ascii_alphabetic() || c == '_' => {}
        _ => return false,
    }
    chars.all(|c| c.is_ascii_alphanumeric() || c == '_')
}

fn handle_references(
    source: &str,
    uri: &Url,
    pos: Position,
    _include_decl: bool,
) -> Vec<Location> {
    // The LSP spec lets us omit the declaration when include_declaration is
    // false, but our textual scan gives us all occurrences uniformly and most
    // clients (VSCode, IntelliJ via LSP4IJ) expect the declaration in the
    // results regardless, so we always return everything.
    let Some(name) = word_at_position(source, pos.line, pos.character) else {
        return Vec::new();
    };
    if !is_valid_identifier(&name) {
        return Vec::new();
    }
    find_identifier_occurrences(source, &name)
        .into_iter()
        .map(|range| Location {
            uri: uri.clone(),
            range,
        })
        .collect()
}

fn handle_rename(
    source: &str,
    uri: &Url,
    pos: Position,
    new_name: &str,
) -> Result<WorkspaceEdit, String> {
    if !is_valid_identifier(new_name) {
        return Err(format!("'{new_name}' is not a valid Ion identifier"));
    }
    let Some(old_name) = word_at_position(source, pos.line, pos.character) else {
        return Err("no identifier under cursor".to_string());
    };
    if !is_valid_identifier(&old_name) {
        return Err(format!("'{old_name}' is not a renameable identifier"));
    }
    if new_name == old_name {
        return Ok(WorkspaceEdit::default());
    }
    let occurrences = find_identifier_occurrences(source, &old_name);
    if occurrences.is_empty() {
        return Err(format!("no occurrences of '{old_name}' found"));
    }
    let edits: Vec<TextEdit> = occurrences
        .into_iter()
        .map(|range| TextEdit {
            range,
            new_text: new_name.to_string(),
        })
        .collect();
    let mut changes = HashMap::new();
    changes.insert(uri.clone(), edits);
    Ok(WorkspaceEdit {
        changes: Some(changes),
        ..Default::default()
    })
}

// ---- Notification handling ----

fn handle_notification(
    conn: &Connection,
    documents: &mut HashMap<Url, String>,
    not: LspNotification,
) {
    match not.method.as_str() {
        DidOpenTextDocument::METHOD => {
            let params: DidOpenTextDocumentParams = serde_json::from_value(not.params).unwrap();
            let uri = params.text_document.uri.clone();
            let text = params.text_document.text.clone();
            documents.insert(uri.clone(), text.clone());
            publish_diagnostics(conn, uri, &text);
        }
        DidChangeTextDocument::METHOD => {
            let params: DidChangeTextDocumentParams = serde_json::from_value(not.params).unwrap();
            let uri = params.text_document.uri.clone();
            if let Some(change) = params.content_changes.into_iter().last() {
                documents.insert(uri.clone(), change.text.clone());
                publish_diagnostics(conn, uri, &change.text);
            }
        }
        DidSaveTextDocument::METHOD => {}
        _ => {}
    }
}

// ---- Diagnostics ----

fn publish_diagnostics(conn: &Connection, uri: Url, source: &str) {
    let diagnostics = compute_diagnostics(source);
    let params = lsp_types::PublishDiagnosticsParams {
        uri,
        diagnostics,
        version: None,
    };
    let not = LspNotification::new(
        notification::PublishDiagnostics::METHOD.to_string(),
        serde_json::to_value(params).unwrap(),
    );
    conn.sender.send(Message::Notification(not)).unwrap();
}

fn compute_diagnostics(source: &str) -> Vec<Diagnostic> {
    let mut diagnostics = Vec::new();

    let mut lexer = Lexer::new(source);
    let tokens = match lexer.tokenize() {
        Ok(tokens) => tokens,
        Err(err) => {
            add_error_diagnostic(&mut diagnostics, &err);
            return diagnostics;
        }
    };

    let mut parser = Parser::new(tokens);
    let output = parser.parse_program_recovering();
    for err in &output.errors {
        diagnostics.push(ion_error_to_diagnostic(err));
    }

    diagnostics
}

fn add_error_diagnostic(diagnostics: &mut Vec<Diagnostic>, err: &IonError) {
    diagnostics.push(ion_error_to_diagnostic(err));
    for extra in &err.additional {
        diagnostics.push(ion_error_to_diagnostic(extra));
    }
}

fn ion_error_to_diagnostic(err: &IonError) -> Diagnostic {
    let line = if err.line > 0 { err.line - 1 } else { 0 };
    let col = if err.col > 0 { err.col - 1 } else { 0 };

    Diagnostic {
        range: Range {
            start: Position::new(line as u32, col as u32),
            end: Position::new(line as u32, (col + 1) as u32),
        },
        severity: Some(DiagnosticSeverity::ERROR),
        source: Some("ion".to_string()),
        message: err.message.clone(),
        ..Default::default()
    }
}

// ---- Document Symbols ----

fn compute_symbols(source: &str) -> Vec<DocumentSymbol> {
    let defs = collect_definitions(source);
    defs.iter().map(def_to_symbol).collect()
}

#[allow(deprecated)]
fn def_to_symbol(def: &Definition) -> DocumentSymbol {
    let kind = match def.kind {
        DefKind::Function => SymbolKind::FUNCTION,
        DefKind::Variable => SymbolKind::VARIABLE,
    };
    DocumentSymbol {
        name: def.name.clone(),
        detail: Some(def.detail.clone()),
        kind,
        tags: None,
        deprecated: None,
        range: Range {
            start: Position::new(def.line, 0),
            end: Position::new(def.line, 0),
        },
        selection_range: Range {
            start: Position::new(def.line, def.col),
            end: Position::new(def.line, def.col + def.name.len() as u32),
        },
        children: None,
    }
}

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

    fn fake_uri() -> Url {
        Url::parse("file:///tmp/test.ion").unwrap()
    }

    #[test]
    fn references_finds_all_occurrences_excluding_strings() {
        let src = "let count = 0;\ncount += 1;\nlet msg = \"count me out\";\n";
        // cursor on `count` in line 0
        let refs = handle_references(src, &fake_uri(), Position::new(0, 6), true);
        assert_eq!(refs.len(), 2, "expected 2 hits, got: {:?}", refs);
        assert_eq!(refs[0].range.start, Position::new(0, 4));
        assert_eq!(refs[1].range.start, Position::new(1, 0));
    }

    #[test]
    fn references_skips_comments() {
        let src = "let x = 1;\n// x is a counter\nx + 1\n";
        let refs = handle_references(src, &fake_uri(), Position::new(0, 4), true);
        assert_eq!(refs.len(), 2);
        assert_eq!(refs[0].range.start, Position::new(0, 4));
        assert_eq!(refs[1].range.start, Position::new(2, 0));
    }

    #[test]
    fn references_returns_empty_for_keyword_or_string() {
        let src = "let x = \"hello\";\n";
        // cursor on `let` keyword: word_at_position returns "let", but lexer
        // tokenizes it as Token::Let, not Token::Ident, so no occurrences.
        let refs = handle_references(src, &fake_uri(), Position::new(0, 1), true);
        assert!(refs.is_empty());
    }

    #[test]
    fn rename_rewrites_every_occurrence() {
        let src = "let count = 0;\ncount += 1;\n";
        let edit = handle_rename(src, &fake_uri(), Position::new(0, 6), "tally").unwrap();
        let changes = edit.changes.expect("changes map");
        let edits = changes.get(&fake_uri()).expect("edits for uri");
        assert_eq!(edits.len(), 2);
        for e in edits {
            assert_eq!(e.new_text, "tally");
        }
    }

    #[test]
    fn rename_rejects_invalid_new_name() {
        let src = "let x = 1;\n";
        let err = handle_rename(src, &fake_uri(), Position::new(0, 4), "1abc").unwrap_err();
        assert!(err.contains("not a valid"));
    }

    #[test]
    fn rename_same_name_is_noop() {
        let src = "let x = 1;\nx + 1\n";
        let edit = handle_rename(src, &fake_uri(), Position::new(0, 4), "x").unwrap();
        assert!(edit.changes.is_none() || edit.changes.as_ref().unwrap().is_empty());
    }

    #[test]
    fn rename_errors_when_cursor_not_on_identifier() {
        let src = "let x = 1;\n";
        // column 14 is past end of line
        let err = handle_rename(src, &fake_uri(), Position::new(0, 14), "y").unwrap_err();
        assert!(err.contains("no identifier"));
    }

    #[test]
    fn is_valid_identifier_basics() {
        assert!(is_valid_identifier("foo"));
        assert!(is_valid_identifier("_bar"));
        assert!(is_valid_identifier("baz_99"));
        assert!(!is_valid_identifier(""));
        assert!(!is_valid_identifier("9foo"));
        assert!(!is_valid_identifier("foo bar"));
        assert!(!is_valid_identifier("foo-bar"));
    }
}
