query_lib.lua

-- Query execution for code graph
-- Implements: usages, deps, callers, callees, hierarchy, impacted
-- Entry kind: library.lua
-- Note: Stage 1 provides stub structure. Queries filled in Stages 2-5.

local graph = require("graph")

-- Forward declarations for recursive helpers
local build_import_subtree
local build_imported_by_subtree
local build_caller_tree
local build_callee_tree
local build_ancestor_chain
local build_descendant_tree

---------------------------------------------------------------------------
-- Find usages (Stage 3)
---------------------------------------------------------------------------

--- Find all references to a symbol across the project
-- @param state IndexState
-- @param symbol_name string  Short name or FQN
-- @param ref_kind string|nil  Filter by reference kind
-- @return {Reference}, Symbol|nil
local function find_usages(state, symbol_name, ref_kind)
    local symbols = graph.find_symbol(state, symbol_name)
    if #symbols == 0 then
        return {}, nil
    end

    -- Collect references for all matching symbols
    local all_refs = {}
    local primary = symbols[1]

    for _, sym in ipairs(symbols) do
        local refs = graph.find_refs_to(state, sym.id)
        for _, ref in ipairs(refs) do
            if not ref_kind or ref_kind == "all" or ref.kind == ref_kind then
                table.insert(all_refs, ref)
            end
        end
    end

    -- Sort by file then line
    table.sort(all_refs, function(a, b)
        if a.from_file ~= b.from_file then
            return a.from_file < b.from_file
        end
        return a.line < b.line
    end)

    return all_refs, primary
end

---------------------------------------------------------------------------
-- File dependencies (Stage 2)
---------------------------------------------------------------------------

--- Build import dependency tree for a file
-- @param state IndexState
-- @param file_path string  File path or symbol name
-- @param direction string  "imports" | "imported_by" | "both"
-- @param max_depth integer
-- @return table tree_node  {name, children, annotation?}
local function build_dep_tree(state, file_path, direction, max_depth)
    max_depth = max_depth or 3
    direction = direction or "both"

    -- If given a symbol name, find its file
    if not state.files[file_path] then
        local syms = graph.find_symbol(state, file_path)
        if #syms > 0 and syms[1].file then
            file_path = syms[1].file
        else
            return {name = file_path, children = {}, annotation = "not found"}
        end
    end

    local result = {imports = nil, imported_by = nil}

    if direction == "imports" or direction == "both" then
        result.imports = build_import_subtree(state, file_path, max_depth, 0, {})
    end

    if direction == "imported_by" or direction == "both" then
        result.imported_by = build_imported_by_subtree(state, file_path, max_depth, 0, {})
    end

    return result
end

--- Build the "imports" direction of the dependency tree
function build_import_subtree(state, file_path, max_depth, depth, seen)
    local info = state.files[file_path]
    if not info then
        return {name = file_path, children = {}}
    end

    local node = {
        name = info.namespace or file_path,
        children = {},
    }

    if depth >= max_depth or seen[file_path] then
        return node
    end
    seen[file_path] = true

    local imports = state.file_imports[file_path] or {}
    for _, fqn in ipairs(imports) do
        -- Find the file that provides this namespace/class
        local target_sym = state.symbols[fqn]
        if target_sym and target_sym.file then
            local child = build_import_subtree(
                state, target_sym.file, max_depth, depth + 1, seen
            )
            table.insert(node.children, child)
        else
            -- Unresolved (external dependency)
            table.insert(node.children, {
                name = fqn,
                children = {},
                annotation = "external",
            })
        end
    end

    return node
end

--- Build the "imported by" direction
function build_imported_by_subtree(state, file_path, max_depth, depth, seen)
    local info = state.files[file_path]
    if not info then
        return {name = file_path, children = {}}
    end

    local node = {
        name = info.namespace or file_path,
        children = {},
    }

    if depth >= max_depth or seen[file_path] then
        return node
    end
    seen[file_path] = true

    -- Find all files that import symbols from this file
    local my_symbols = state.symbols_by_file[file_path] or {}
    local importing_files = {}

    for _, sym_id in ipairs(my_symbols) do
        local refs = graph.find_refs_to(state, sym_id)
        for _, ref in ipairs(refs) do
            if ref.from_file ~= file_path and not importing_files[ref.from_file] then
                importing_files[ref.from_file] = true
            end
        end
    end

    for imp_file, _ in pairs(importing_files) do
        local child = build_imported_by_subtree(
            state, imp_file, max_depth, depth + 1, seen
        )
        table.insert(node.children, child)
    end

    -- Sort children by name
    table.sort(node.children, function(a, b) return a.name < b.name end)

    return node
