"""Agent orchestration models for the Automation Engine.

Stores pipeline runs and per-agent execution state in the database
so Celery workers can resume after interruption and the orchestrator
can query progress without reading JSON files.

Migrations
----------
This module creates one table: ``automation_agentrun``.
A second migration adds ``automation_pipelinerun`` when the pipeline
runner is first used.
"""

from __future__ import annotations

from django.db import models
from django.utils.translation import gettext_lazy as _

from simorgh.core.models import TimeStampedModel, UUIDModel


# ---------------------------------------------------------------------------
# Choices
# ---------------------------------------------------------------------------

class PipelineStatus(models.TextChoices):
    IDLE      = "idle",      _("Idle — not started")
    RUNNING   = "running",   _("Running")
    PAUSED    = "paused",    _("Paused — awaiting manual gate")
    COMPLETED = "completed", _("Completed")
    FAILED    = "failed",    _("Failed")


class AgentRunStatus(models.TextChoices):
    PENDING   = "pending",   _("Pending — waiting for dependencies")
    RUNNING   = "running",   _("Running")
    COMPLETED = "completed", _("Completed successfully")
    FAILED    = "failed",    _("Failed")
    SKIPPED   = "skipped",   _("Skipped")


class GateType(models.TextChoices):
    AUTO           = "auto",           _("Automatic validation")
    MANUAL_REVIEW  = "manual_review",  _("Manual review required")
    CTO_APPROVAL   = "cto_approval",   _("CTO approval required")


class GateStatus(models.TextChoices):
    PENDING         = "pending",         _("Pending")
    PASSED          = "passed",          _("Passed")
    FAILED          = "failed",          _("Failed")
    AWAITING_REVIEW = "awaiting_review", _("Awaiting human review")


# ---------------------------------------------------------------------------
# PipelineRun — one per module pipeline execution
# ---------------------------------------------------------------------------

class PipelineRun(UUIDModel, TimeStampedModel):
    """Top-level record of a pipeline execution for one module.

    Created when ``run_agent_pipeline`` management command is invoked.
    Tracks overall status, current step, and QA iteration count.
    """

    module_name = models.CharField(
        max_length=128,
        db_index=True,
        help_text="Module name, e.g. 'inbox_referral'.",
    )
    pipeline_ref = models.CharField(
        max_length=512,
        help_text="Relative path to pipeline JSON, e.g. '.agents/pipelines/inbox-referral.pipeline.json'.",
    )
    status = models.CharField(
        max_length=20,
        choices=PipelineStatus.choices,
        default=PipelineStatus.IDLE,
        db_index=True,
    )
    current_step = models.PositiveSmallIntegerField(
        default=0,
        help_text="0-based index of the currently executing step.",
    )
    qa_iteration = models.PositiveSmallIntegerField(
        default=0,
        help_text="Current QA rejection loop iteration (0-3).",
    )
    cto_approved = models.BooleanField(
        default=False,
        help_text="Whether CTO has approved the final review.",
    )
    celery_workflow_id = models.CharField(
        max_length=255,
        blank=True,
        default="",
        help_text="Celery workflow ID (chain/chord) for cancellation/tracking.",
    )
    default_model = models.CharField(
        max_length=128,
        blank=True,
        default="deepseek-v4-pro",
        help_text="Default LLM model for all agents in this pipeline, e.g. 'deepseek-v4-pro', 'gpt-4o', 'claude-3.5-sonnet'.",
    )
    started_at = models.DateTimeField(null=True, blank=True)
    completed_at = models.DateTimeField(null=True, blank=True)
    error_summary = models.JSONField(
        default=dict,
        blank=True,
        help_text="Aggregated error info if status is 'failed'.",
    )

    class Meta:
        db_table = "automation_pipelinerun"
        ordering = ["-created_at"]
        verbose_name = _("Pipeline Run")
        verbose_name_plural = _("Pipeline Runs")

    def __str__(self) -> str:
        return f"PipelineRun({self.module_name} — {self.status})"


# ---------------------------------------------------------------------------
# AgentRun — one row per agent execution inside a pipeline
# ---------------------------------------------------------------------------

