import pytest

from blends.models import SyntaxGraph
from blends.stack.attributes import (
    STACK_GRAPH_IS_DEFINITION,
    STACK_GRAPH_IS_REFERENCE,
    STACK_GRAPH_KIND,
    STACK_GRAPH_SYMBOL,
)
from blends.stack.forward_stitcher import (
    ForwardStitcher,
    ForwardStitcherConfig,
)
from blends.stack.node_kinds import StackGraphNodeKind
from blends.stack.partial_path_db import PartialPathDb
from blends.stack.view import StackGraphView
from test.stack.test_helpers import make_path, view_for_graph
from test.stack.test_helpers import record as make_record

EXPECTED_MAX_PHASES = 100


def _create_test_graph_view(num_join_nodes: int = 1) -> StackGraphView:
    graph = SyntaxGraph()
    graph.add_node("ref")
    graph.update_node_attribute("ref", STACK_GRAPH_KIND, StackGraphNodeKind.PUSH_SYMBOL.value)
    graph.update_node_attribute("ref", STACK_GRAPH_SYMBOL, "x")
    graph.update_node_attribute("ref", STACK_GRAPH_IS_REFERENCE, value=True)

    for i in range(num_join_nodes):
        node_name = f"join_{i}"
        graph.add_node(node_name)
        graph.update_node_attribute(node_name, STACK_GRAPH_KIND, StackGraphNodeKind.SCOPE.value)

    graph.add_node("def")
    graph.update_node_attribute("def", STACK_GRAPH_KIND, StackGraphNodeKind.POP_SYMBOL.value)
    graph.update_node_attribute("def", STACK_GRAPH_SYMBOL, "x")
    graph.update_node_attribute("def", STACK_GRAPH_IS_DEFINITION, value=True)

    return view_for_graph(graph)


@pytest.mark.blends_test_group("stack_unittesting")
def test_cycle_guard_can_be_disabled() -> None:
    view = _create_test_graph_view()
    ref_index = view.nid_to_index["ref"]
    join_index = view.nid_to_index["join_0"]

    db = PartialPathDb()
    db.add_file(
        make_record(
            (
                make_path(ref_index, join_index),
                make_path(join_index, join_index),
            )
        )
    )

    max_phases = 10
    stitcher = ForwardStitcher.from_references(
        view=view,
        file_handle=1,
        db=db,
        reference_nodes=(ref_index,),
        config=ForwardStitcherConfig(enable_cycle_guard=False, max_phases=max_phases),
    )

    result = stitcher.process_next_phase()
    while not stitcher.is_complete():
        result = stitcher.process_next_phase()

    assert result.stats.cycle_guard_hits == 0
    assert result.stats.phases_executed == max_phases


@pytest.mark.blends_test_group("stack_unittesting")
def test_cycle_guard_with_variable_paths() -> None:
    view = _create_test_graph_view()
    ref_index = view.nid_to_index["ref"]
    join_index = view.nid_to_index["join_0"]

    db = PartialPathDb()
    db.add_file(
        make_record(
            (
                make_path(ref_index, join_index),
                make_path(join_index, join_index),
            )
        )
    )

    stitcher = ForwardStitcher.from_references(
        view=view,
        file_handle=1,
        db=db,
        reference_nodes=(ref_index,),
        config=ForwardStitcherConfig(enable_cycle_guard=True, max_phases=10),
    )

    result = stitcher.process_next_phase()
    while not stitcher.is_complete():
        result = stitcher.process_next_phase()

    assert stitcher.is_complete()
    assert result.stats.cycle_guard_hits > 0


@pytest.mark.blends_test_group("stack_unittesting")
def test_queue_size_limit_prunes_excess_paths() -> None:
    view = _create_test_graph_view(num_join_nodes=10)
    ref_index = view.nid_to_index["ref"]
    def_index = view.nid_to_index["def"]

    paths = []
    for i in range(10):
        join_index = view.nid_to_index[f"join_{i}"]
        paths.append(make_path(ref_index, join_index))
        paths.append(make_path(join_index, def_index))

    db = PartialPathDb()
    db.add_file(make_record(tuple(paths)))

    stitcher = ForwardStitcher.from_references(
        view=view,
        file_handle=1,
        db=db,
        reference_nodes=(ref_index,),
        config=ForwardStitcherConfig(max_queue_size=3),
    )

    result = stitcher.process_next_phase()

    assert result.stats.queue_pruned_count > 0


