solidity_language_server/

goto.rs

use serde_json::Value;
use std::collections::HashMap;
use tower_lsp::lsp_types::{Location, Position, Range, Url};
use tree_sitter::{Node, Parser};

use crate::types::{NodeId, SourceLoc};

#[derive(Debug, Clone)]
pub struct NodeInfo {
    pub src: String,
    pub name_location: Option<String>,
    pub name_locations: Vec<String>,
    pub referenced_declaration: Option<NodeId>,
    pub node_type: Option<String>,
    pub member_location: Option<String>,
    pub absolute_path: Option<String>,
}

/// All AST child keys to traverse (Solidity + Yul).
pub const CHILD_KEYS: &[&str] = &[
    "AST",
    "arguments",
    "baseContracts",
    "baseExpression",
    "baseName",
    "baseType",
    "block",
    "body",
    "components",
    "condition",
    "declarations",
    "endExpression",
    "errorCall",
    "eventCall",
    "expression",
    "externalCall",
    "falseBody",
    "falseExpression",
    "file",
    "foreign",
    "functionName",
    "indexExpression",
    "initialValue",
    "initializationExpression",
    "keyType",
    "leftExpression",
    "leftHandSide",
    "libraryName",
    "literals",
    "loopExpression",
    "members",
    "modifierName",
    "modifiers",
    "name",
    "names",
    "nodes",
    "options",
    "overrides",
    "parameters",
    "pathNode",
    "post",
    "pre",
    "returnParameters",
    "rightExpression",
    "rightHandSide",
    "startExpression",
    "statements",
    "storageLayout",
    "subExpression",
    "subdenomination",
    "symbolAliases",
    "trueBody",
    "trueExpression",
    "typeName",
    "unitAlias",
    "value",
    "valueType",
    "variableNames",
    "variables",
];

fn push_if_node_or_array<'a>(tree: &'a Value, key: &str, stack: &mut Vec<&'a Value>) {
    if let Some(value) = tree.get(key) {
        match value {
            Value::Array(arr) => {
                stack.extend(arr);
            }
            Value::Object(_) => {
                stack.push(value);
            }
            _ => {}
        }
    }
}

/// Maps `"offset:length:fileId"` src strings from Yul externalReferences
/// to the Solidity declaration node id they refer to.
pub type ExternalRefs = HashMap<String, NodeId>;

/// Pre-computed AST index. Built once when an AST enters the cache,
/// then reused on every goto/references/rename/hover request.
#[derive(Debug, Clone)]
pub struct CachedBuild {
    pub ast: Value,
    pub nodes: HashMap<String, HashMap<NodeId, NodeInfo>>,
    pub path_to_abs: HashMap<String, String>,
    pub external_refs: ExternalRefs,
    pub id_to_path_map: HashMap<String, String>,
    /// Pre-built gas index from contract output. Built once, reused by
    /// hover, inlay hints, and code lens.
    pub gas_index: crate::gas::GasIndex,
    /// Pre-built hint lookup per file. Built once, reused on every
    /// inlay hint request (avoids O(n²) declaration resolution per request).
    pub hint_index: crate::inlay_hints::HintIndex,
    /// Pre-built documentation index from solc userdoc/devdoc.
    /// Merged and keyed by selector for fast hover lookup.
    pub doc_index: crate::hover::DocIndex,
    /// The text_cache version this build was created from.
    /// Used to detect dirty files (unsaved edits since last build).
    pub build_version: i32,
}

impl CachedBuild {
    /// Build the index from normalized AST output.
    ///
    /// Canonical shape:
    /// - `sources[path] = { id, ast }`
    /// - `contracts[path][name] = { abi, evm, ... }`
    /// - `source_id_to_path = { "0": "path", ... }`
    pub fn new(ast: Value, build_version: i32) -> Self {
        let (nodes, path_to_abs, external_refs) = if let Some(sources) = ast.get("sources") {
            cache_ids(sources)
        } else {
            (HashMap::new(), HashMap::new(), HashMap::new())
        };

        let id_to_path_map = ast
            .get("source_id_to_path")
            .and_then(|v| v.as_object())
            .map(|obj| {
                obj.iter()
                    .map(|(k, v)| (k.clone(), v.as_str().unwrap_or("").to_string()))
                    .collect()
            })
            .unwrap_or_default();

        let gas_index = crate::gas::build_gas_index(&ast);

        let hint_index = if let Some(sources) = ast.get("sources") {
            crate::inlay_hints::build_hint_index(sources)
        } else {
            HashMap::new()
        };

        let doc_index = crate::hover::build_doc_index(&ast);

        Self {
            ast,
            nodes,
            path_to_abs,
            external_refs,
            id_to_path_map,
            gas_index,
            hint_index,
            doc_index,
            build_version,
        }
    }
}

type Type = (
    HashMap<String, HashMap<NodeId, NodeInfo>>,
    HashMap<String, String>,
    ExternalRefs,
);

