from __future__ import annotations

from pathlib import Path

import pytest
from django.core.exceptions import ValidationError

from agents.providers import DeterministicCodingProvider
from agents.replay_arena import ReplayArena, replay_experiment_registry
from control_plane.agents.management.commands.seed_core_agents import Command as SeedAgentsCommand
from control_plane.agents.models import Agent, AgentRole, AgentVersion, ExperimentComparison, ExperimentVariant, ImprovementCandidate, ProgenyExperiment, PromotionStatus, ReplayResult, ReplayRun
from control_plane.projects.models import CommitRecord, Milestone, Project, ProjectPlan, Task, TaskStatus
from graph.bootstrap import champion_task_execution_graph_v1
from graph.models import ExecutionGraphDefinition, ExecutionGraphVersion, ExecutionGraphVersionStatus, GraphRun, GraphRunStatus
from graph.native_runtime import NativeGraphRuntime
from graph.progeny_replay import progeny_replay_experiment_graph_v1
from graph.task_execution import task_execution_graph_v1
from model_router.router import ModelRouter
from tests.test_m2_autonomous_loop import create_disposable_django_repo


def source_task(repo: Path) -> Task:
    project = Project.objects.create(name="Replay Source", goal="Source", repository_path=str(repo))
    plan = ProjectPlan.objects.create(project=project, version=1, goal="Source")
    milestone = Milestone.objects.create(project=project, plan=plan, key="R1", title="Replay", goal="Replay")
    return Task.objects.create(
        project=project,
        milestone=milestone,
        task_type="implementation",
        status=TaskStatus.COMPLETE,
        goal='Add a /health endpoint returning JSON {"status": "ok"} and add tests.',
        acceptance_criteria=["/health returns ok", "tests pass"],
    )


def candidate(target_type: str = "EXECUTION_GRAPH") -> ImprovementCandidate:
    return ImprovementCandidate.objects.create(target_type=target_type, hypothesis="Compare challenger against champion", recommended_route="Progeny Graph Evolution")


def frozen_dataset(tmp_path: Path, arena: ReplayArena, *, count: int = 1):
    dataset = arena.create_dataset("Replay Dataset", selection_criteria={"task_type": "implementation"})
    cases = []
    for index in range(count):
        root = tmp_path / f"repo-{index}"
        root.mkdir()
        task = source_task(create_disposable_django_repo(root))
        cases.append(arena.add_case_from_task(dataset, task))
    arena.freeze_dataset(dataset)
    return dataset, cases


def experiment_with_variants(tmp_path: Path, *, count: int = 1, minimum: int = 1):
    SeedAgentsCommand().handle()
    arena = ReplayArena(ModelRouter({"qwen": DeterministicCodingProvider()}), test_command=["python", "manage.py", "test"])
    dataset, cases = frozen_dataset(tmp_path, arena, count=count)
    experiment = arena.create_experiment(candidate(), dataset, success_criteria={"minimum_replay_cases": minimum})
    champion = arena.add_champion(experiment, graph_version=champion_task_execution_graph_v1())
    challenger = arena.add_challenger(experiment, graph_version=arena.ensure_static_analysis_graph_challenger())
    return arena, dataset, cases, experiment, champion, challenger


def test_replay_dataset_freeze_makes_membership_immutable_and_replacement_versions(tmp_path: Path) -> None:
    arena = ReplayArena()
    dataset, _ = frozen_dataset(tmp_path, arena)
    new_root = tmp_path / "new-repo"
    new_root.mkdir()

    with pytest.raises(ValidationError):
        arena.add_case_from_task(dataset, source_task(create_disposable_django_repo(new_root)))

    replacement = arena.create_replacement_dataset_version(dataset)

    assert dataset.status == "FROZEN"
    assert replacement.version == dataset.version + 1
    assert replacement.status == "DRAFT"
    assert replacement.metadata["replaces_dataset_id"] == str(dataset.id)


def test_experiment_variants_are_immutable_and_preserve_target_lineage(tmp_path: Path) -> None:
    arena, _, _, experiment, champion, challenger = experiment_with_variants(tmp_path)

    assert champion.role == "CHAMPION"
    assert champion.execution_graph_version.status == ExecutionGraphVersionStatus.CHAMPION
    assert challenger.role == "CHALLENGER"
    assert challenger.execution_graph_version.version == 2
    assert challenger.configuration_snapshot["graph_spec"]["metadata"]["parent_version"] == 1

    challenger.target_reference = "changed"
    with pytest.raises(ValidationError):
        challenger.save()