@pytest.mark.blends_test_group("stack_unittesting")
def test_queue_pruning_is_deterministic() -> None:
    view = _create_test_graph_view(num_join_nodes=5)
    ref_index = view.nid_to_index["ref"]
    def_index = view.nid_to_index["def"]

    paths = []
    for i in range(5):
        join_index = view.nid_to_index[f"join_{i}"]
        paths.append(make_path(ref_index, join_index))
        paths.append(make_path(join_index, def_index))

    config = ForwardStitcherConfig(max_queue_size=2)

    results = []
    for _ in range(2):
        db = PartialPathDb()
        db.add_file(make_record(tuple(paths)))

        stitcher = ForwardStitcher.from_references(
            view=view,
            file_handle=1,
            db=db,
            reference_nodes=(ref_index,),
            config=config,
        )

        result = stitcher.process_next_phase()
        while not stitcher.is_complete():
            result = stitcher.process_next_phase()
        results.append(result)

    assert results[0].stats.queue_pruned_count == results[1].stats.queue_pruned_count
    assert results[0].stats.complete_paths_found == results[1].stats.complete_paths_found


@pytest.mark.blends_test_group("stack_unittesting")
def test_max_work_per_phase_stops_expansion() -> None:
    view = _create_test_graph_view()
    ref_index = view.nid_to_index["ref"]
    def_index = view.nid_to_index["def"]

    db = PartialPathDb()
    db.add_file(make_record((make_path(ref_index, def_index),)))

    stitcher = ForwardStitcher.from_references(
        view=view,
        file_handle=1,
        db=db,
        reference_nodes=(ref_index,),
        config=ForwardStitcherConfig(max_work_per_phase=0),
    )

    result = stitcher.process_next_phase()

    assert result.stats.candidates_considered == 0
    assert not stitcher.is_complete()


@pytest.mark.blends_test_group("stack_unittesting")
def test_max_total_work_stops_across_phases() -> None:
    view = _create_test_graph_view()
    ref_index = view.nid_to_index["ref"]
    join_index = view.nid_to_index["join_0"]
    def_index = view.nid_to_index["def"]

    paths = (
        make_path(ref_index, join_index),
        make_path(join_index, def_index),
    )

    db_limited = PartialPathDb()
    db_limited.add_file(make_record(paths))
    stitcher_limited = ForwardStitcher.from_references(
        view=view,
        file_handle=1,
        db=db_limited,
        reference_nodes=(ref_index,),
        config=ForwardStitcherConfig(max_total_work=1),
    )

    result_limited = stitcher_limited.process_next_phase()
    while not stitcher_limited.is_complete():
        result_limited = stitcher_limited.process_next_phase()

    db_unlimited = PartialPathDb()
    db_unlimited.add_file(make_record(paths))
    stitcher_unlimited = ForwardStitcher.from_references(
        view=view,
        file_handle=1,
        db=db_unlimited,
        reference_nodes=(ref_index,),
        config=ForwardStitcherConfig(),
    )

    result_unlimited = stitcher_unlimited.process_next_phase()
    while not stitcher_unlimited.is_complete():
        result_unlimited = stitcher_unlimited.process_next_phase()

    assert stitcher_limited.is_complete()
    assert result_limited.stats.candidates_considered < result_unlimited.stats.candidates_considered
    assert result_limited.stats.complete_paths_found == 0

    assert stitcher_unlimited.is_complete()
    assert result_unlimited.stats.complete_paths_found > 0


@pytest.mark.blends_test_group("stack_unittesting")
def test_max_phases_stops_expansion() -> None:
    view = _create_test_graph_view()
    ref_index = view.nid_to_index["ref"]
    join_index = view.nid_to_index["join_0"]
    def_index = view.nid_to_index["def"]

    paths = (
        make_path(ref_index, join_index),
        make_path(join_index, def_index),
    )

    db_limited = PartialPathDb()
    db_limited.add_file(make_record(paths))
    stitcher_limited = ForwardStitcher.from_references(
        view=view,
        file_handle=1,
        db=db_limited,
        reference_nodes=(ref_index,),
        config=ForwardStitcherConfig(max_phases=1),
    )

    result_limited = stitcher_limited.process_next_phase()
    while not stitcher_limited.is_complete():
        result_limited = stitcher_limited.process_next_phase()

    db_unlimited = PartialPathDb()
    db_unlimited.add_file(make_record(paths))
    stitcher_unlimited = ForwardStitcher.from_references(
        view=view,
        file_handle=1,
        db=db_unlimited,
        reference_nodes=(ref_index,),
        config=ForwardStitcherConfig(),
    )

    result_unlimited = stitcher_unlimited.process_next_phase()
    while not stitcher_unlimited.is_complete():
        result_unlimited = stitcher_unlimited.process_next_phase()

    assert stitcher_limited.is_complete()
    assert result_limited.stats.phases_executed == 1
    assert result_limited.stats.complete_paths_found == 0

    assert stitcher_unlimited.is_complete()
    assert result_unlimited.stats.phases_executed > 1
    assert result_unlimited.stats.complete_paths_found > 0