pub fn cache_ids(sources: &Value) -> Type {
    let mut nodes: HashMap<String, HashMap<NodeId, NodeInfo>> = HashMap::new();
    let mut path_to_abs: HashMap<String, String> = HashMap::new();
    let mut external_refs: ExternalRefs = HashMap::new();

    if let Some(sources_obj) = sources.as_object() {
        for (path, source_data) in sources_obj {
            if let Some(ast) = source_data.get("ast") {
                // Get the absolute path for this file
                let abs_path = ast
                    .get("absolutePath")
                    .and_then(|v| v.as_str())
                    .unwrap_or(path)
                    .to_string();

                path_to_abs.insert(path.clone(), abs_path.clone());

                // Initialize the nodes map for this file
                if !nodes.contains_key(&abs_path) {
                    nodes.insert(abs_path.clone(), HashMap::new());
                }

                if let Some(id) = ast.get("id").and_then(|v| v.as_u64())
                    && let Some(src) = ast.get("src").and_then(|v| v.as_str())
                {
                    nodes.get_mut(&abs_path).unwrap().insert(
                        NodeId(id),
                        NodeInfo {
                            src: src.to_string(),
                            name_location: None,
                            name_locations: vec![],
                            referenced_declaration: None,
                            node_type: ast
                                .get("nodeType")
                                .and_then(|v| v.as_str())
                                .map(|s| s.to_string()),
                            member_location: None,
                            absolute_path: ast
                                .get("absolutePath")
                                .and_then(|v| v.as_str())
                                .map(|s| s.to_string()),
                        },
                    );
                }

                let mut stack = vec![ast];

                while let Some(tree) = stack.pop() {
                    if let Some(raw_id) = tree.get("id").and_then(|v| v.as_u64())
                        && let Some(src) = tree.get("src").and_then(|v| v.as_str())
                    {
                        let id = NodeId(raw_id);
                        // Check for nameLocation first
                        let mut name_location = tree
                            .get("nameLocation")
                            .and_then(|v| v.as_str())
                            .map(|s| s.to_string());

                        // Check for nameLocations array and use appropriate element
                        // For IdentifierPath (qualified names like D.State), use the last element (the actual identifier)
                        // For other nodes, use the first element
                        if name_location.is_none()
                            && let Some(name_locations) = tree.get("nameLocations")
                            && let Some(locations_array) = name_locations.as_array()
                            && !locations_array.is_empty()
                        {
                            let node_type = tree.get("nodeType").and_then(|v| v.as_str());
                            if node_type == Some("IdentifierPath") {
                                name_location = locations_array
                                    .last()
                                    .and_then(|v| v.as_str())
                                    .map(|s| s.to_string());
                            } else {
                                name_location = locations_array[0].as_str().map(|s| s.to_string());
                            }
                        }

                        let name_locations = if let Some(name_locations) = tree.get("nameLocations")
                            && let Some(locations_array) = name_locations.as_array()
                        {
                            locations_array
                                .iter()
                                .filter_map(|v| v.as_str().map(|s| s.to_string()))
                                .collect()
                        } else {
                            vec![]
                        };

                        let mut final_name_location = name_location;
                        if final_name_location.is_none()
                            && let Some(member_loc) =
                                tree.get("memberLocation").and_then(|v| v.as_str())
                        {
                            final_name_location = Some(member_loc.to_string());
                        }

                        let node_info = NodeInfo {
                            src: src.to_string(),
                            name_location: final_name_location,
                            name_locations,
                            referenced_declaration: tree
                                .get("referencedDeclaration")
                                .and_then(|v| v.as_u64())
                                .map(NodeId),
                            node_type: tree
                                .get("nodeType")
                                .and_then(|v| v.as_str())
                                .map(|s| s.to_string()),
                            member_location: tree
                                .get("memberLocation")
                                .and_then(|v| v.as_str())
                                .map(|s| s.to_string()),
                            absolute_path: tree
                                .get("absolutePath")
                                .and_then(|v| v.as_str())
                                .map(|s| s.to_string()),
                        };

                        nodes.get_mut(&abs_path).unwrap().insert(id, node_info);

                        // Collect externalReferences from InlineAssembly nodes
                        if tree.get("nodeType").and_then(|v| v.as_str()) == Some("InlineAssembly")
                            && let Some(ext_refs) =
                                tree.get("externalReferences").and_then(|v| v.as_array())
                        {
                            for ext_ref in ext_refs {
                                if let Some(src_str) = ext_ref.get("src").and_then(|v| v.as_str())
                                    && let Some(decl_id) =
                                        ext_ref.get("declaration").and_then(|v| v.as_u64())
                                {
                                    external_refs.insert(src_str.to_string(), NodeId(decl_id));
                                }
                            }
                        }
                    }

                    for key in CHILD_KEYS {
                        push_if_node_or_array(tree, key, &mut stack);
                    }
                }
            }
        }
    }

    (nodes, path_to_abs, external_refs)
}

pub fn pos_to_bytes(source_bytes: &[u8], position: Position) -> usize {
    let text = String::from_utf8_lossy(source_bytes);
    crate::utils::position_to_byte_offset(&text, position)
}

pub fn bytes_to_pos(source_bytes: &[u8], byte_offset: usize) -> Option<Position> {
    let text = String::from_utf8_lossy(source_bytes);
    let pos = crate::utils::byte_offset_to_position(&text, byte_offset);
    Some(pos)
}

/// Convert a `"offset:length:fileId"` src string to an LSP Location.
pub fn src_to_location(src: &str, id_to_path: &HashMap<String, String>) -> Option<Location> {
    let loc = SourceLoc::parse(src)?;
    let file_path = id_to_path.get(&loc.file_id_str())?;

    let absolute_path = if std::path::Path::new(file_path).is_absolute() {
        std::path::PathBuf::from(file_path)
    } else {
        std::env::current_dir().ok()?.join(file_path)
    };

    let source_bytes = std::fs::read(&absolute_path).ok()?;
    let start_pos = bytes_to_pos(&source_bytes, loc.offset)?;
    let end_pos = bytes_to_pos(&source_bytes, loc.end())?;
    let uri = Url::from_file_path(&absolute_path).ok()?;

    Some(Location {
        uri,
        range: Range {
            start: start_pos,
            end: end_pos,
        },
    })
}

pub fn goto_bytes(
    nodes: &HashMap<String, HashMap<NodeId, NodeInfo>>,
    path_to_abs: &HashMap<String, String>,
    id_to_path: &HashMap<String, String>,
    external_refs: &ExternalRefs,
    uri: &str,
    position: usize,
) -> Option<(String, usize, usize)> {
    let path = match uri.starts_with("file://") {
        true => &uri[7..],
        false => uri,
    };

    // Get absolute path for this file
    let abs_path = path_to_abs.get(path)?;

    // Get nodes for the current file only
    let current_file_nodes = nodes.get(abs_path)?;

    // Build reverse map: file_path -> file_id for filtering external refs by current file
    let path_to_file_id: HashMap<&str, &str> = id_to_path
        .iter()
        .map(|(id, p)| (p.as_str(), id.as_str()))
        .collect();

    // Determine the file id for the current file
    // path_to_abs maps filesystem path -> absolutePath (e.g. "src/libraries/SwapMath.sol")
    // id_to_path maps file_id -> relative path (e.g. "34" -> "src/libraries/SwapMath.sol")
    let current_file_id = path_to_file_id.get(abs_path.as_str());

    // Check if cursor is on a Yul external reference first
    for (src_str, decl_id) in external_refs {
        let Some(src_loc) = SourceLoc::parse(src_str) else {
            continue;
        };

        // Only consider external refs in the current file
        if let Some(file_id) = current_file_id {
            if src_loc.file_id_str() != *file_id {
                continue;
            }
        } else {
            continue;
        }

        if src_loc.offset <= position && position < src_loc.end() {
            // Found a Yul external reference — resolve to the declaration target
            let mut target_node: Option<&NodeInfo> = None;
            for file_nodes in nodes.values() {
                if let Some(node) = file_nodes.get(decl_id) {
                    target_node = Some(node);
                    break;
                }
            }
            let node = target_node?;
            let loc_str = node.name_location.as_deref().unwrap_or(&node.src);
            let loc = SourceLoc::parse(loc_str)?;
            let file_path = id_to_path.get(&loc.file_id_str())?.clone();
            return Some((file_path, loc.offset, loc.length));
        }
    }

    let mut refs = HashMap::new();

    // Only consider nodes from the current file that have references
    for (id, content) in current_file_nodes {
        if content.referenced_declaration.is_none() {
            continue;
        }

        let Some(src_loc) = SourceLoc::parse(&content.src) else {
            continue;
        };

        if src_loc.offset <= position && position < src_loc.end() {
            let diff = src_loc.length;
            if !refs.contains_key(&diff) || refs[&diff] <= *id {
                refs.insert(diff, *id);
            }
        }
    }

    if refs.is_empty() {
        // Check if we're on the string part of an import statement
        // ImportDirective nodes have absolutePath pointing to the imported file
        let tmp = current_file_nodes.iter();
        for (_id, content) in tmp {
            if content.node_type == Some("ImportDirective".to_string()) {
                let Some(src_loc) = SourceLoc::parse(&content.src) else {
                    continue;
                };

                if src_loc.offset <= position
                    && position < src_loc.end()
                    && let Some(import_path) = &content.absolute_path
                {
                    return Some((import_path.clone(), 0, 0));
                }
            }
        }
        return None;
    }

    // Find the reference with minimum diff (most specific)
    let min_diff = *refs.keys().min()?;
    let chosen_id = refs[&min_diff];
    let ref_id = current_file_nodes[&chosen_id].referenced_declaration?;

    // Search for the referenced declaration across all files
    let mut target_node: Option<&NodeInfo> = None;
    for file_nodes in nodes.values() {
        if let Some(node) = file_nodes.get(&ref_id) {
            target_node = Some(node);
            break;
        }
    }

    let node = target_node?;

    // Get location from nameLocation or src
    let loc_str = node.name_location.as_deref().unwrap_or(&node.src);
    let loc = SourceLoc::parse(loc_str)?;
    let file_path = id_to_path.get(&loc.file_id_str())?.clone();

    Some((file_path, loc.offset, loc.length))
}