def test_agent_version_variant_supported_with_policy_snapshot(tmp_path: Path) -> None:
    SeedAgentsCommand().handle()
    arena = ReplayArena(ModelRouter({"qwen": DeterministicCodingProvider()}), test_command=["python", "manage.py", "test"])
    dataset, _ = frozen_dataset(tmp_path, arena)
    coder = Agent.objects.get(role=AgentRole.CODER)
    challenger_agent = AgentVersion.objects.create(agent=coder, version=99, model="qwen", system_contract="Coder challenger", promotion_status=PromotionStatus.CHALLENGER, context_policy={"window": "same"})
    experiment = arena.create_experiment(candidate("AGENT"), dataset, success_criteria={"minimum_replay_cases": 1})

    variant = arena.add_challenger(experiment, agent_version=challenger_agent)

    assert variant.target_type == "AGENT"
    assert variant.agent_version == challenger_agent
    assert variant.configuration_snapshot["system_contract"] == "Coder challenger"


def test_replay_execution_uses_fresh_worktree_and_preserves_source_task(tmp_path: Path) -> None:
    arena, _, cases, experiment, champion, _ = experiment_with_variants(tmp_path)
    source = cases[0].source_task

    run = arena.run_case(experiment, cases[0], champion)

    source.refresh_from_db()
    assert source.status == TaskStatus.COMPLETE
    assert run.status == "COMPLETE", run.graph_run.metadata
    assert run.replay_task != source
    assert run.graph_run.execution_graph_version == champion.execution_graph_version
    assert run.metadata["production_safe"] is True
    assert run.metadata["commit_label"] == "REPLAY / EXPERIMENTAL"
    assert CommitRecord.objects.filter(task=source).count() == 0
    assert CommitRecord.objects.filter(task=run.replay_task).count() == 1


def test_comparison_normalizes_metrics_and_detects_paired_regressions(tmp_path: Path) -> None:
    arena, _, cases, experiment, champion, challenger = experiment_with_variants(tmp_path, count=2, minimum=2)
    for replay_case in cases:
        champion_run = ReplayRun.objects.create(experiment=experiment, variant=champion, replay_case=replay_case, status="COMPLETE")
        ReplayResult.objects.create(replay_run=champion_run, completion_status="COMPLETE", accepted_candidate=True, metrics={"accepted_candidate": 1, "review_pass": 1, "judge_pass": 1, "runtime_seconds": 10})
    first = ReplayRun.objects.create(experiment=experiment, variant=challenger, replay_case=cases[0], status="COMPLETE")
    ReplayResult.objects.create(replay_run=first, completion_status="COMPLETE", accepted_candidate=True, metrics={"accepted_candidate": 1, "review_pass": 1, "judge_pass": 1, "runtime_seconds": 8})
    second = ReplayRun.objects.create(experiment=experiment, variant=challenger, replay_case=cases[1], status="FAILED", failure_classification="VARIANT_FAILURE")
    ReplayResult.objects.create(replay_run=second, completion_status="FAILED", accepted_candidate=False, metrics={"accepted_candidate": 0, "review_pass": 0, "judge_pass": 0, "runtime_seconds": 8})

    comparison = arena.compare(experiment)

    assert comparison.aggregate_metrics["CHAMPION"]["accepted_candidate_rate"] == 1
    assert comparison.aggregate_metrics["CHALLENGER"]["accepted_candidate_rate"] == 0.5
    assert comparison.paired_outcomes["BOTH_PASS"] == [str(cases[0].id)]
    assert comparison.paired_outcomes["CHAMPION_ONLY_PASS"] == [str(cases[1].id)]
    assert comparison.regression_cases == [str(cases[1].id)]
    assert comparison.verdict == "REJECT_RECOMMENDED"


def test_infrastructure_failures_are_excluded_from_quality_comparison(tmp_path: Path) -> None:
    arena, _, cases, experiment, champion, challenger = experiment_with_variants(tmp_path, count=1, minimum=1)
    champion_run = ReplayRun.objects.create(experiment=experiment, variant=champion, replay_case=cases[0], status="FAILED", failure_classification="INFRASTRUCTURE_FAILURE")
    ReplayResult.objects.create(replay_run=champion_run, completion_status="FAILED", accepted_candidate=False, metrics={"accepted_candidate": 0})
    challenger_run = ReplayRun.objects.create(experiment=experiment, variant=challenger, replay_case=cases[0], status="COMPLETE")
    ReplayResult.objects.create(replay_run=challenger_run, completion_status="COMPLETE", accepted_candidate=True, metrics={"accepted_candidate": 1})

    comparison = arena.compare(experiment)

    assert comparison.aggregate_metrics["CHAMPION"] == {"case_count": 0}
    assert comparison.verdict in {"INCONCLUSIVE", "RUN_MORE_REPLAYS"}


