from __future__ import annotations

import json
import shutil
from pathlib import Path

from django.core.management import call_command
from django.core.management.base import BaseCommand, CommandError

from agents.replay_arena import ReplayArena
from control_plane.agents.management.commands.seed_core_agents import Command as SeedAgentsCommand
from control_plane.agents.models import ImprovementCandidate
from control_plane.projects.models import Milestone, Project, ProjectPlan, Task, TaskStatus
from control_plane.resources.models import Resource
from graph.bootstrap import champion_task_execution_graph_v1
from model_router.providers import QwenProvider
from model_router.router import ModelRouter


class Command(BaseCommand):
    help = "Opt-in smoke test: run a tiny real-Qwen Replay Arena champion/challenger experiment."

    def add_arguments(self, parser):
        parser.add_argument("--cases", type=int, default=2)
        parser.add_argument("--root", default="/tmp/artifex-replay-arena-smoke")

    def handle(self, *args, **options):
        case_count = int(options["cases"])
        if case_count < 1:
            raise CommandError("--cases must be at least 1")
        resource = Resource.objects.filter(provider="local_inference", is_active=True).first()
        if resource is None:
            raise CommandError("No Qwen/local_inference resource configured. Run seed_spark_resources first.")
        SeedAgentsCommand().handle()
        arena = ReplayArena(ModelRouter({"qwen": QwenProvider(resource)}, persist_requests=True), test_command=["python", "manage.py", "test"])
        dataset = arena.create_dataset(
            "qwen-replay-arena-smoke",
            description="Real Qwen Replay Arena smoke",
            selection_criteria={"source": "bounded smoke", "case_count": case_count},
        )
        root = Path(options["root"])
        cases = []
        for index in range(case_count):
            repo = root / f"case-{index}"
            if repo.exists():
                shutil.rmtree(repo)
            call_command("create_disposable_django_repo", str(repo), verbosity=0)
            goal = 'Add a /health endpoint returning JSON {"status": "ok"} and add tests.'
            project = Project.objects.create(name=f"Replay Source {index}", goal=goal, repository_path=str(repo))
            plan = ProjectPlan.objects.create(project=project, version=1, goal=goal)
            milestone = Milestone.objects.create(project=project, plan=plan, key="R", title="Replay", goal="Replay")
            task = Task.objects.create(
                project=project,
                milestone=milestone,
                task_type="implementation",
                status=TaskStatus.COMPLETE,
                goal=goal,
                acceptance_criteria=["/health returns ok", "tests pass"],
            )
            cases.append(arena.add_case_from_task(dataset, task))
        arena.freeze_dataset(dataset)
        candidate = ImprovementCandidate.objects.create(
            target_type="EXECUTION_GRAPH",
            hypothesis="Insert deterministic static-analysis node before Reviewer to reduce rework risk",
            recommended_route="Progeny Graph Evolution",
        )
        experiment = arena.create_experiment(candidate, dataset, success_criteria={"minimum_replay_cases": 3})
        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())
        arena.run_experiment(experiment, max_cases=case_count)
        comparison = experiment.comparison
        payload = {
            "dataset_id": str(dataset.id),
            "dataset_version": dataset.version,
            "case_count": dataset.cases.count(),
            "frozen": dataset.status,
            "experiment_id": str(experiment.id),
            "hypothesis": experiment.hypothesis,
            "target_type": experiment.target_type,
            "champion": f"{champion.configuration_snapshot['graph']} v{champion.configuration_snapshot['version']}",
            "challenger": f"{challenger.configuration_snapshot['graph']} v{challenger.configuration_snapshot['version']}",
            "case_ids": [str(case.id) for case in cases],
            "source_task_ids": [str(case.source_task_id) for case in cases],
            "baseline_shas": [case.repository_baseline_ref for case in cases],
            "metrics": comparison.aggregate_metrics,
            "paired": comparison.paired_outcomes,
            "regressions": comparison.regression_cases,
            "verdict": comparison.verdict,
            "reasons": comparison.reasons,
        }
        self.stdout.write(json.dumps(payload, indent=2, default=str))