pub fn goto_declaration(
    ast_data: &Value,
    file_uri: &Url,
    position: Position,
    source_bytes: &[u8],
) -> Option<Location> {
    let sources = ast_data.get("sources")?;
    let id_to_path = ast_data.get("source_id_to_path")?.as_object()?;

    let id_to_path_map: HashMap<String, String> = id_to_path
        .iter()
        .map(|(k, v)| (k.clone(), v.as_str().unwrap_or("").to_string()))
        .collect();

    let (nodes, path_to_abs, external_refs) = cache_ids(sources);
    let byte_position = pos_to_bytes(source_bytes, position);

    if let Some((file_path, location_bytes, length)) = goto_bytes(
        &nodes,
        &path_to_abs,
        &id_to_path_map,
        &external_refs,
        file_uri.as_ref(),
        byte_position,
    ) {
        let target_file_path = std::path::Path::new(&file_path);
        let absolute_path = if target_file_path.is_absolute() {
            target_file_path.to_path_buf()
        } else {
            std::env::current_dir().ok()?.join(target_file_path)
        };

        if let Ok(target_source_bytes) = std::fs::read(&absolute_path)
            && let Some(start_pos) = bytes_to_pos(&target_source_bytes, location_bytes)
            && let Some(end_pos) = bytes_to_pos(&target_source_bytes, location_bytes + length)
            && let Ok(target_uri) = Url::from_file_path(&absolute_path)
        {
            return Some(Location {
                uri: target_uri,
                range: Range {
                    start: start_pos,
                    end: end_pos,
                },
            });
        }
    };

    None
}

/// Name-based AST goto — resolves by searching cached AST nodes for identifiers
/// matching `name` in the current file, then following `referencedDeclaration`.
///
/// Unlike `goto_declaration` which matches by byte offset (breaks on dirty files),
/// this reads the identifier text from the **built source** (on disk) at each node's
/// `src` range and compares it to the cursor name. Works on dirty files because the
/// AST node relationships (referencedDeclaration) are still valid — only the byte
/// offsets in the current buffer are stale.
/// `byte_hint` is the cursor's byte offset in the dirty buffer, used to pick
/// the closest matching node when multiple nodes share the same name (overloads).
pub fn goto_declaration_by_name(
    cached_build: &CachedBuild,
    file_uri: &Url,
    name: &str,
    byte_hint: usize,
) -> Option<Location> {
    let path = match file_uri.as_ref().starts_with("file://") {
        true => &file_uri.as_ref()[7..],
        false => file_uri.as_ref(),
    };
    let abs_path = cached_build.path_to_abs.get(path)?;
    // Read the built source from disk to extract identifier text at src ranges
    let built_source = std::fs::read_to_string(abs_path).ok()?;

    // Collect all matching nodes: (distance_to_hint, span_size, ref_id)
    let mut candidates: Vec<(usize, usize, NodeId)> = Vec::new();

    let mut tmp = {
        let this = cached_build.nodes.get(abs_path)?;
        this.iter()
    };
    while let Some((_id, node)) = tmp.next() {
        let ref_id = match node.referenced_declaration {
            Some(id) => id,
            None => continue,
        };

        // Parse the node's src to get the byte range in the built source
        let Some(src_loc) = SourceLoc::parse(&node.src) else {
            continue;
        };
        let start = src_loc.offset;
        let length = src_loc.length;

        if start + length > built_source.len() {
            continue;
        }

        let node_text = &built_source[start..start + length];

        // Check if this node's text matches the name we're looking for.
        // For simple identifiers, the text equals the name directly.
        // For member access (e.g. `x.toInt128()`), check if the text contains
        // `.name(` or ends with `.name`.
        let matches = node_text == name
            || node_text.contains(&format!(".{name}("))
            || node_text.ends_with(&format!(".{name}"));

        if matches {
            // Distance from the byte_hint (cursor in dirty buffer) to the
            // node's src range. The closest node is most likely the one the
            // cursor is on, even if byte offsets shifted slightly.
            let distance = if byte_hint >= start && byte_hint < start + length {
                0 // cursor is inside this node's range
            } else if byte_hint < start {
                start - byte_hint
            } else {
                byte_hint - (start + length)
            };
            candidates.push((distance, length, ref_id));
        }
    }

    // Sort by distance (closest to cursor hint), then by span size (narrowest)
    candidates.sort_by_key(|&(dist, span, _)| (dist, span));
    let ref_id = candidates.first()?.2;

    // Find the declaration node across all files
    let mut target_node: Option<&NodeInfo> = None;
    for file_nodes in cached_build.nodes.values() {
        if let Some(node) = file_nodes.get(&ref_id) {
            target_node = Some(node);
            break;
        }
    }

    let node = target_node?;

    // Parse the target's nameLocation or src
    let loc_str = node.name_location.as_deref().unwrap_or(&node.src);
    let loc = SourceLoc::parse(loc_str)?;

    let file_path = cached_build.id_to_path_map.get(&loc.file_id_str())?;
    let location_bytes = loc.offset;
    let length = loc.length;

    let target_file_path = std::path::Path::new(file_path);
    let absolute_path = if target_file_path.is_absolute() {
        target_file_path.to_path_buf()
    } else {
        std::env::current_dir().ok()?.join(target_file_path)
    };

    let target_source_bytes = std::fs::read(&absolute_path).ok()?;
    let start_pos = bytes_to_pos(&target_source_bytes, location_bytes)?;
    let end_pos = bytes_to_pos(&target_source_bytes, location_bytes + length)?;
    let target_uri = Url::from_file_path(&absolute_path).ok()?;

    Some(Location {
        uri: target_uri,
        range: Range {
            start: start_pos,
            end: end_pos,
        },
    })
}