def test_verdicts_cover_promote_reject_inconclusive_and_more_replays(tmp_path: Path) -> None:
    arena, _, _, experiment, _, _ = experiment_with_variants(tmp_path, count=1, minimum=3)
    promote, _ = arena.judge_experiment(experiment, {"DELTA": {"accepted_candidate_rate": 0.2}}, {"case_count": 3, "CHAMPION_ONLY_PASS": []})
    reject, _ = arena.judge_experiment(experiment, {"DELTA": {"accepted_candidate_rate": 0.2}}, {"case_count": 3, "CHAMPION_ONLY_PASS": ["case"]})
    more, _ = arena.judge_experiment(experiment, {"DELTA": {"accepted_candidate_rate": 0.2}}, {"case_count": 1, "CHAMPION_ONLY_PASS": []})
    inconclusive, _ = arena.judge_experiment(experiment, {"DELTA": {"accepted_candidate_rate": 0.0, "median_runtime_seconds": 0.0}}, {"case_count": 3, "CHAMPION_ONLY_PASS": []})

    assert promote == "PROMOTE_RECOMMENDED"
    assert reject == "REJECT_RECOMMENDED"
    assert more == "RUN_MORE_REPLAYS"
    assert inconclusive == "INCONCLUSIVE"


def test_budget_cap_stops_experiment_cleanly(tmp_path: Path) -> None:
    arena, _, _, experiment, _, _ = experiment_with_variants(tmp_path, count=2)
    experiment.metadata = {"budget": {"maximum_replay_cases": 1, "maximum_model_requests": 0}}
    experiment.save(update_fields=["metadata", "updated_at"])

    arena.run_experiment(experiment)
    experiment.refresh_from_db()

    assert experiment.status == "FAILED"
    assert experiment.metadata["incomplete_reason"] == "budget_exceeded"


def test_graph_promotion_requires_approval_and_preserves_historical_lineage(tmp_path: Path) -> None:
    arena, _, cases, experiment, champion, challenger = experiment_with_variants(tmp_path, count=1, minimum=1)
    historical = GraphRun.objects.create(execution_graph_version=champion.execution_graph_version, current_node="complete", status=GraphRunStatus.COMPLETE)
    champion_run = ReplayRun.objects.create(experiment=experiment, variant=champion, replay_case=cases[0], status="COMPLETE")
    ReplayResult.objects.create(replay_run=champion_run, completion_status="COMPLETE", accepted_candidate=True, metrics={"accepted_candidate": 1, "runtime_seconds": 10})
    challenger_run = ReplayRun.objects.create(experiment=experiment, variant=challenger, replay_case=cases[0], status="COMPLETE")
    ReplayResult.objects.create(replay_run=challenger_run, completion_status="COMPLETE", accepted_candidate=True, metrics={"accepted_candidate": 1, "runtime_seconds": 5})
    comparison = arena.compare(experiment)

    assert challenger.execution_graph_version.status == ExecutionGraphVersionStatus.CHALLENGER
    arena.approve_promotion(comparison, actor="tester")
    historical.refresh_from_db()
    champion.execution_graph_version.refresh_from_db()
    challenger.execution_graph_version.refresh_from_db()

    assert champion.execution_graph_version.status == ExecutionGraphVersionStatus.RETIRED
    assert challenger.execution_graph_version.status == ExecutionGraphVersionStatus.CHAMPION
    assert historical.execution_graph_version_id == champion.execution_graph_version_id


def test_progeny_replay_experiment_graph_pauses_at_human_decision() -> None:
    spec = progeny_replay_experiment_graph_v1()
    definition = ExecutionGraphDefinition.objects.create(name=spec.name, graph_type=spec.graph_type)
    version = ExecutionGraphVersion.objects.create(graph=definition, version=1, status=ExecutionGraphVersionStatus.CHAMPION, graph_spec=spec.to_dict())
    dataset = ReplayArena().create_dataset("Pause Dataset")
    experiment = ProgenyExperiment.objects.create(replay_dataset=dataset, target_type="EXECUTION_GRAPH", hypothesis="Pause", status="COMPLETE")
    graph_run = GraphRun.objects.create(execution_graph_version=version, current_node="await_human_decision", metadata={"experiment_id": str(experiment.id)})
    runtime = NativeGraphRuntime(replay_experiment_registry(ReplayArena()))

    runtime.run_until_terminal_or_paused(graph_run)
    graph_run.refresh_from_db()
    assert graph_run.status == GraphRunStatus.PAUSED

    runtime.signal_now(str(graph_run.id), {"action": "approve", "actor": "tester"})
    runtime.run_until_terminal_or_paused(graph_run)
    graph_run.refresh_from_db()
    assert graph_run.status == GraphRunStatus.COMPLETE