end

---------------------------------------------------------------------------
-- Caller chain (Stage 4)
---------------------------------------------------------------------------

--- Find callers of a symbol, transitively
-- @param state IndexState
-- @param symbol_name string
-- @param max_depth integer
-- @return table tree_node
local function find_callers(state, symbol_name, max_depth)
    max_depth = max_depth or 3
    local symbols = graph.find_symbol(state, symbol_name)
    if #symbols == 0 then
        return {name = symbol_name, children = {}, annotation = "not found"}
    end

    local primary = symbols[1]
    return build_caller_tree(state, primary.id, max_depth, 0, {})
end

function build_caller_tree(state, symbol_id, max_depth, depth, seen)
    local sym = state.symbols[symbol_id]
    local node = {
        name = symbol_id,
        children = {},
        annotation = sym and string.format("L%d", sym.line) or nil,
    }

    if depth >= max_depth or seen[symbol_id] then
        return node
    end
    seen[symbol_id] = true

    -- Find all references that call this symbol
    local refs = graph.find_refs_to(state, symbol_id)
    local callers = {}

    for _, ref in ipairs(refs) do
        if (ref.kind == "method_call" or ref.kind == "static_call"
            or ref.kind == "static_method_call" or ref.kind == "function_call")
            and ref.from_symbol then
            if not callers[ref.from_symbol] then
                callers[ref.from_symbol] = true
            end
        end
    end

    for caller_id, _ in pairs(callers) do
        local child = build_caller_tree(state, caller_id, max_depth, depth + 1, seen)
        child.annotation = "called by"
        table.insert(node.children, child)
    end

    table.sort(node.children, function(a, b) return a.name < b.name end)
    return node
end

---------------------------------------------------------------------------
-- Callee chain (Stage 4)
---------------------------------------------------------------------------

--- Find what a symbol calls, transitively
-- @param state IndexState
-- @param symbol_name string
-- @param max_depth integer
-- @return table tree_node
local function find_callees(state, symbol_name, max_depth)
    max_depth = max_depth or 2
    local symbols = graph.find_symbol(state, symbol_name)
    if #symbols == 0 then
        return {name = symbol_name, children = {}, annotation = "not found"}
    end

    local primary = symbols[1]
    return build_callee_tree(state, primary.id, max_depth, 0, {})
end

function build_callee_tree(state, symbol_id, max_depth, depth, seen)
    local sym = state.symbols[symbol_id]
    local node = {
        name = symbol_id,
        children = {},
        annotation = sym and string.format("L%d-%d", sym.line, sym.end_line or sym.line) or nil,
    }

    if depth >= max_depth or seen[symbol_id] then
        return node
    end
    seen[symbol_id] = true

    -- Find all calls made by this symbol
    local refs = graph.find_refs_from(state, symbol_id)
    local callees = {}

    for _, ref in ipairs(refs) do
        if ref.kind == "method_call" or ref.kind == "static_call"
            or ref.kind == "static_method_call" or ref.kind == "function_call" then
            local target = ref.to_resolved or ref.to_name
            if target and not callees[target] then
                callees[target] = ref
            end
        end
    end

    for target_id, ref in pairs(callees) do
        if state.symbols[target_id] then
            local child = build_callee_tree(state, target_id, max_depth, depth + 1, seen)
            table.insert(node.children, child)
        else
            table.insert(node.children, {
                name = target_id,
                children = {},
                annotation = "unresolved",
            })
        end
    end

    table.sort(node.children, function(a, b) return a.name < b.name end)
    return node
end

---------------------------------------------------------------------------
-- Hierarchy (Stage 3)
---------------------------------------------------------------------------