// ── Tree-sitter enhanced goto ──────────────────────────────────────────────

/// Context extracted from the cursor position via tree-sitter.
#[derive(Debug, Clone)]
pub struct CursorContext {
    /// The identifier text under the cursor.
    pub name: String,
    /// Enclosing function name (if any).
    pub function: Option<String>,
    /// Enclosing contract/interface/library name (if any).
    pub contract: Option<String>,
    /// Object in a member access expression (e.g. `SqrtPriceMath` in
    /// `SqrtPriceMath.getAmount0Delta`). Set when the cursor is on the
    /// property side of a dot expression.
    pub object: Option<String>,
    /// Number of arguments at the call site (for overload disambiguation).
    /// Set when the cursor is on a function name inside a `call_expression`.
    pub arg_count: Option<usize>,
    /// Inferred argument types at the call site (e.g. `["uint160", "uint160", "int128"]`).
    /// `None` entries mean the type couldn't be inferred for that argument.
    pub arg_types: Vec<Option<String>>,
}

/// Parse Solidity source with tree-sitter.
fn ts_parse(source: &str) -> Option<tree_sitter::Tree> {
    let mut parser = Parser::new();
    parser
        .set_language(&tree_sitter_solidity::LANGUAGE.into())
        .expect("failed to load Solidity grammar");
    parser.parse(source, None)
}

/// Validate that the text at a goto target location matches the expected name.
///
/// Used to reject tree-sitter results that land on the wrong identifier.
/// AST results are NOT validated because the AST can legitimately resolve
/// to a different name (e.g. `.selector` → error declaration).
pub fn validate_goto_target(target_source: &str, location: &Location, expected_name: &str) -> bool {
    let line = location.range.start.line as usize;
    let start_col = location.range.start.character as usize;
    let end_col = location.range.end.character as usize;

    if let Some(line_text) = target_source.lines().nth(line)
        && end_col <= line_text.len()
    {
        return &line_text[start_col..end_col] == expected_name;
    }
    // Can't read target — assume valid
    true
}

/// Find the deepest named node at the given byte offset.
fn ts_node_at_byte(node: Node, byte: usize) -> Option<Node> {
    if byte < node.start_byte() || byte >= node.end_byte() {
        return None;
    }
    let mut cursor = node.walk();
    for child in node.children(&mut cursor) {
        if child.start_byte() <= byte
            && byte < child.end_byte()
            && let Some(deeper) = ts_node_at_byte(child, byte)
        {
            return Some(deeper);
        }
    }
    Some(node)
}

/// Get the identifier name from a node (first `identifier` child or the node itself).
fn ts_child_id_text<'a>(node: Node<'a>, source: &'a str) -> Option<&'a str> {
    let mut cursor = node.walk();
    node.children(&mut cursor)
        .find(|c| c.kind() == "identifier" && c.is_named())
        .map(|c| &source[c.byte_range()])
}

/// Infer the type of an expression node using tree-sitter.
///
/// For identifiers, walks up to find the variable declaration and extracts its type.
/// For literals, infers the type from the literal kind.
/// For function calls, returns None (would need return type resolution).
fn infer_argument_type<'a>(arg_node: Node<'a>, source: &'a str) -> Option<String> {
    // Unwrap call_argument → get inner expression
    let expr = if arg_node.kind() == "call_argument" {
        let mut c = arg_node.walk();
        arg_node.children(&mut c).find(|ch| ch.is_named())?
    } else {
        arg_node
    };

    match expr.kind() {
        "identifier" => {
            let var_name = &source[expr.byte_range()];
            // Walk up scopes to find the variable declaration
            find_variable_type(expr, source, var_name)
        }
        "number_literal" | "decimal_number" | "hex_number" => Some("uint256".into()),
        "boolean_literal" => Some("bool".into()),
        "string_literal" | "hex_string_literal" => Some("string".into()),
        _ => None,
    }
}

/// Find the type of a variable by searching upward through enclosing scopes.
///
/// Looks for `parameter`, `variable_declaration`, and `state_variable_declaration`
/// nodes whose identifier matches the variable name.
fn find_variable_type(from: Node, source: &str, var_name: &str) -> Option<String> {
    let mut scope = from.parent();
    while let Some(node) = scope {
        match node.kind() {
            "function_definition" | "modifier_definition" | "constructor_definition" => {
                // Check parameters
                let mut c = node.walk();
                for child in node.children(&mut c) {
                    if child.kind() == "parameter"
                        && let Some(id) = ts_child_id_text(child, source)
                        && id == var_name
                    {
                        // Extract the type from this parameter
                        let mut pc = child.walk();
                        return child
                            .children(&mut pc)
                            .find(|c| {
                                matches!(
                                    c.kind(),
                                    "type_name"
                                        | "primitive_type"
                                        | "user_defined_type"
                                        | "mapping"
                                )
                            })
                            .map(|t| source[t.byte_range()].trim().to_string());
                    }
                }
            }
            "function_body" | "block_statement" | "unchecked_block" => {
                // Check local variable declarations
                let mut c = node.walk();
                for child in node.children(&mut c) {
                    if child.kind() == "variable_declaration_statement"
                        || child.kind() == "variable_declaration"
                    {
                        if let Some(id) = ts_child_id_text(child, source)
                            && id == var_name
                        {
                            let mut pc = child.walk();
                            return child
                                .children(&mut pc)
                                .find(|c| {
                                    matches!(
                                        c.kind(),
                                        "type_name"
                                            | "primitive_type"
                                            | "user_defined_type"
                                            | "mapping"
                                    )
                                })
                                .map(|t| source[t.byte_range()].trim().to_string());
                        }
                    }
                }
            }
            "contract_declaration" | "library_declaration" | "interface_declaration" => {
                // Check state variables
                if let Some(body) = ts_find_child(node, "contract_body") {
                    let mut c = body.walk();
                    for child in body.children(&mut c) {
                        if child.kind() == "state_variable_declaration"
                            && let Some(id) = ts_child_id_text(child, source)
                            && id == var_name
                        {
                            let mut pc = child.walk();
                            return child
                                .children(&mut pc)
                                .find(|c| {
                                    matches!(
                                        c.kind(),
                                        "type_name"
                                            | "primitive_type"
                                            | "user_defined_type"
                                            | "mapping"
                                    )
                                })
                                .map(|t| source[t.byte_range()].trim().to_string());
                        }
                    }
                }
            }
            _ => {}
        }
        scope = node.parent();
    }
    None
}

