from pathlib import Path
from unittest.mock import patch

from blends import get_graphs_from_path
from blends.models import SyntaxGraph
from blends.stack.forward_stitcher import ForwardStitcherConfig
from blends.stack.partial_path import (
    PartialPath,
    PartialPathDatabase,
    PartialPathFileRecord,
    PartialPathIndexRequest,
    PartialPathLimits,
    PartialScopeStack,
    PartialSymbolStack,
    PerFilePartialPathStats,
    StackGraphViewMetadata,
    index_partial_paths_for_file,
)
from blends.stack.partial_path.variables import ScopeStackVariable, SymbolStackVariable
from blends.stack.single_file_resolver import (
    SingleFileResolutionRequest,
    resolve_definitions_single_file,
)
from blends.stack.view import StackGraphView

INDEX_MAX_WORK = 50_000
QUERY_MAX_PHASES = 100
QUERY_MAX_WORK = 10_000
STACK_GRAPH_TEST_DATA_DIR = Path(__file__).parent.parent / "data" / "test_files" / "stack_graph"


def setup_multi_file_database(
    files: dict[str, int],
) -> tuple[dict[int, StackGraphView], PartialPathDatabase]:
    database = PartialPathDatabase()
    views: dict[int, StackGraphView] = {}
    for filename, file_handle in files.items():
        path = STACK_GRAPH_TEST_DATA_DIR / filename
        logical = f"test/data/test_files/stack_graph/{filename}"
        with patch(
            "blends.ctx.ctx.has_feature_flag",
            side_effect=lambda flag: flag == "StackGraph",
        ):
            graphs = get_graphs_from_path(path)
        assert graphs.syntax_graph is not None
        views[file_handle] = StackGraphView.from_syntax_graph(graphs.syntax_graph, path=logical)
        index_partial_paths_for_file(
            graphs.syntax_graph,
            request=PartialPathIndexRequest(
                file_path=logical,
                file_handle=file_handle,
                limits=PartialPathLimits(max_work_per_phase=INDEX_MAX_WORK),
            ),
            database=database,
        )
    return views, database


def view_for_graph(graph: SyntaxGraph, *, path: str = "test.py") -> StackGraphView:
    return StackGraphView.from_syntax_graph(graph, path=path)


def make_path(start: int, end: int) -> PartialPath:
    symbol_var = SymbolStackVariable.initial()
    scope_var = ScopeStackVariable.initial()
    return PartialPath(
        start_node_index=start,
        end_node_index=end,
        symbol_stack_precondition=PartialSymbolStack.from_variable(symbol_var),
        symbol_stack_postcondition=PartialSymbolStack.from_variable(symbol_var),
        scope_stack_precondition=PartialScopeStack.from_variable(scope_var),
        scope_stack_postcondition=PartialScopeStack.from_variable(scope_var),
    )


def record(
    paths: tuple[PartialPath, ...],
    *,
    file_handle: int = 1,
    file_path: str = "test.py",
    symbol_id_to_string: tuple[str, ...] = ("x",),
) -> PartialPathFileRecord:
    metadata = StackGraphViewMetadata(
        file_handles=(file_path,),
        symbol_id_to_string=symbol_id_to_string,
    )
    return PartialPathFileRecord(
        file_handle=file_handle,
        file_path=file_path,
        view_snapshot=metadata,
        paths=paths,
        stats=PerFilePartialPathStats(),
    )


def setup_stack_resolution(test_file: str) -> tuple[StackGraphView, PartialPathFileRecord]:
    test_file_path = STACK_GRAPH_TEST_DATA_DIR / test_file
    file_path = f"test/data/test_files/stack_graph/{test_file}"

    with patch(
        "blends.ctx.ctx.has_feature_flag",
        side_effect=lambda flag: flag == "StackGraph",
    ):
        graphs = get_graphs_from_path(test_file_path)

    assert graphs.syntax_graph is not None

    view = StackGraphView.from_syntax_graph(graphs.syntax_graph, path=file_path)

    database = PartialPathDatabase()
    per_file = index_partial_paths_for_file(
        graphs.syntax_graph,
        request=PartialPathIndexRequest(
            file_path=file_path,
            file_handle=1,
            limits=PartialPathLimits(max_work_per_phase=INDEX_MAX_WORK),
        ),
        database=database,
    )
    assert per_file.paths, "Expected at least one partial path"

    file_record = database.get_file(1)
    assert file_record is not None

    return view, file_record


def find_symbol_nodes(
    view: StackGraphView,
    *,
    symbol_id: int,
    is_definition: bool,
) -> list[int]:
    matches: list[int] = []
    flags = view.node_is_definition if is_definition else view.node_is_reference
    for node_index, flag in enumerate(flags):
        if not flag:
            continue
        if view.node_symbol_id[node_index] == symbol_id:
            matches.append(node_index)
    return matches


def resolve_symbol(
    view: StackGraphView,
    file_record: PartialPathFileRecord,
    symbol: str,
    ref_index: int = 0,
) -> list[int]:
    symbol_id = view.symbols.symbol_to_id.get(symbol)
    if symbol_id is None:
        return []

    ref_nodes = find_symbol_nodes(view, symbol_id=symbol_id, is_definition=False)
    if ref_index >= len(ref_nodes):
        return []

    request = SingleFileResolutionRequest(
        file_handle=1,
        record=file_record,
        ref_node_index=ref_nodes[ref_index],
        config=ForwardStitcherConfig(
            max_phases=QUERY_MAX_PHASES,
            max_work_per_phase=QUERY_MAX_WORK,
            max_total_work=QUERY_MAX_WORK,
        ),
    )
    result = resolve_definitions_single_file(view, request=request)
    assert isinstance(result, list)
    return result


def count_defs_refs(view: StackGraphView, symbol: str) -> tuple[int, int]:
    symbol_id = view.symbols.symbol_to_id.get(symbol)
    if symbol_id is None:
        return (0, 0)
    defs = find_symbol_nodes(view, symbol_id=symbol_id, is_definition=True)
    refs = find_symbol_nodes(view, symbol_id=symbol_id, is_definition=False)
    return (len(defs), len(refs))


def get_definition_precedences(view: StackGraphView, symbol: str) -> list[int]:
    symbol_id = view.symbols.symbol_to_id.get(symbol)
    if symbol_id is None:
        return []
    def_nodes = find_symbol_nodes(view, symbol_id=symbol_id, is_definition=True)
    return [view.node_precedence[node_index] for node_index in def_nodes]
