from __future__ import annotations

from django.db import models

from control_plane.common import TimestampedModel


class AgentRole(models.TextChoices):
    PROJECT_ARCHAEOLOGIST = "PROJECT_ARCHAEOLOGIST"
    PLANNER = "PLANNER"
    CODER = "CODER"
    REVIEWER = "REVIEWER"
    PROJECT_JUDGE = "PROJECT_JUDGE"


class PromotionStatus(models.TextChoices):
    CANDIDATE = "CANDIDATE"
    CHALLENGER = "CHALLENGER"
    CHAMPION = "CHAMPION"
    REJECTED = "REJECTED"


class Agent(TimestampedModel):
    name = models.CharField(max_length=200, unique=True)
    role = models.CharField(max_length=80, choices=AgentRole.choices)
    description = models.TextField(blank=True)
    champion_version = models.ForeignKey(
        "AgentVersion", on_delete=models.SET_NULL, null=True, blank=True, related_name="championed_by"
    )


class AgentVersion(TimestampedModel):
    agent = models.ForeignKey(Agent, on_delete=models.CASCADE, related_name="versions")
    version = models.PositiveIntegerField()
    model = models.CharField(max_length=120)
    system_contract = models.TextField()
    capabilities = models.JSONField(default=list, blank=True)
    tools = models.JSONField(default=list, blank=True)
    permissions = models.JSONField(default=dict, blank=True)
    context_policy = models.JSONField(default=dict, blank=True)
    workflow = models.JSONField(default=dict, blank=True)
    retry_policy = models.JSONField(default=dict, blank=True)
    evaluator = models.JSONField(default=dict, blank=True)
    benchmark_status = models.CharField(max_length=32, default="UNBENCHMARKED")
    promotion_status = models.CharField(max_length=32, choices=PromotionStatus.choices, default=PromotionStatus.CANDIDATE)

    class Meta:
        constraints = [models.UniqueConstraint(fields=["agent", "version"], name="unique_agent_version")]


class AgentPlan(TimestampedModel):
    agent = models.ForeignKey(Agent, on_delete=models.CASCADE, related_name="plans", null=True, blank=True)
    name = models.CharField(max_length=200)
    role = models.CharField(max_length=80, choices=AgentRole.choices)
    model = models.CharField(max_length=120)
    system_contract = models.TextField()
    capabilities = models.JSONField(default=list, blank=True)
    tools = models.JSONField(default=list, blank=True)
    permissions = models.JSONField(default=dict, blank=True)
    context_policy = models.JSONField(default=dict, blank=True)
    workflow = models.JSONField(default=dict, blank=True)
    benchmarks = models.JSONField(default=list, blank=True)
    success_criteria = models.JSONField(default=dict, blank=True)
    created_by = models.CharField(max_length=120, default="sol")


class AgentRun(TimestampedModel):
    agent_version = models.ForeignKey(AgentVersion, on_delete=models.PROTECT, related_name="runs")
    project = models.ForeignKey("projects.Project", on_delete=models.CASCADE, null=True, blank=True)
    task = models.ForeignKey("projects.Task", on_delete=models.SET_NULL, null=True, blank=True)
    status = models.CharField(max_length=32, default="QUEUED")
    input_contract = models.JSONField(default=dict, blank=True)
    output_contract = models.JSONField(default=dict, blank=True)
    metrics = models.JSONField(default=dict, blank=True)


class BenchmarkRun(TimestampedModel):
    champion = models.ForeignKey(AgentVersion, on_delete=models.PROTECT, related_name="champion_benchmarks")
    challenger = models.ForeignKey(AgentVersion, on_delete=models.PROTECT, related_name="challenger_benchmarks")
    benchmark_set = models.JSONField(default=list, blank=True)
    metrics = models.JSONField(default=dict, blank=True)
    decision = models.CharField(max_length=32, default="PENDING")