/// Infer argument types at a call site by examining each `call_argument` child.
fn infer_call_arg_types(call_node: Node, source: &str) -> Vec<Option<String>> {
    let mut cursor = call_node.walk();
    call_node
        .children(&mut cursor)
        .filter(|c| c.kind() == "call_argument")
        .map(|arg| infer_argument_type(arg, source))
        .collect()
}

/// Pick the best overload from multiple declarations based on argument types.
///
/// Strategy:
/// 1. If only one declaration, return it.
/// 2. Filter by argument count first.
/// 3. Among count-matched declarations, score by how many argument types match.
/// 4. Return the highest-scoring declaration.
fn best_overload<'a>(
    decls: &'a [TsDeclaration],
    arg_count: Option<usize>,
    arg_types: &[Option<String>],
) -> Option<&'a TsDeclaration> {
    if decls.len() == 1 {
        return decls.first();
    }
    if decls.is_empty() {
        return None;
    }

    // Filter to only function declarations (skip parameters, variables, etc.)
    let func_decls: Vec<&TsDeclaration> =
        decls.iter().filter(|d| d.param_count.is_some()).collect();

    if func_decls.is_empty() {
        return decls.first();
    }

    // If we have arg_count, filter by it
    let count_matched: Vec<&&TsDeclaration> = if let Some(ac) = arg_count {
        let matched: Vec<_> = func_decls
            .iter()
            .filter(|d| d.param_count == Some(ac))
            .collect();
        if matched.len() == 1 {
            return Some(matched[0]);
        }
        if matched.is_empty() {
            // No count match — fall back to all
            func_decls.iter().collect()
        } else {
            matched
        }
    } else {
        func_decls.iter().collect()
    };

    // Score each candidate by how many argument types match parameter types
    if !arg_types.is_empty() {
        let mut best: Option<(&TsDeclaration, usize)> = None;
        for &&decl in &count_matched {
            let score = arg_types
                .iter()
                .zip(decl.param_types.iter())
                .filter(|(arg_ty, param_ty)| {
                    if let Some(at) = arg_ty {
                        at == param_ty.as_str()
                    } else {
                        false
                    }
                })
                .count();
            if best.is_none() || score > best.unwrap().1 {
                best = Some((decl, score));
            }
        }
        if let Some((decl, _)) = best {
            return Some(decl);
        }
    }

    // Fallback: return first count-matched or first overall
    count_matched.first().map(|d| **d).or(decls.first())
}

/// Extract cursor context: the identifier under the cursor and its ancestor names.
///
/// Walks up the tree-sitter parse tree to find the enclosing function and contract.
pub fn cursor_context(source: &str, position: Position) -> Option<CursorContext> {
    let tree = ts_parse(source)?;
    let byte = pos_to_bytes(source.as_bytes(), position);
    let leaf = ts_node_at_byte(tree.root_node(), byte)?;

    // The leaf should be an identifier (or we find the nearest identifier)
    let id_node = if leaf.kind() == "identifier" {
        leaf
    } else {
        // Check parent — cursor might be just inside a node that contains an identifier
        let parent = leaf.parent()?;
        if parent.kind() == "identifier" {
            parent
        } else {
            return None;
        }
    };

    let name = source[id_node.byte_range()].to_string();
    let mut function = None;
    let mut contract = None;

    // Detect member access: if the identifier is the `property` side of a
    // member_expression (e.g. `SqrtPriceMath.getAmount0Delta`), extract
    // the object name so the caller can resolve cross-file.
    let object = id_node.parent().and_then(|parent| {
        if parent.kind() == "member_expression" {
            let prop = parent.child_by_field_name("property")?;
            // Only set object when cursor is on the property, not the object side
            if prop.id() == id_node.id() {
                let obj = parent.child_by_field_name("object")?;
                Some(source[obj.byte_range()].to_string())
            } else {
                None
            }
        } else {
            None
        }
    });

    // Count arguments and infer types at the call site for overload disambiguation.
    // Walk up from the identifier to find an enclosing `call_expression`,
    // then count its `call_argument` children and infer their types.
    let (arg_count, arg_types) = {
        let mut node = id_node.parent();
        let mut result = (None, vec![]);
        while let Some(n) = node {
            if n.kind() == "call_expression" {
                let types = infer_call_arg_types(n, source);
                result = (Some(types.len()), types);
                break;
            }
            node = n.parent();
        }
        result
    };

    // Walk ancestors
    let mut current = id_node.parent();
    while let Some(node) = current {
        match node.kind() {
            "function_definition" | "modifier_definition" if function.is_none() => {
                function = ts_child_id_text(node, source).map(String::from);
            }
            "constructor_definition" if function.is_none() => {
                function = Some("constructor".into());
            }
            "contract_declaration" | "interface_declaration" | "library_declaration"
                if contract.is_none() =>
            {
                contract = ts_child_id_text(node, source).map(String::from);
            }
            _ => {}
        }
        current = node.parent();
    }

    Some(CursorContext {
        name,
        function,
        contract,
        object,
        arg_count,
        arg_types,
    })
}

/// Information about a declaration found by tree-sitter.
#[derive(Debug, Clone)]
pub struct TsDeclaration {
    /// Position range of the declaration identifier.
    pub range: Range,
    /// What kind of declaration (contract, function, state_variable, etc.).
    pub kind: &'static str,
    /// Container name (contract/struct that owns this declaration).
    pub container: Option<String>,
    /// Number of parameters (for function/modifier declarations).
    pub param_count: Option<usize>,
    /// Parameter type signature (e.g. `["uint160", "uint160", "int128"]`).
    /// Used for overload disambiguation.
    pub param_types: Vec<String>,
}

/// Find all declarations of a name in a source file using tree-sitter.
///
/// Scans the parse tree for declaration nodes (state variables, functions, events,
/// errors, structs, enums, contracts, etc.) whose identifier matches `name`.
pub fn find_declarations_by_name(source: &str, name: &str) -> Vec<TsDeclaration> {
    let tree = match ts_parse(source) {
        Some(t) => t,
        None => return vec![],
    };
    let mut results = Vec::new();
    collect_declarations(tree.root_node(), source, name, None, &mut results);
    results
}

