"""
Management command to run the Agent Orchestration Pipeline via Celery.

Usage
-----
.. code-block:: sh

    # Start a full pipeline for a module
    python manage.py run_agent_pipeline inbox_referral

    # Dry-run (validate without dispatching Celery tasks)
    python manage.py run_agent_pipeline inbox_referral --dry-run

    # Resume a paused pipeline (after manual gate review)
    python manage.py run_agent_pipeline inbox_referral --resume

    # Resume a specific pipeline run by UUID
    python manage.py run_agent_pipeline --resume --pipeline-run-id <uuid>

    # Check status of a pipeline
    python manage.py run_agent_pipeline inbox_referral --status

    # List all pipeline runs
    python manage.py run_agent_pipeline --list

    # Retry failed agents in a pipeline
    python manage.py run_agent_pipeline inbox_referral --retry-failed

How it works
------------
1. Reads ``.agents/schemas/<module>.definition.json``
2. Reads ``.agents/pipelines/<module>.pipeline.json``
3. Creates ``PipelineRun`` and ``AgentRun`` database rows
4. Builds a Celery ``chain`` of ``chord`` workflows
5. Dispatches the workflow asynchronously
6. Each agent task loads its prompt template, upstream outputs,
   calls the LLM (or dry-run), and saves output to
   ``.docs/generated/<module>/``
"""

from __future__ import annotations

import json
import uuid
from pathlib import Path
from typing import Any

from django.core.management.base import BaseCommand, CommandError
from django.utils import timezone

AGENTS_DIR = Path(__file__).resolve().parent.parent.parent.parent.parent.parent / ".agents"
SCHEMAS_DIR = AGENTS_DIR / "schemas"
PIPELINES_DIR = AGENTS_DIR / "pipelines"
AGENT_ENV_FILE = AGENTS_DIR / ".env"


def _load_agent_env() -> dict[str, str]:
    """Load agent .env file into a dict."""
    env = {}
    if not AGENT_ENV_FILE.exists():
        return env
    with open(AGENT_ENV_FILE, encoding="utf-8") as f:
        for line in f:
            line = line.strip()
            if not line or line.startswith("#") or "=" not in line:
                continue
            key, _, value = line.partition("=")
            key = key.strip()
            value = value.strip().strip('"').strip("'")
            if key and value:
                env[key] = value
    return env