@pytest.mark.blends_test_group("stack_unittesting")
def test_synthetic_blowup_graph_bounded_by_queue_limit() -> None:
    view = _create_test_graph_view(num_join_nodes=20)
    ref_index = view.nid_to_index["ref"]
    def_index = view.nid_to_index["def"]

    paths = []
    for i in range(20):
        join_index = view.nid_to_index[f"join_{i}"]
        paths.append(make_path(ref_index, join_index))
        paths.append(make_path(join_index, def_index))

    db = PartialPathDb()
    db.add_file(make_record(tuple(paths)))

    stitcher = ForwardStitcher.from_references(
        view=view,
        file_handle=1,
        db=db,
        reference_nodes=(ref_index,),
        config=ForwardStitcherConfig(
            max_queue_size=5,
            max_phases=10,
        ),
    )

    result = stitcher.process_next_phase()
    while not stitcher.is_complete():
        result = stitcher.process_next_phase()

    assert result.stats.queue_pruned_count > 0
    assert stitcher.is_complete()


@pytest.mark.blends_test_group("stack_unittesting")
def test_synthetic_cycle_graph_terminates() -> None:
    view = _create_test_graph_view()
    ref_index = view.nid_to_index["ref"]
    join_index = view.nid_to_index["join_0"]
    def_index = view.nid_to_index["def"]

    db = PartialPathDb()
    db.add_file(
        make_record(
            (
                make_path(ref_index, join_index),
                make_path(join_index, join_index),
                make_path(join_index, def_index),
            )
        )
    )

    stitcher = ForwardStitcher.from_references(
        view=view,
        file_handle=1,
        db=db,
        reference_nodes=(ref_index,),
        config=ForwardStitcherConfig(
            enable_cycle_guard=True,
            max_phases=100,
        ),
    )

    phases_run = 0
    result = stitcher.process_next_phase()
    phases_run += 1
    while not stitcher.is_complete():
        result = stitcher.process_next_phase()
        phases_run += 1

    assert phases_run < EXPECTED_MAX_PHASES
    assert stitcher.is_complete()
    assert result.stats.cycle_guard_hits > 0


@pytest.mark.blends_test_group("stack_unittesting")
def test_all_limits_work_together() -> None:
    view = _create_test_graph_view(num_join_nodes=10)
    ref_index = view.nid_to_index["ref"]
    def_index = view.nid_to_index["def"]

    paths = []
    for i in range(10):
        join_index = view.nid_to_index[f"join_{i}"]
        paths.append(make_path(ref_index, join_index))
        paths.append(make_path(join_index, join_index))
        paths.append(make_path(join_index, def_index))

    max_work_per_phase = 50
    max_total_work = 200
    max_queue_size = 10
    max_phases = 20

    db = PartialPathDb()
    db.add_file(make_record(tuple(paths)))

    stitcher = ForwardStitcher.from_references(
        view=view,
        file_handle=1,
        db=db,
        reference_nodes=(ref_index,),
        config=ForwardStitcherConfig(
            enable_cycle_guard=True,
            max_work_per_phase=max_work_per_phase,
            max_total_work=max_total_work,
            max_queue_size=max_queue_size,
            max_phases=max_phases,
        ),
    )

    result = stitcher.process_next_phase()
    while not stitcher.is_complete():
        result = stitcher.process_next_phase()

    assert stitcher.is_complete()
    assert result.stats.phases_executed <= max_phases

    max_candidates_per_extension = len(paths)
    assert result.stats.candidates_considered <= max_total_work + max_candidates_per_extension
    assert result.stats.cycle_guard_hits > 0 or result.stats.queue_pruned_count > 0