fn collect_declarations(
    node: Node,
    source: &str,
    name: &str,
    container: Option<&str>,
    out: &mut Vec<TsDeclaration>,
) {
    let mut cursor = node.walk();
    for child in node.children(&mut cursor) {
        if !child.is_named() {
            continue;
        }
        match child.kind() {
            "contract_declaration" | "interface_declaration" | "library_declaration" => {
                if let Some(id_name) = ts_child_id_text(child, source) {
                    if id_name == name {
                        out.push(TsDeclaration {
                            range: id_range(child),
                            kind: child.kind(),
                            container: container.map(String::from),
                            param_count: None,
                            param_types: vec![],
                        });
                    }
                    // Recurse into contract body
                    if let Some(body) = ts_find_child(child, "contract_body") {
                        collect_declarations(body, source, name, Some(id_name), out);
                    }
                }
            }
            "function_definition" | "modifier_definition" => {
                if let Some(id_name) = ts_child_id_text(child, source) {
                    if id_name == name {
                        let types = parameter_type_signature(child, source);
                        out.push(TsDeclaration {
                            range: id_range(child),
                            kind: child.kind(),
                            container: container.map(String::from),
                            param_count: Some(types.len()),
                            param_types: types.into_iter().map(String::from).collect(),
                        });
                    }
                    // Check function parameters
                    collect_parameters(child, source, name, container, out);
                    // Recurse into function body for local variables
                    if let Some(body) = ts_find_child(child, "function_body") {
                        collect_declarations(body, source, name, container, out);
                    }
                }
            }
            "constructor_definition" => {
                if name == "constructor" {
                    let types = parameter_type_signature(child, source);
                    out.push(TsDeclaration {
                        range: ts_range(child),
                        kind: "constructor_definition",
                        container: container.map(String::from),
                        param_count: Some(types.len()),
                        param_types: types.into_iter().map(String::from).collect(),
                    });
                }
                // Check constructor parameters
                collect_parameters(child, source, name, container, out);
                if let Some(body) = ts_find_child(child, "function_body") {
                    collect_declarations(body, source, name, container, out);
                }
            }
            "state_variable_declaration" | "variable_declaration" => {
                if let Some(id_name) = ts_child_id_text(child, source)
                    && id_name == name
                {
                    out.push(TsDeclaration {
                        range: id_range(child),
                        kind: child.kind(),
                        container: container.map(String::from),
                        param_count: None,
                        param_types: vec![],
                    });
                }
            }
            "struct_declaration" => {
                if let Some(id_name) = ts_child_id_text(child, source) {
                    if id_name == name {
                        out.push(TsDeclaration {
                            range: id_range(child),
                            kind: "struct_declaration",
                            container: container.map(String::from),
                            param_count: None,
                            param_types: vec![],
                        });
                    }
                    if let Some(body) = ts_find_child(child, "struct_body") {
                        collect_declarations(body, source, name, Some(id_name), out);
                    }
                }
            }
            "enum_declaration" => {
                if let Some(id_name) = ts_child_id_text(child, source) {
                    if id_name == name {
                        out.push(TsDeclaration {
                            range: id_range(child),
                            kind: "enum_declaration",
                            container: container.map(String::from),
                            param_count: None,
                            param_types: vec![],
                        });
                    }
                    // Check enum values
                    if let Some(body) = ts_find_child(child, "enum_body") {
                        let mut ecur = body.walk();
                        for val in body.children(&mut ecur) {
                            if val.kind() == "enum_value" && &source[val.byte_range()] == name {
                                out.push(TsDeclaration {
                                    range: ts_range(val),
                                    kind: "enum_value",
                                    container: Some(id_name.to_string()),
                                    param_count: None,
                                    param_types: vec![],
                                });
                            }
                        }
                    }
                }
            }
            "event_definition" | "error_declaration" => {
                if let Some(id_name) = ts_child_id_text(child, source)
                    && id_name == name
                {
                    out.push(TsDeclaration {
                        range: id_range(child),
                        kind: child.kind(),
                        container: container.map(String::from),
                        param_count: None,
                        param_types: vec![],
                    });
                }
            }
            "user_defined_type_definition" => {
                if let Some(id_name) = ts_child_id_text(child, source)
                    && id_name == name
                {
                    out.push(TsDeclaration {
                        range: id_range(child),
                        kind: "user_defined_type_definition",
                        container: container.map(String::from),
                        param_count: None,
                        param_types: vec![],
                    });
                }
            }
            // Recurse into blocks, if-else, loops, etc.
            _ => {
                collect_declarations(child, source, name, container, out);
            }
        }
    }
}

/// Extract the type signature from a function's parameters.
///
/// Returns a list of type strings, e.g. `["uint160", "uint160", "int128"]`.
/// For complex types (mappings, arrays, user-defined), returns the full
/// text of the type node.
fn parameter_type_signature<'a>(node: Node<'a>, source: &'a str) -> Vec<&'a str> {
    let mut cursor = node.walk();
    node.children(&mut cursor)
        .filter(|c| c.kind() == "parameter")
        .filter_map(|param| {
            let mut pc = param.walk();
            param
                .children(&mut pc)
                .find(|c| {
                    matches!(
                        c.kind(),
                        "type_name" | "primitive_type" | "user_defined_type" | "mapping"
                    )
                })
                .map(|t| source[t.byte_range()].trim())
        })
        .collect()
}

/// Collect parameter declarations from a function/constructor node.
fn collect_parameters(
    node: Node,
    source: &str,
    name: &str,
    container: Option<&str>,
    out: &mut Vec<TsDeclaration>,
) {
    let mut cursor = node.walk();
    for child in node.children(&mut cursor) {
        if child.kind() == "parameter"
            && let Some(id_name) = ts_child_id_text(child, source)
            && id_name == name
        {
            out.push(TsDeclaration {
                range: id_range(child),
                kind: "parameter",
                container: container.map(String::from),
                param_count: None,
                param_types: vec![],
            });
        }
    }
}

/// Tree-sitter range helper.
fn ts_range(node: Node) -> Range {
    let s = node.start_position();
    let e = node.end_position();
    Range {
        start: Position::new(s.row as u32, s.column as u32),
        end: Position::new(e.row as u32, e.column as u32),
    }
}

/// Get the range of the identifier child within a declaration node.
fn id_range(node: Node) -> Range {
    let mut cursor = node.walk();
    node.children(&mut cursor)
        .find(|c| c.kind() == "identifier" && c.is_named())
        .map(|c| ts_range(c))
        .unwrap_or_else(|| ts_range(node))
}

fn ts_find_child<'a>(node: Node<'a>, kind: &str) -> Option<Node<'a>> {
    let mut cursor = node.walk();
    node.children(&mut cursor).find(|c| c.kind() == kind)
}