class AgentRun(UUIDModel, TimeStampedModel):
    """Immutable record of a single agent's execution within a pipeline run.

    One row per (pipeline_run, agent_id).  The unique constraint ensures
    a given agent is only executed once per pipeline run.

    Celery tasks update this row as execution progresses:
    ``PENDING → RUNNING → COMPLETED | FAILED``.
    """

    pipeline_run = models.ForeignKey(
        PipelineRun,
        on_delete=models.CASCADE,
        related_name="agent_runs",
    )
    agent_id = models.CharField(
        max_length=3,
        help_text="Two-digit agent ID, e.g. '01', '07', '19'.",
    )
    agent_name = models.CharField(
        max_length=128,
        help_text="Human-readable agent name, e.g. 'Domain Architect'.",
    )
    agent_file = models.CharField(
        max_length=256,
        help_text="Relative path to agent .md file, e.g. '.agents/agents/01-domain-architect.md'.",
    )
    step_index = models.PositiveSmallIntegerField(
        help_text="Pipeline step index this agent belongs to.",
    )
    parallel_group = models.CharField(
        max_length=32,
        blank=True,
        default="",
        help_text="Parallel group label, e.g. 'group-a'. Empty for sequential.",
    )
    depends_on = models.JSONField(
        default=list,
        help_text="List of agent IDs this agent waits for.",
    )
    output_file = models.CharField(
        max_length=512,
        help_text="Expected output .md file path.",
    )
    model_name = models.CharField(
        max_length=128,
        blank=True,
        default="",
        help_text="LLM model for this agent. Overrides PipelineRun.default_model. e.g. 'deepseek-v4-pro', 'gpt-4o'. Empty = use pipeline default.",
    )
    status = models.CharField(
        max_length=20,
        choices=AgentRunStatus.choices,
        default=AgentRunStatus.PENDING,
        db_index=True,
    )
    retry_count = models.PositiveSmallIntegerField(
        default=0,
        help_text="Number of retry attempts (max 3).",
    )
    task_id = models.CharField(
        max_length=255,
        blank=True,
        default="",
        help_text="Celery task ID for this agent's execution.",
    )
    started_at = models.DateTimeField(null=True, blank=True)
    completed_at = models.DateTimeField(null=True, blank=True)
    error_message = models.TextField(
        blank=True,
        default="",
        help_text="Last error message if status is 'failed'.",
    )
    output_summary = models.JSONField(
        default=dict,
        blank=True,
        help_text="Lightweight metadata about the output (line count, validation result).",
    )

    class Meta:
        db_table = "automation_agentrun"
        ordering = ["pipeline_run", "step_index", "agent_id"]
        constraints = [
            models.UniqueConstraint(
                fields=["pipeline_run", "agent_id"],
                name="uq_agentrun_pipeline_agent",
            ),
        ]
        verbose_name = _("Agent Run")
        verbose_name_plural = _("Agent Runs")

    def __str__(self) -> str:
        return f"AgentRun({self.agent_id} {self.agent_name} — {self.status})"

    def transition_to(self, new_status: AgentRunStatus, **kwargs) -> None:
        """Safe state transition with timestamp bookkeeping."""
        self.status = new_status
        if new_status == AgentRunStatus.RUNNING and not self.started_at:
            from django.utils import timezone
            self.started_at = timezone.now()
        if new_status in (AgentRunStatus.COMPLETED, AgentRunStatus.FAILED):
            from django.utils import timezone
            self.completed_at = timezone.now()
        for field, value in kwargs.items():
            setattr(self, field, value)
        self.save(update_fields=["status", "started_at", "completed_at", *kwargs.keys()])


# ---------------------------------------------------------------------------
# ApprovalGate — records gate evaluation results
# ---------------------------------------------------------------------------

class ApprovalGate(UUIDModel, TimeStampedModel):
    """Record of a gate evaluation within a pipeline run."""

    pipeline_run = models.ForeignKey(
        PipelineRun,
        on_delete=models.CASCADE,
        related_name="gates",
    )
    after_agent_id = models.CharField(
        max_length=3,
        help_text="Agent ID after which this gate fires.",
    )
    gate_type = models.CharField(
        max_length=20,
        choices=GateType.choices,
    )
    status = models.CharField(
        max_length=20,
        choices=GateStatus.choices,
        default=GateStatus.PENDING,
    )
    description = models.TextField(blank=True, default="")
    reviewed_by = models.CharField(
        max_length=255,
        blank=True,
        default="",
        help_text="Human reviewer identifier.",
    )
    reviewed_at = models.DateTimeField(null=True, blank=True)
    notes = models.TextField(blank=True, default="")

    class Meta:
        db_table = "automation_approvalgate"
        ordering = ["pipeline_run", "created_at"]
        verbose_name = _("Approval Gate")
        verbose_name_plural = _("Approval Gates")

    def __str__(self) -> str:
        return f"Gate({self.after_agent_id} {self.gate_type} — {self.status})"