"""Core call graph logic."""

from __future__ import annotations

from collections import defaultdict
from typing import TYPE_CHECKING

if TYPE_CHECKING:
    from pathlib import Path

from vibelexity.call_graph.parsers.python import PythonCallParser
from vibelexity.models import CallGraph, CallEdge, FunctionDefinition
from vibelexity.parsers.python import parse


def _collect_function_definitions(
    files: list[Path], parser: PythonCallParser
) -> tuple[dict[tuple[str, str], FunctionDefinition], list[FunctionDefinition]]:
    functions: dict[tuple[str, str], FunctionDefinition] = {}
    all_functions: list[FunctionDefinition] = []

    for filepath in files:
        source = filepath.read_text()
        root = parse(source)
        defs = parser.extract_function_definitions(root, source, str(filepath))
        all_functions.extend(defs)
        for d in defs:
            functions[(d.name, d.filepath)] = d

    return functions, all_functions


def _collect_call_edges(
    files: list[Path],
    parser: PythonCallParser,
    all_functions: list[FunctionDefinition],
) -> list[CallEdge]:
    calls: list[CallEdge] = []

    for filepath in files:
        source = filepath.read_text()
        root = parse(source)
        file_calls = parser.extract_calls(root, source, str(filepath), all_functions)
        calls.extend(file_calls)

    return calls


def _resolve_callee(
    call: CallEdge, all_functions: list[FunctionDefinition]
) -> CallEdge:
    callee_defs = [f for f in all_functions if f.name == call.callee_function]
    if not callee_defs:
        call.call_type = "external"
        return call

    same_file_defs = [f for f in callee_defs if f.filepath == call.caller_file]
    if same_file_defs:
        callee = same_file_defs[0]
    else:
        callee = callee_defs[0]

    call.callee_file = callee.filepath
    call.callee_line = callee.line
    return call


def resolve_calls(
    calls: list[CallEdge], all_functions: list[FunctionDefinition]
) -> list[CallEdge]:
    resolved_calls: list[CallEdge] = []
    for call in calls:
        if call.call_type == "external":
            resolved_calls.append(call)
        else:
            resolved_calls.append(_resolve_callee(call, all_functions))
    return resolved_calls


def build_call_graph(files: list[Path], lang: str = "python") -> CallGraph:
    parser = PythonCallParser()

    functions, all_functions = _collect_function_definitions(files, parser)
    calls = _collect_call_edges(files, parser, all_functions)
    resolved_calls = resolve_calls(calls, all_functions)

    reverse_calls: dict[str, list[CallEdge]] = defaultdict(list)
    for call in resolved_calls:
        reverse_calls[call.callee_function].append(call)

    return CallGraph(
        functions=functions, calls=resolved_calls, reverse_calls=dict(reverse_calls)
    )


def find_call_cycles(graph: CallGraph) -> list[list]:
    colors = {k: "WHITE" for k in graph.functions.keys()}
    cycles: list[list[tuple[str, str]]] = []

    def dfs(key: tuple[str, str], path: list, edges: list) -> None:
        colors[key] = "GRAY"
        path.append(key)

        func_name = key[0]
        calls = [
            c
            for c in graph.calls
            if c.caller_function == func_name and c.call_type == "local"
        ]
        for call in calls:
            callee_key = (
                (call.callee_function, call.callee_file) if call.callee_file else None
            )
            if not callee_key:
                continue
            if callee_key in colors:
                if colors[callee_key] == "GRAY":
                    cycles.append(path.copy() + [callee_key])
                elif colors[callee_key] == "WHITE":
                    dfs(callee_key, list(path), list(edges))

        colors[key] = "BLACK"

    for key in graph.functions.keys():
        if colors[key] == "WHITE":
            dfs(key, [], [])

    return cycles


def find_unused_functions(graph: CallGraph) -> list[FunctionDefinition]:
    used = {c.callee_function for c in graph.calls}

    unused: list[FunctionDefinition] = []
    for (name, _), func_def in graph.functions.items():
        if name not in used and not name.startswith("_"):
            unused.append(func_def)

    return unused