/// Tree-sitter enhanced goto definition.
///
/// Uses tree-sitter to find the identifier under the cursor and its scope,
/// then resolves via the CompletionCache (for cross-file/semantic resolution),
/// and finally uses tree-sitter to find the declaration position in the target file.
///
/// Falls back to None if resolution fails — caller should try the existing AST-based path.
pub fn goto_definition_ts(
    source: &str,
    position: Position,
    file_uri: &Url,
    completion_cache: &crate::completion::CompletionCache,
    text_cache: &HashMap<String, (i32, String)>,
) -> Option<Location> {
    let ctx = cursor_context(source, position)?;

    // Member access: cursor is on `getAmount0Delta` in `SqrtPriceMath.getAmount0Delta`.
    // Look up the object (SqrtPriceMath) in the completion cache to find its file,
    // then search that file for the member declaration.
    // When multiple overloads exist, disambiguate by argument count and types.
    if let Some(obj_name) = &ctx.object {
        if let Some(path) = find_file_for_contract(completion_cache, obj_name, file_uri) {
            let target_source = read_target_source(&path, text_cache)?;
            let target_uri = Url::from_file_path(&path).ok()?;
            let decls = find_declarations_by_name(&target_source, &ctx.name);
            if let Some(d) = best_overload(&decls, ctx.arg_count, &ctx.arg_types) {
                return Some(Location {
                    uri: target_uri,
                    range: d.range,
                });
            }
        }
        // Object might be in the same file (e.g. a struct or contract in this file)
        let decls = find_declarations_by_name(source, &ctx.name);
        if let Some(d) = best_overload(&decls, ctx.arg_count, &ctx.arg_types) {
            return Some(Location {
                uri: file_uri.clone(),
                range: d.range,
            });
        }
    }

    // Step 1: Try to resolve via CompletionCache to find which file + name the declaration is in.
    // Use the scope chain by names: find the contract scope, then resolve the name.
    let resolved = resolve_via_cache(&ctx, file_uri, completion_cache);

    match resolved {
        Some(ResolvedTarget::SameFile) => {
            // Declaration is in the same file — find it with tree-sitter
            find_best_declaration(source, &ctx, file_uri)
        }
        Some(ResolvedTarget::OtherFile { path, name }) => {
            // Declaration is in another file — read target source and find by name
            let target_source = read_target_source(&path, text_cache);
            let target_source = target_source?;
            let target_uri = Url::from_file_path(&path).ok()?;
            let decls = find_declarations_by_name(&target_source, &name);
            decls.first().map(|d| Location {
                uri: target_uri,
                range: d.range,
            })
        }
        None => {
            // CompletionCache couldn't resolve — try same-file tree-sitter lookup as fallback
            find_best_declaration(source, &ctx, file_uri)
        }
    }
}

#[derive(Debug)]
enum ResolvedTarget {
    /// Declaration is in the same file as the usage.
    SameFile,
    /// Declaration is in a different file.
    OtherFile { path: String, name: String },
}

/// Try to resolve an identifier using the CompletionCache.
///
/// Finds the scope by matching ancestor names (contract, function) against
/// the cache's scope data, then resolves the name to a type and traces
/// back to the declaring file.
fn resolve_via_cache(
    ctx: &CursorContext,
    file_uri: &Url,
    cache: &crate::completion::CompletionCache,
) -> Option<ResolvedTarget> {
    // Find the contract scope node_id by name
    let contract_scope = ctx
        .contract
        .as_ref()
        .and_then(|name| cache.name_to_node_id.get(name.as_str()))
        .copied();

    // Try scope-based resolution: look in the contract's scope_declarations
    if let Some(contract_id) = contract_scope {
        // Check function scope if we're inside one
        if let Some(func_name) = &ctx.function {
            // Find the function scope: look for a scope whose parent is this contract
            // and which has a declaration for this function name
            if let Some(func_scope_id) = find_function_scope(cache, contract_id, func_name) {
                // Check declarations in this function scope first
                if let Some(decls) = cache.scope_declarations.get(&func_scope_id)
                    && decls.iter().any(|d| d.name == ctx.name)
                {
                    return Some(ResolvedTarget::SameFile);
                }
            }
        }

        // Check contract scope declarations (state variables, functions)
        if let Some(decls) = cache.scope_declarations.get(&contract_id)
            && decls.iter().any(|d| d.name == ctx.name)
        {
            return Some(ResolvedTarget::SameFile);
        }

        // Check inherited contracts (C3 linearization)
        if let Some(bases) = cache.linearized_base_contracts.get(&contract_id) {
            for &base_id in bases.iter().skip(1) {
                if let Some(decls) = cache.scope_declarations.get(&base_id)
                    && decls.iter().any(|d| d.name == ctx.name)
                {
                    // Found in a base contract — find which file it's in
                    // Reverse lookup: base_id → contract name → file
                    let base_name = cache
                        .name_to_node_id
                        .iter()
                        .find(|&(_, &id)| id == base_id)
                        .map(|(name, _)| name.clone());

                    if let Some(base_name) = base_name
                        && let Some(path) = find_file_for_contract(cache, &base_name, file_uri)
                    {
                        return Some(ResolvedTarget::OtherFile {
                            path,
                            name: ctx.name.clone(),
                        });
                    }
                    // Base contract might be in the same file
                    return Some(ResolvedTarget::SameFile);
                }
            }
        }
    }

    // Check if the name is a contract/library/interface name
    if cache.name_to_node_id.contains_key(&ctx.name) {
        // Could be same file or different file — check if it's in the current file
        if let Some(path) = find_file_for_contract(cache, &ctx.name, file_uri) {
            let current_path = file_uri.to_file_path().ok()?;
            let current_str = current_path.to_str()?;
            if path == current_str || path.ends_with(current_str) || current_str.ends_with(&path) {
                return Some(ResolvedTarget::SameFile);
            }
            return Some(ResolvedTarget::OtherFile {
                path,
                name: ctx.name.clone(),
            });
        }
        return Some(ResolvedTarget::SameFile);
    }

    // Flat fallback — name_to_type knows about it but we can't determine the file
    if cache.name_to_type.contains_key(&ctx.name) {
        return Some(ResolvedTarget::SameFile);
    }

    None
}

/// Find the scope node_id for a function within a contract.
fn find_function_scope(
    cache: &crate::completion::CompletionCache,
    contract_id: NodeId,
    func_name: &str,
) -> Option<NodeId> {
    // Look for a scope whose parent is the contract and which is a function scope.
    // The function name should appear as a declaration in the contract scope,
    // and the function's own scope is the one whose parent is the contract.
    for (&scope_id, &parent_id) in &cache.scope_parent {
        if parent_id == contract_id {
            // This scope's parent is our contract — it might be a function scope.
            // Check if this scope has declarations (functions/blocks do).
            // We also check if the contract declares a function with this name.
            if let Some(contract_decls) = cache.scope_declarations.get(&contract_id)
                && contract_decls.iter().any(|d| d.name == func_name)
            {
                // Found a child scope of the contract — could be the function.
                // Check if this scope_id has child scopes or declarations
                // that match what we'd expect for a function body.
                if cache.scope_declarations.contains_key(&scope_id)
                    || cache.scope_parent.values().any(|&p| p == scope_id)
                {
                    return Some(scope_id);
                }
            }
        }
    }
    None
}