class Command(BaseCommand):
    help = __doc__

    def add_arguments(self, parser):
        parser.add_argument(
            "module",
            nargs="?",
            help="Module name, e.g. 'inbox_referral'.",
        )
        parser.add_argument(
            "--dry-run",
            action="store_true",
            help="Validate pipeline without dispatching Celery tasks.",
        )
        parser.add_argument(
            "--resume",
            action="store_true",
            help="Resume a paused or interrupted pipeline.",
        )
        parser.add_argument(
            "--pipeline-run-id",
            help="Resume a specific pipeline run by UUID.",
        )
        parser.add_argument(
            "--status",
            action="store_true",
            help="Show status of the latest pipeline run for the module.",
        )
        parser.add_argument(
            "--list",
            action="store_true",
            help="List all pipeline runs.",
        )
        parser.add_argument(
            "--retry-failed",
            action="store_true",
            help="Retry all failed agents in the latest pipeline run.",
        )

    def handle(self, **options):
        if options["list"]:
            self._list_pipelines()
            return

        if options["status"]:
            if not options["module"]:
                raise CommandError("--status requires a module name.")
            self._show_status(options["module"])
            return

        if options["resume"]:
            self._resume_pipeline(options)
            return

        if options["retry_failed"]:
            if not options["module"]:
                raise CommandError("--retry-failed requires a module name.")
            self._retry_failed(options["module"])
            return

        if not options["module"]:
            raise CommandError("Module name is required. Usage: manage.py run_agent_pipeline <module>")

        module_name = options["module"]
        self._start_pipeline(module_name, dry_run=options["dry_run"])

    # ------------------------------------------------------------------
    # Start pipeline
    # ------------------------------------------------------------------

    def _start_pipeline(self, module_name: str, dry_run: bool = False) -> None:
        """Validate inputs, create DB records, dispatch Celery workflow."""
        # Validate module definition exists
        def_path = SCHEMAS_DIR / f"{module_name}.definition.json"
        if not def_path.exists():
            raise CommandError(
                f"Module definition not found: {def_path}\n"
                f"Create it at .agents/schemas/{module_name}.definition.json"
            )

        # Validate pipeline exists
        pipe_path = PIPELINES_DIR / f"{module_name}.pipeline.json"
        if not pipe_path.exists():
            raise CommandError(
                f"Pipeline not found: {pipe_path}\n"
                f"Create it at .agents/pipelines/{module_name}.pipeline.json"
            )

        # Load and validate
        with open(def_path, encoding="utf-8") as f:
            module_def = json.load(f)
        with open(pipe_path, encoding="utf-8") as f:
            pipeline = json.load(f)

        module_info = module_def.get("module", {})
        steps = pipeline.get("execution_order", [])

        if not steps:
            raise CommandError(f"Pipeline '{module_name}' has no execution_order steps.")

        if dry_run:
            self._dry_run_validation(module_name, module_info, steps)
            return

        # Resolve default model: module definition > .env > hardcoded default
        agent_env = _load_agent_env()
        default_model = (
            module_info.get("default_model", "")
            or agent_env.get("AGENT_DEFAULT_MODEL", "")
            or "deepseek-v4-pro"
        )

        # Create PipelineRun
        from simorgh.apps.automation.agent_models import PipelineRun, PipelineStatus

        pipeline_run = PipelineRun.objects.create(
            module_name=module_name,
            pipeline_ref=f".agents/pipelines/{module_name}.pipeline.json",
            status=PipelineStatus.IDLE,
            default_model=default_model,
        )

        # Seed AgentRun rows
        from simorgh.apps.automation.agent_tasks import seed_agent_runs

        count = seed_agent_runs(str(pipeline_run.public_id), module_name)

        self.stdout.write(self.style.SUCCESS(
            f"\n{'='*60}\n"
            f"  Pipeline Run Created\n"
            f"  {'='*60}\n"
            f"  Module:      {module_name}\n"
            f"  Model:       {default_model}\n"
            f"  Description: {module_info.get('description', 'N/A')}\n"
            f"  Pipeline ID: {pipeline_run.public_id}\n"
            f"  Agents:      {count}\n"
            f"  Steps:       {len(steps)}\n"
            f"  {'='*60}\n"
        ))

        # Show agent list
        self.stdout.write("  Agents in execution order:\n")
        for step in steps:
            is_parallel = step.get("parallel", False)
            agent_ids = step.get("agents", [step["agent_id"]]) if is_parallel else [step["agent_id"]]
            label = "∥" if is_parallel else "→"
            self.stdout.write(f"    {label} {', '.join(agent_ids):20s}  {step.get('description', '')[:60]}")

        # Dispatch Celery workflow
        from simorgh.apps.automation.agent_tasks import start_pipeline_task

        self.stdout.write(f"\n  ⏳ Dispatching Celery workflow...\n")

        result = start_pipeline_task.delay(
            module_name=module_name,
            pipeline_run_id=str(pipeline_run.public_id),
        )

        self.stdout.write(self.style.SUCCESS(
            f"  ✅ Pipeline dispatched!\n"
            f"  Celery Workflow ID: {result.id}\n"
            f"  Monitor: celery -A config events\n"
            f"  Status:  manage.py run_agent_pipeline {module_name} --status\n"
            f"  Resume:  manage.py run_agent_pipeline {module_name} --resume\n"
            f"{'='*60}\n"
        ))

    # ------------------------------------------------------------------
    # Dry-run validation
    # ------------------------------------------------------------------

    def _dry_run_validation(self, module_name: str, module_info: dict, steps: list) -> None:
        """Validate pipeline structure without dispatching tasks."""
        self.stdout.write(self.style.NOTICE(f"\n🔍 DRY RUN — Validating pipeline for '{module_name}'\n"))

        errors: list[str] = []
        warnings: list[str] = []
        agent_ids_seen: set[str] = set()

        for i, step in enumerate(steps):
            is_parallel = step.get("parallel", False)
            agent_ids = step.get("agents", [step["agent_id"]]) if is_parallel else [step["agent_id"]]
            depends_on = step.get("depends_on", [])

            for aid in agent_ids:
                # Check agent prompt file exists
                prompt_files = list(AGENTS_DIR.glob(f"agents/{aid}-*.md"))
                if not prompt_files:
                    errors.append(f"Step {i}: Agent {aid} prompt file not found in .agents/agents/")
                else:
                    self.stdout.write(f"  ✅ Agent {aid}: {prompt_files[0].name}")

                # Check dependency order
                for dep in depends_on:
                    if dep not in agent_ids_seen:
                        errors.append(f"Step {i}: Agent {aid} depends on {dep} which hasn't run yet")

                # Check for missing agents
                if aid in agent_ids_seen:
                    warnings.append(f"Step {i}: Agent {aid} appears more than once")
                agent_ids_seen.add(aid)

            # Check parallel group consistency
            if is_parallel:
                pg = step.get("parallel_group", "")
                if not pg:
                    warnings.append(f"Step {i}: Parallel step has no parallel_group name")
                for aid in agent_ids:
                    for dep in depends_on:
                        if dep in agent_ids:
                            errors.append(f"Step {i}: Agent {aid} in parallel group depends on co-group agent {dep}")

        # Summary
        self.stdout.write(f"\n  Agents found: {len(agent_ids_seen)}")
        self.stdout.write(f"  Steps: {len(steps)}")
        self.stdout.write(f"  Errors: {len(errors)}")
        self.stdout.write(f"  Warnings: {len(warnings)}")

        if errors:
            self.stdout.write(self.style.ERROR("\n❌ ERRORS:"))
            for e in errors:
                self.stdout.write(f"  • {e}")
            raise CommandError("Pipeline validation failed. Fix errors above.")

        if warnings:
            self.stdout.write(self.style.WARNING("\n⚠️  WARNINGS:"))
            for w in warnings:
                self.stdout.write(f"  • {w}")

        self.stdout.write(self.style.SUCCESS(
            "\n✅ Pipeline validation passed! Ready to run:\n"
            f"   manage.py run_agent_pipeline {module_name}\n"
        ))

    # ------------------------------------------------------------------
    # Resume pipeline
    # ------------------------------------------------------------------

    def _resume_pipeline(self, options: dict) -> None:
        """Resume a paused or interrupted pipeline."""
        from simorgh.apps.automation.agent_models import (
            AgentRun,
            AgentRunStatus,
            PipelineRun,
            PipelineStatus,
        )

        pipeline_run_id = options.get("pipeline_run_id")
        module_name = options.get("module")

        if pipeline_run_id:
            try:
                pr = PipelineRun.objects.get(public_id=pipeline_run_id)
            except PipelineRun.DoesNotExist:
                raise CommandError(f"Pipeline run not found: {pipeline_run_id}")
        elif module_name:
            pr = PipelineRun.objects.filter(
                module_name=module_name,
            ).order_by("-created_at").first()
            if not pr:
                raise CommandError(f"No pipeline runs found for module '{module_name}'")
        else:
            raise CommandError("Provide --module or --pipeline-run-id to resume.")

        if pr.status not in (PipelineStatus.PAUSED, PipelineStatus.FAILED):
            raise CommandError(
                f"Pipeline {pr.public_id} is '{pr.status}', not 'paused' or 'failed'. "
                f"Cannot resume."
            )

        self.stdout.write(f"Resuming pipeline {pr.public_id} (module={pr.module_name}, status={pr.status})")

        # Re-dispatch from current step
        from simorgh.apps.automation.agent_tasks import build_pipeline_workflow

        pr.status = PipelineStatus.RUNNING
        pr.save(update_fields=["status"])

        workflow = build_pipeline_workflow(pr.module_name, str(pr.public_id))
        result = workflow.apply_async()

        pr.celery_workflow_id = result.id
        pr.save(update_fields=["celery_workflow_id"])

        self.stdout.write(self.style.SUCCESS(
            f"✅ Pipeline resumed! Workflow ID: {result.id}"
        ))

    # ------------------------------------------------------------------
    # Show status
    # ------------------------------------------------------------------

    def _show_status(self, module_name: str) -> None:
        """Display pipeline run status with per-agent progress."""
        from simorgh.apps.automation.agent_models import PipelineRun

        pr = PipelineRun.objects.filter(
            module_name=module_name,
        ).order_by("-created_at").first()

        if not pr:
            self.stdout.write(self.style.WARNING(f"No pipeline runs found for '{module_name}'."))
            return

        self.stdout.write(f"\n{'='*60}")
        self.stdout.write(f"  Pipeline: {pr.module_name}")
        self.stdout.write(f"  Run ID:   {pr.public_id}")
        self.stdout.write(f"  Model:    {pr.default_model}")
        self.stdout.write(f"  Status:   {pr.status}")
        self.stdout.write(f"  Step:     {pr.current_step}")
        self.stdout.write(f"  QA Loop:  {pr.qa_iteration}/3")
        self.stdout.write(f"  CTO:      {'✅ Approved' if pr.cto_approved else '⏳ Pending'}")
        self.stdout.write(f"  Started:  {pr.started_at}")
        self.stdout.write(f"  {'='*60}\n")

        # Per-agent status
        status_icons = {
            "pending": "⏳",
            "running": "🔄",
            "completed": "✅",
            "failed": "❌",
            "skipped": "⏭️",
        }

        for ar in pr.agent_runs.all():
            icon = status_icons.get(ar.status, "❓")
            pg = f" [{ar.parallel_group}]" if ar.parallel_group else ""
            model_info = f" [{ar.model_name}]" if ar.model_name else f" [{pr.default_model}]"
            self.stdout.write(
                f"  {icon} Agent {ar.agent_id}: {ar.agent_name:30s} {ar.status:12s}{pg}{model_info}"
            )
            if ar.error_message:
                self.stdout.write(f"     ↳ Error: {ar.error_message[:100]}")

        self.stdout.write("")

    # ------------------------------------------------------------------
    # List pipelines
    # ------------------------------------------------------------------

    def _list_pipelines(self) -> None:
        """List all pipeline runs."""
        from simorgh.apps.automation.agent_models import PipelineRun

        runs = PipelineRun.objects.all().order_by("-created_at")[:20]
        if not runs:
            self.stdout.write("No pipeline runs found.")
            return

        self.stdout.write(f"\n{'Module':25s} {'Run ID':38s} {'Status':12s} {'Started'}")
        self.stdout.write("-" * 95)
        for pr in runs:
            started = pr.started_at.strftime("%Y-%m-%d %H:%M") if pr.started_at else "—"
            self.stdout.write(
                f"{pr.module_name:25s} {str(pr.public_id):38s} {pr.status:12s} {started}"
            )
        self.stdout.write("")

    # ------------------------------------------------------------------
    # Retry failed
    # ------------------------------------------------------------------

    def _retry_failed(self, module_name: str) -> None:
        """Retry all failed agents in the latest pipeline run."""
        from simorgh.apps.automation.agent_models import AgentRunStatus, PipelineRun, PipelineStatus
        from simorgh.apps.automation.agent_tasks import execute_agent_task

        pr = PipelineRun.objects.filter(
            module_name=module_name,
        ).order_by("-created_at").first()

        if not pr:
            raise CommandError(f"No pipeline runs found for '{module_name}'.")

        failed_agents = pr.agent_runs.filter(status=AgentRunStatus.FAILED)
        if not failed_agents.exists():
            self.stdout.write(self.style.SUCCESS("No failed agents to retry."))
            return

        self.stdout.write(f"Retrying {failed_agents.count()} failed agent(s)...\n")

        pr.status = PipelineStatus.RUNNING
        pr.save(update_fields=["status"])

        for ar in failed_agents:
            self.stdout.write(f"  🔄 Retrying Agent {ar.agent_id} ({ar.agent_name})...")
            task = execute_agent_task.delay(
                pipeline_run_id=str(pr.public_id),
                agent_id=ar.agent_id,
                module_name=module_name,
                depends_on=ar.depends_on,
                output_file=ar.output_file,
            )
            ar.task_id = task.id
            ar.retry_count += 1
            ar.status = AgentRunStatus.PENDING
            ar.save(update_fields=["task_id", "retry_count", "status"])

        self.stdout.write(self.style.SUCCESS(f"\n✅ {failed_agents.count()} agent(s) re-queued."))