from __future__ import annotations

import pytest
from django.core.exceptions import ValidationError
from django.db import IntegrityError

from control_plane.agents.models import ProgenySignal
from graph.models import ExecutionGraphDefinition, ExecutionGraphVersion, ExecutionGraphVersionStatus, GraphRun, GraphRunStatus
from graph.registry import NodeHandlerRegistry, NodeResult
from graph.spec import ExecutionGraphSpec, outgoing_edges
from graph.task_execution import task_execution_graph_v1


class FixtureNode:
    node_type = "fixture"
    idempotent = True
    replay_safe = True
    destructive = False

    def run(self, context: object) -> NodeResult:
        return NodeResult("COMPLETE", "success", {"ok": True})


def test_task_execution_graph_v1_is_serializable_and_valid() -> None:
    spec = task_execution_graph_v1()

    payload = spec.to_dict()
    restored = ExecutionGraphSpec.from_dict(payload)
    restored.validate()

    assert restored.name == "task_execution"
    assert restored.version == 1
    assert restored.entry == "claim_task"
    assert restored.nodes["coder"].metadata["contains_internal_tool_loop"] is True
    assert {edge.condition for edge in outgoing_edges(restored, "review")} == {"PASS", "REWORK_REQUIRED", "REJECTED"}
    assert {edge.condition for edge in outgoing_edges(restored, "retry_or_fail")} == {"retry_available", "retry_exhausted"}


def test_node_registry_resolves_handlers_without_arbitrary_execution() -> None:
    registry = NodeHandlerRegistry()
    registry.register(FixtureNode())

    assert registry.get("fixture").run(object()).output_metadata == {"ok": True}
    with pytest.raises(KeyError):
        registry.get("missing")


def test_graph_models_persist_versioned_run_state() -> None:
    spec = task_execution_graph_v1()
    definition = ExecutionGraphDefinition.objects.create(
        name=spec.name,
        graph_type=spec.graph_type,
        description="Task execution graph",
    )
    version = ExecutionGraphVersion.objects.create(
        graph=definition,
        version=spec.version,
        status=ExecutionGraphVersionStatus.CHAMPION,
        graph_spec=spec.to_dict(),
        metadata={"immutable_after_use": True},
    )

    run = GraphRun.objects.create(
        execution_graph_version=version,
        status=GraphRunStatus.PENDING,
        current_node=spec.entry,
        metadata={"node_count": len(spec.nodes)},
    )

    assert run.execution_graph_version.graph.name == "task_execution"
    assert run.execution_graph_version.status == ExecutionGraphVersionStatus.CHAMPION
    assert run.metadata["node_count"] == 12


def test_champion_graph_version_spec_is_immutable() -> None:
    spec = task_execution_graph_v1()
    definition = ExecutionGraphDefinition.objects.create(name="immutable_task_execution", graph_type=spec.graph_type)
    version = ExecutionGraphVersion.objects.create(graph=definition, version=1, status=ExecutionGraphVersionStatus.CHAMPION, graph_spec=spec.to_dict())

    version.graph_spec = {**version.graph_spec, "entry": "different"}

    with pytest.raises(ValidationError):
        version.save()


def test_graph_changes_require_a_new_version() -> None:
    spec = task_execution_graph_v1()
    definition = ExecutionGraphDefinition.objects.create(name="versioned_task_execution", graph_type=spec.graph_type)
    v1 = ExecutionGraphVersion.objects.create(graph=definition, version=1, status=ExecutionGraphVersionStatus.CHAMPION, graph_spec=spec.to_dict())
    changed_spec = spec.to_dict()
    changed_spec["metadata"] = {**changed_spec["metadata"], "change": "new node planned"}

    v2 = ExecutionGraphVersion.objects.create(graph=definition, version=2, status=ExecutionGraphVersionStatus.CHALLENGER, graph_spec=changed_spec)

    assert v1.graph_spec != v2.graph_spec
    assert v2.status == ExecutionGraphVersionStatus.CHALLENGER


def test_only_one_champion_version_per_graph_definition() -> None:
    spec = task_execution_graph_v1()
    definition = ExecutionGraphDefinition.objects.create(name="single_champion_task_execution", graph_type=spec.graph_type)
    ExecutionGraphVersion.objects.create(graph=definition, version=1, status=ExecutionGraphVersionStatus.CHAMPION, graph_spec=spec.to_dict())

    with pytest.raises((IntegrityError, ValidationError)):
        ExecutionGraphVersion.objects.create(graph=definition, version=2, status=ExecutionGraphVersionStatus.CHAMPION, graph_spec=spec.to_dict())


def test_commit_lineage_preserves_exact_graph_version() -> None:
    spec = task_execution_graph_v1()
    definition = ExecutionGraphDefinition.objects.create(name="lineage_preserved_task_execution", graph_type=spec.graph_type)
    v1 = ExecutionGraphVersion.objects.create(graph=definition, version=1, status=ExecutionGraphVersionStatus.CHAMPION, graph_spec=spec.to_dict())
    v2 = ExecutionGraphVersion.objects.create(graph=definition, version=2, status=ExecutionGraphVersionStatus.CHALLENGER, graph_spec=spec.to_dict())
    run = GraphRun.objects.create(execution_graph_version=v1, status=GraphRunStatus.COMPLETE, current_node="complete")

    v2.status = ExecutionGraphVersionStatus.RETIRED
    v2.save()
    run.refresh_from_db()

    assert run.execution_graph_version == v1
    assert run.execution_graph_version.version == 1


def test_progeny_signal_can_reference_graph_lineage() -> None:
    spec = task_execution_graph_v1()
    definition = ExecutionGraphDefinition.objects.create(name="lineage_graph", graph_type=spec.graph_type)
    version = ExecutionGraphVersion.objects.create(graph=definition, version=1, status=ExecutionGraphVersionStatus.CHAMPION, graph_spec=spec.to_dict())
    run = GraphRun.objects.create(execution_graph_version=version, current_node=spec.entry)
    node_run = run.node_runs.create(node_id="coder", node_type="coder", visit_index=1)

    signal = ProgenySignal.objects.create(
        source="runtime",
        severity="high",
        failure_category="MODEL_OUTPUT_INVALID",
        summary="Malformed model JSON",
        graph_run=run,
        graph_node_run=node_run,
        execution_graph_version=version,
    )

    assert signal.graph_run == run
    assert signal.graph_node_run == node_run
    assert signal.execution_graph_version == version