from __future__ import annotations

from dataclasses import dataclass

from control_plane.agents.models import Agent, AgentPlan, AgentVersion, BenchmarkRun, PromotionStatus
from control_plane.events.bus import EventBus
from control_plane.events.models import EventType


@dataclass(frozen=True)
class BenchmarkDecision:
    decision: str
    metrics: dict[str, float]


class ProgenyService:
    def __init__(self, bus: EventBus | None = None) -> None:
        self.bus = bus or EventBus()

    def create_candidate_from_plan(self, plan: AgentPlan) -> AgentVersion:
        agent = plan.agent or Agent.objects.create(name=plan.name, role=plan.role)
        next_version = (agent.versions.order_by("-version").first().version + 1) if agent.versions.exists() else 1
        version = AgentVersion.objects.create(
            agent=agent,
            version=next_version,
            model=plan.model,
            system_contract=plan.system_contract,
            capabilities=plan.capabilities,
            tools=plan.tools,
            permissions=plan.permissions,
            context_policy=plan.context_policy,
            workflow=plan.workflow,
            retry_policy={"max_retries": 2},
            evaluator=plan.success_criteria,
            promotion_status=PromotionStatus.CHALLENGER,
        )
        if agent.champion_version_id is None:
            agent.champion_version = version
            version.promotion_status = PromotionStatus.CHAMPION
            version.save(update_fields=["promotion_status", "updated_at"])
            agent.save(update_fields=["champion_version", "updated_at"])
        self.bus.publish(EventType.AGENT_CREATED, actor="progeny", payload={"agent": agent.name, "version": version.version})
        return version

    def replay_benchmark(self, champion: AgentVersion, challenger: AgentVersion, benchmark_set: list[dict[str, object]]) -> BenchmarkRun:
        metrics = self._score(challenger, benchmark_set)
        champion_metrics = self._score(champion, benchmark_set)
        decision = "PROMOTE" if metrics["completion_rate"] >= champion_metrics["completion_rate"] and metrics["test_pass_rate"] >= champion_metrics["test_pass_rate"] else "REJECT"
        return BenchmarkRun.objects.create(
            champion=champion,
            challenger=challenger,
            benchmark_set=benchmark_set,
            metrics={"champion": champion_metrics, "challenger": metrics},
            decision=decision,
        )

    def promote_or_reject(self, run: BenchmarkRun) -> BenchmarkDecision:
        challenger = run.challenger
        agent = challenger.agent
        challenger_metrics = run.metrics["challenger"]
        if run.decision == "PROMOTE":
            if agent.champion_version_id:
                old = agent.champion_version
                old.promotion_status = PromotionStatus.CANDIDATE
                old.save(update_fields=["promotion_status", "updated_at"])
            challenger.promotion_status = PromotionStatus.CHAMPION
            challenger.save(update_fields=["promotion_status", "updated_at"])
            agent.champion_version = challenger
            agent.save(update_fields=["champion_version", "updated_at"])
            self.bus.publish(EventType.AGENT_PROMOTED, actor="progeny", payload={"agent": agent.name, "version": challenger.version})
            return BenchmarkDecision("PROMOTED", challenger_metrics)
        challenger.promotion_status = PromotionStatus.REJECTED
        challenger.save(update_fields=["promotion_status", "updated_at"])
        return BenchmarkDecision("REJECTED", challenger_metrics)

    def _score(self, version: AgentVersion, benchmark_set: list[dict[str, object]]) -> dict[str, float]:
        if not benchmark_set:
            return {"completion_rate": 0.0, "test_pass_rate": 0.0, "review_acceptance": 0.0, "tokens": 0.0, "runtime": 0.0}
        base = 1.0 if "do not self-certify" in version.system_contract.lower() else 0.8
        return {
            "completion_rate": base,
            "test_pass_rate": base,
            "review_acceptance": base,
            "tokens": float(len(version.system_contract.split())),
            "runtime": float(len(benchmark_set)),
        }