/// Find the file path for a contract by searching the CompletionCache's path_to_file_id.
fn find_file_for_contract(
    cache: &crate::completion::CompletionCache,
    contract_name: &str,
    _file_uri: &Url,
) -> Option<String> {
    // The completion cache doesn't directly map contract → file.
    // But scope_ranges + path_to_file_id can help.
    // For now, check if the contract's node_id appears in any scope_range,
    // then map file_id back to path.
    let node_id = cache.name_to_node_id.get(contract_name)?;
    let scope_range = cache.scope_ranges.iter().find(|r| r.node_id == *node_id)?;
    let file_id = scope_range.file_id;

    // Reverse lookup: file_id → path
    cache
        .path_to_file_id
        .iter()
        .find(|&(_, &fid)| fid == file_id)
        .map(|(path, _)| path.clone())
}

/// Read source for a target file — prefer text_cache (open buffers), fallback to disk.
fn read_target_source(path: &str, text_cache: &HashMap<String, (i32, String)>) -> Option<String> {
    // Try text_cache by URI
    let uri = Url::from_file_path(path).ok()?;
    if let Some((_, content)) = text_cache.get(&uri.to_string()) {
        return Some(content.clone());
    }
    // Fallback to disk
    std::fs::read_to_string(path).ok()
}

/// Find the best matching declaration in the same file.
fn find_best_declaration(source: &str, ctx: &CursorContext, file_uri: &Url) -> Option<Location> {
    let decls = find_declarations_by_name(source, &ctx.name);
    if decls.is_empty() {
        return None;
    }

    // If there's only one declaration, use it
    if decls.len() == 1 {
        return Some(Location {
            uri: file_uri.clone(),
            range: decls[0].range,
        });
    }

    // Multiple declarations — prefer the one in the same contract
    if let Some(contract_name) = &ctx.contract
        && let Some(d) = decls
            .iter()
            .find(|d| d.container.as_deref() == Some(contract_name))
    {
        return Some(Location {
            uri: file_uri.clone(),
            range: d.range,
        });
    }

    // Fallback: return first declaration
    Some(Location {
        uri: file_uri.clone(),
        range: decls[0].range,
    })
}

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

    #[test]
    fn test_cursor_context_state_var() {
        let source = r#"
contract Token {
    uint256 public totalSupply;
    function mint(uint256 amount) public {
        totalSupply += amount;
    }
}
"#;
        // Cursor on `totalSupply` inside mint (line 4, col 8)
        let ctx = cursor_context(source, Position::new(4, 8)).unwrap();
        assert_eq!(ctx.name, "totalSupply");
        assert_eq!(ctx.function.as_deref(), Some("mint"));
        assert_eq!(ctx.contract.as_deref(), Some("Token"));
    }

    #[test]
    fn test_cursor_context_top_level() {
        let source = r#"
contract Foo {}
contract Bar {}
"#;
        // Cursor on `Foo` (line 1, col 9) — the identifier of the contract declaration
        let ctx = cursor_context(source, Position::new(1, 9)).unwrap();
        assert_eq!(ctx.name, "Foo");
        assert!(ctx.function.is_none());
        // The identifier `Foo` is a child of contract_declaration, so contract is set
        assert_eq!(ctx.contract.as_deref(), Some("Foo"));
    }

    #[test]
    fn test_find_declarations() {
        let source = r#"
contract Token {
    uint256 public totalSupply;
    function mint(uint256 amount) public {
        totalSupply += amount;
    }
}
"#;
        let decls = find_declarations_by_name(source, "totalSupply");
        assert_eq!(decls.len(), 1);
        assert_eq!(decls[0].kind, "state_variable_declaration");
        assert_eq!(decls[0].container.as_deref(), Some("Token"));
    }

    #[test]
    fn test_find_declarations_multiple_contracts() {
        let source = r#"
contract A {
    uint256 public value;
}
contract B {
    uint256 public value;
}
"#;
        let decls = find_declarations_by_name(source, "value");
        assert_eq!(decls.len(), 2);
        assert_eq!(decls[0].container.as_deref(), Some("A"));
        assert_eq!(decls[1].container.as_deref(), Some("B"));
    }

    #[test]
    fn test_find_declarations_enum_value() {
        let source = "contract Foo { enum Status { Active, Paused } }";
        let decls = find_declarations_by_name(source, "Active");
        assert_eq!(decls.len(), 1);
        assert_eq!(decls[0].kind, "enum_value");
        assert_eq!(decls[0].container.as_deref(), Some("Status"));
    }

    #[test]
    fn test_cursor_context_short_param() {
        let source = r#"
contract Shop {
    uint256 public TAX;
    constructor(uint256 price, uint16 tax, uint16 taxBase) {
        TAX = tax;
    }
}
"#;
        // Cursor on `tax` usage at line 4, col 14 (TAX = tax;)
        let ctx = cursor_context(source, Position::new(4, 14)).unwrap();
        assert_eq!(ctx.name, "tax");
        assert_eq!(ctx.contract.as_deref(), Some("Shop"));

        // Cursor on `TAX` at line 4, col 8
        let ctx2 = cursor_context(source, Position::new(4, 8)).unwrap();
        assert_eq!(ctx2.name, "TAX");

        // Parameters are found as declarations
        let decls = find_declarations_by_name(source, "tax");
        assert_eq!(decls.len(), 1);
        assert_eq!(decls[0].kind, "parameter");

        let decls_tax_base = find_declarations_by_name(source, "taxBase");
        assert_eq!(decls_tax_base.len(), 1);
        assert_eq!(decls_tax_base[0].kind, "parameter");

        let decls_price = find_declarations_by_name(source, "price");
        assert_eq!(decls_price.len(), 1);
        assert_eq!(decls_price[0].kind, "parameter");

        // State variable is also found
        let decls_tax_upper = find_declarations_by_name(source, "TAX");
        assert_eq!(decls_tax_upper.len(), 1);
        assert_eq!(decls_tax_upper[0].kind, "state_variable_declaration");
    }

    #[test]
    fn test_find_best_declaration_same_contract() {
        let source = r#"
contract A { uint256 public x; }
contract B { uint256 public x; }
"#;
        let ctx = CursorContext {
            name: "x".into(),
            function: None,
            contract: Some("B".into()),
            object: None,
            arg_count: None,
            arg_types: vec![],
        };
        let uri = Url::parse("file:///test.sol").unwrap();
        let loc = find_best_declaration(source, &ctx, &uri).unwrap();
        // Should pick B's x (line 2), not A's x (line 1)
        assert_eq!(loc.range.start.line, 2);
    }
}