from unittest.mock import patch

import pytest

from blends import get_graphs_from_path
from blends.stack.node_kinds import StackGraphNodeKind
from blends.stack.partial_path import (
    PartialPathDatabase,
    PartialPathIndexRequest,
    PartialPathLimits,
    index_partial_paths_for_file,
)
from blends.stack.view import ROOT_INDEX, StackGraphView
from test.stack.test_helpers import STACK_GRAPH_TEST_DATA_DIR

INDEX_MAX_WORK = 50_000


def _load_file(
    test_file: str,
) -> tuple[StackGraphView, PartialPathDatabase]:
    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()
    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,
    )

    return view, database


@pytest.mark.blends_test_group("stack_unittesting")
def test_file_a_has_module_definition_path_from_root() -> None:
    view, _db = _load_file("cross_file_a.py")

    assert view.node_kind[ROOT_INDEX] == StackGraphNodeKind.ROOT.value

    root_successors = [idx for idx, _prec in view.outgoing[ROOT_INDEX]]
    assert len(root_successors) >= 1

    has_pop_for_module = False
    for succ_idx in root_successors:
        kind = view.node_kind[succ_idx]
        if kind == StackGraphNodeKind.POP_SYMBOL.value:
            symbol_id = view.node_symbol_id[succ_idx]
            assert symbol_id is not None
            symbol = view.symbols.id_to_symbol[symbol_id]
            if "cross_file_a" in symbol:
                has_pop_for_module = True

    assert has_pop_for_module, (
        "Expected ROOT to have an outgoing pop_symbol edge containing 'cross_file_a'"
    )


@pytest.mark.blends_test_group("stack_unittesting")
def test_file_b_has_import_reference_chain_to_root() -> None:
    view, _db = _load_file("cross_file_b.py")

    assert view.node_kind[ROOT_INDEX] == StackGraphNodeKind.ROOT.value

    push_nodes_to_root = []
    for node_idx in range(len(view.node_kind)):
        kind = view.node_kind[node_idx]
        if kind != StackGraphNodeKind.PUSH_SYMBOL.value:
            continue
        successors = [idx for idx, _prec in view.outgoing[node_idx]]
        if ROOT_INDEX in successors:
            push_nodes_to_root.append(node_idx)

    assert len(push_nodes_to_root) >= 1, "Expected at least one push_symbol node with edge to ROOT"

    pushed_symbols = []
    for node_idx in push_nodes_to_root:
        symbol_id = view.node_symbol_id[node_idx]
        if symbol_id is not None:
            pushed_symbols.append(view.symbols.id_to_symbol[symbol_id])

    assert any("cross_file_a" in sym for sym in pushed_symbols), (
        f"Expected a push toward ROOT with 'cross_file_a', got {pushed_symbols}"
    )


@pytest.mark.blends_test_group("stack_unittesting")
def test_partial_paths_are_indexed_for_both_files() -> None:
    _view_a, db_a = _load_file("cross_file_a.py")
    _view_b, db_b = _load_file("cross_file_b.py")

    record_a = db_a.get_file(1)
    record_b = db_b.get_file(1)

    assert record_a is not None
    assert record_b is not None
    assert len(record_a.paths) > 0, "File A should have partial paths"
    assert len(record_b.paths) > 0, "File B should have partial paths"


@pytest.mark.blends_test_group("stack_unittesting")
def test_definition_and_reference_paths_have_compatible_symbols() -> None:
    view_a, db_a = _load_file("cross_file_a.py")
    view_b, db_b = _load_file("cross_file_b.py")

    record_a = db_a.get_file(1)
    record_b = db_b.get_file(1)
    assert record_a is not None
    assert record_b is not None

    def_precondition_symbols: set[str] = set()
    for path in record_a.paths:
        if path.start_node_index == ROOT_INDEX:
            for scoped_sym in path.symbol_stack_precondition.symbols:
                symbol = view_a.symbols.id_to_symbol[scoped_sym.symbol_id]
                def_precondition_symbols.add(symbol)

    ref_postcondition_symbols: set[str] = set()
    for path in record_b.paths:
        if path.end_node_index == ROOT_INDEX:
            for scoped_sym in path.symbol_stack_postcondition.symbols:
                symbol = view_b.symbols.id_to_symbol[scoped_sym.symbol_id]
                ref_postcondition_symbols.add(symbol)

    shared_symbols = def_precondition_symbols & ref_postcondition_symbols
    assert len(shared_symbols) > 0, (
        f"Expected shared symbols between definition preconditions {def_precondition_symbols} "
        f"and reference postconditions {ref_postcondition_symbols}"
    )