--- Build inheritance/implementation tree
-- @param state IndexState
-- @param symbol_name string
-- @param direction string  "up" | "down" | "both"
-- @return table
local function build_hierarchy(state, symbol_name, direction)
    direction = direction or "both"
    local symbols = graph.find_symbol(state, symbol_name)
    if #symbols == 0 then
        return {name = symbol_name, up = {}, down = {}}
    end

    local primary = symbols[1]
    local result = {name = primary.id, kind = primary.kind}

    if direction == "up" or direction == "both" then
        result.ancestors = build_ancestor_chain(state, primary.id, {})
    end

    if direction == "down" or direction == "both" then
        result.descendants = build_descendant_tree(state, primary.id, 10, 0, {})
    end

    return result
end

function build_ancestor_chain(state, symbol_id, seen)
    if seen[symbol_id] then return {} end
    seen[symbol_id] = true

    local chain = {}
    local parent_id = graph.parent_of(state, symbol_id)
    if parent_id then
        table.insert(chain, parent_id)
        local parent_chain = build_ancestor_chain(state, parent_id, seen)
        for _, p in ipairs(parent_chain) do
            table.insert(chain, p)
        end
    end
    return chain
end

function build_descendant_tree(state, symbol_id, max_depth, depth, seen)
    local node = {name = symbol_id, children = {}}

    if depth >= max_depth or seen[symbol_id] then return node end
    seen[symbol_id] = true

    local children_ids = graph.children_of(state, symbol_id)
    for _, child_id in ipairs(children_ids) do
        local child = build_descendant_tree(state, child_id, max_depth, depth + 1, seen)
        table.insert(node.children, child)
    end

    table.sort(node.children, function(a, b) return a.name < b.name end)
    return node
end

---------------------------------------------------------------------------
-- Impact analysis (Stage 5)
---------------------------------------------------------------------------

--- Find files impacted by changing a file/symbol
-- @param state IndexState
-- @param target string  File path or symbol name
-- @param max_depth integer
-- @return table  {direct: {files}, transitive: {files}}
local function find_impacted(state, target, max_depth)
    max_depth = max_depth or 2

    -- Resolve to file
    local file_path = target
    if not state.files[target] then
        local syms = graph.find_symbol(state, target)
        if #syms > 0 and syms[1].file then
            file_path = syms[1].file
        else
            return {direct = {}, transitive = {}}
        end
    end

    local direct = {}
    local all = {}
    all[file_path] = true

    -- Find direct dependents
    local my_symbols = state.symbols_by_file[file_path] or {}
    for _, sym_id in ipairs(my_symbols) do
        local refs = graph.find_refs_to(state, sym_id)
        for _, ref in ipairs(refs) do
            if ref.from_file ~= file_path then
                if not direct[ref.from_file] then
                    direct[ref.from_file] = {
                        file = ref.from_file,
                        refs = {},
                    }
                end
                table.insert(direct[ref.from_file].refs, ref)
                all[ref.from_file] = true
            end
        end
    end

    -- Transitive dependents
    local transitive = {}
    if max_depth > 1 then
        local frontier = {}
        for f, _ in pairs(direct) do
            table.insert(frontier, f)
        end

        for d = 2, max_depth do
            local next_frontier = {}
            for _, dep_file in ipairs(frontier) do
                local dep_syms = state.symbols_by_file[dep_file] or {}
                for _, sym_id in ipairs(dep_syms) do
                    local refs = graph.find_refs_to(state, sym_id)
                    for _, ref in ipairs(refs) do
                        if ref.from_file ~= dep_file and not all[ref.from_file] then
                            all[ref.from_file] = true
                            transitive[ref.from_file] = {
                                file = ref.from_file,
                                via = dep_file,
                                depth = d,
                            }
                            table.insert(next_frontier, ref.from_file)
                        end
                    end
                end
            end
            frontier = next_frontier
        end
    end

    return {
        source = file_path,
        direct = direct,
        transitive = transitive,
    }
end

return {
    find_usages = find_usages,
    build_dep_tree = build_dep_tree,
    find_callers = find_callers,
    find_callees = find_callees,
    build_hierarchy = build_hierarchy,
    find_impacted = find_impacted,
}