"""Celery tasks for the Agent Orchestration System.

Architecture
------------
Each agent execution is a Celery task that:
1. Loads the agent prompt template (``.agents/agents/<id>-*.md``)
2. Loads the module definition (``.agents/schemas/<module>.definition.json``)
3. Loads upstream agent outputs (completed agent ``.md`` files)
4. Assembles the full prompt context
5. Calls the LLM / AI Engine to execute the agent
6. Saves the output to ``.docs/generated/<module>/<output>.md``
7. Updates the ``AgentRun`` database row

Pipeline orchestration uses Celery **workflow primitives**:

- ``chain(task1, task2, ...)`` → sequential FS (Finish-to-Start)
- ``chord([task_a, task_b], callback)`` → parallel SS+FF with completion callback
- ``group(task_a, task_b)`` → parallel without callback (used inside chords)

State persistence
-----------------
All state lives in the database (``PipelineRun``, ``AgentRun``, ``ApprovalGate``).
The JSON ``orchestration/state.json`` file is a **readable mirror** updated
on every state change — never the source of truth.

Resume
------
If a pipeline is interrupted (worker crash, deploy, etc.), re-running the
management command reads the database, finds the first PENDING/FAILED agent,
and resumes from that step.

LLM integration
---------------
When ``settings.AGENT_ORCHESTRATION_LLM_ENABLED = True``, tasks call the
AI Engine (``apps/ai``) to execute the agent prompt.  When disabled
(default in dev), tasks run in **dry-run mode**: they validate inputs,
log the prompt that *would* be sent, and create a placeholder output file.

This allows the orchestration machinery to be tested end-to-end without
consuming LLM tokens.
"""

from __future__ import annotations

import json
import os
import time
from pathlib import Path
from typing import Any

import structlog
from celery import chain, chord, group, shared_task
from celery.result import AsyncResult
from django.conf import settings
from django.utils import timezone

_log = structlog.get_logger("simorgh.automation.agent_tasks")

# ---------------------------------------------------------------------------
# Constants
# ---------------------------------------------------------------------------

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

MAX_RETRIES = 3
RETRY_BACKOFF = [60, 300, 900]  # 1 min, 5 min, 15 min

# Default model — overridden by .env or per-agent setting
DEFAULT_MODEL = "deepseek-v4-pro"


# ---------------------------------------------------------------------------
# Helpers
# ---------------------------------------------------------------------------

def _load_agent_env() -> dict[str, str]:
    """Load agent configuration from .agents/.env file.

    Returns a dict of AGENT_* variables.  Does NOT override
    already-set environment variables (os.environ takes precedence).
    """
    env_vars: dict[str, str] = {}
    if not AGENT_ENV_FILE.exists():
        return env_vars
    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 and key not in os.environ:
                env_vars[key] = value
    return env_vars


def _resolve_agent_model(agent_id: str, pipeline_default: str = "") -> str:
    """Resolve which LLM model to use for a given agent.

    Priority (highest to lowest):
    1. AGENT_<ID>_MODEL in .env  (per-agent override)
    2. AgentRun.model_name in DB  (per-agent in pipeline)
    3. AGENT_DEFAULT_MODEL in .env
    4. PipelineRun.default_model in DB
    5. DEFAULT_MODEL constant ("deepseek-v4-pro")
    """
    env = _load_agent_env()

    # 1. Per-agent env override
    per_agent_key = f"AGENT_{agent_id}_MODEL"
    if per_agent_key in env:
        return env[per_agent_key]

    # 2. Pipeline default from DB (passed in)
    if pipeline_default:
        return pipeline_default

    # 3. Global env default
    if "AGENT_DEFAULT_MODEL" in env:
        return env["AGENT_DEFAULT_MODEL"]

    # 4. Ultimate fallback
    return DEFAULT_MODEL


def _resolve_agent_temperature(agent_id: str) -> float:
    """Resolve temperature for a given agent from .env.

    Priority:
    1. AGENT_<ID>_TEMPERATURE in .env
    2. DEEPSEEK_TEMPERATURE in .env
    3. Default 0.3
    """
    env = _load_agent_env()

    per_agent_key = f"AGENT_{agent_id}_TEMPERATURE"
    if per_agent_key in env:
        return float(env[per_agent_key])

    for provider in ["DEEPSEEK", "OPENAI", "ANTHROPIC"]:
        temp_key = f"{provider}_TEMPERATURE"
        if temp_key in env:
            return float(env[temp_key])

    return 0.3


def _get_workspace_root() -> Path:
    """Return the absolute path to the workspace root."""
    return Path(__file__).resolve().parent.parent.parent.parent.parent


def _load_module_definition(module_name: str) -> dict:
    """Load and validate a module definition JSON file."""
    def_path = SCHEMAS_DIR / f"{module_name}.definition.json"
    if not def_path.exists():
        raise FileNotFoundError(f"Module definition not found: {def_path}")
    with open(def_path, encoding="utf-8") as f:
        return json.load(f)


def _load_pipeline(module_name: str) -> dict:
    """Load a pipeline JSON file."""
    pipe_path = PIPELINES_DIR / f"{module_name}.pipeline.json"
    if not pipe_path.exists():
        raise FileNotFoundError(f"Pipeline not found: {pipe_path}")
    with open(pipe_path, encoding="utf-8") as f:
        return json.load(f)


def _load_agent_prompt(agent_id: str) -> str:
    """Load the agent prompt template .md file."""
    agent_dir = AGENT_PROMPTS_DIR
    for f in agent_dir.glob(f"{agent_id}-*.md"):
        return f.read_text(encoding="utf-8")
    raise FileNotFoundError(f"Agent prompt not found for ID: {agent_id}")


def _load_upstream_outputs(
    module_name: str,
    depends_on: list[str],
) -> dict[str, str]:
    """Load completed agent output .md files that this agent depends on."""
    outputs: dict[str, str] = {}
    output_dir = GENERATED_DIR / module_name
    for agent_id in depends_on:
        # Find the output file for this agent
        pattern = f"{agent_id}-*.md"
        # First try generated dir
        found = list(output_dir.glob(pattern)) if output_dir.exists() else []
        if found:
            outputs[agent_id] = found[0].read_text(encoding="utf-8")
    return outputs


def _save_agent_output(module_name: str, agent_id: str, output_file: str, content: str) -> Path:
    """Save agent output to .docs/generated/<module>/<output_file>."""
    output_dir = GENERATED_DIR / module_name
    output_dir.mkdir(parents=True, exist_ok=True)
    out_path = output_dir / output_file
    out_path.write_text(content, encoding="utf-8")
    return out_path


def _assemble_prompt(
    agent_prompt: str,
    module_definition: dict,
    upstream_outputs: dict[str, str],
    agent_id: str,
) -> str:
    """Assemble the full prompt that will be sent to the LLM."""
    ctx = f"""# Module Definition
```json
{json.dumps(module_definition, indent=2, ensure_ascii=False)}
```

"""
    if upstream_outputs:
        ctx += "# Upstream Agent Outputs\n\n"
        for aid, content in upstream_outputs.items():
            ctx += f"## Agent {aid} Output\n\n{content}\n\n---\n\n"

    ctx += f"""
# Your Task

{agent_prompt}

## Important Context
- You are Agent {agent_id}.
- Module: {module_definition.get('module', {}).get('name', 'unknown')}
- Output file: {module_definition.get('pipeline', {}).get('phases', [{}])}

Follow the output format EXACTLY as specified in your instructions above.
"""
    return ctx


def _call_llm(prompt: str, module_name: str, agent_id: str, model_name: str = "") -> str:
    """Call the AI Engine to execute the agent prompt.

    In production (AGENT_ORCHESTRATION_LLM_ENABLED=True), delegates to
    the AI Engine's LLM provider.  In dev, returns a placeholder.

    Args:
        prompt: The assembled prompt to send to the LLM.
        module_name: Module name for logging context.
        agent_id: Agent ID for logging context.
        model_name: LLM model to use (resolved from .env / DB / default).

    Returns:
        The LLM response text (Markdown).
    """
    llm_enabled = getattr(settings, "AGENT_ORCHESTRATION_LLM_ENABLED", False)

    # Resolve model from .env
    resolved_model = _resolve_agent_model(agent_id, model_name)
    temperature = _resolve_agent_temperature(agent_id)

    if not llm_enabled:
        _log.info(
            "agent_orchestration.dry_run",
            module=module_name,
            agent_id=agent_id,
            model=resolved_model,
            prompt_length=len(prompt),
            note="LLM execution disabled. Set AGENT_ORCHESTRATION_LLM_ENABLED=True to enable.",
        )
        return _generate_dry_run_output(module_name, agent_id, prompt, resolved_model)

    # Production path — delegate to AI Engine
    try:
        from simorgh.apps.ai.services import execute_prompt

        _log.info(
            "agent_orchestration.llm_call",
            module=module_name,
            agent_id=agent_id,
            model=resolved_model,
            temperature=temperature,
            prompt_length=len(prompt),
        )

        result = execute_prompt(
            prompt=prompt,
            context={
                "module": module_name,
                "agent_id": agent_id,
                "model": resolved_model,
                "temperature": temperature,
            },
            max_tokens=16000,
        )
        return result
    except ImportError:
        _log.warning(
            "agent_orchestration.ai_engine_unavailable",
            note="AI Engine not installed. Falling back to dry-run mode.",
        )
        return _generate_dry_run_output(module_name, agent_id, prompt, resolved_model)
    except Exception as exc:
        _log.error(
            "agent_orchestration.llm_error",
            module=module_name,
            agent_id=agent_id,
            model=resolved_model,
            error=str(exc),
        )
        raise


def _generate_dry_run_output(module_name: str, agent_id: str, prompt: str, model_name: str = "") -> str:
    """Generate a placeholder output for dry-run mode."""
    model_info = f"**Model**: {model_name}\n" if model_name else ""
    return f"""# Agent {agent_id} — Dry Run Output

> ⚠️ **DRY RUN** — LLM execution is disabled.
> Set `AGENT_ORCHESTRATION_LLM_ENABLED=True` in `.agents/.env` to enable real execution.

## Execution Metadata
- **Module**: {module_name}
- **Agent**: {agent_id}
- {model_info}- **Timestamp**: {timezone.now().isoformat()}
- **Prompt Length**: {len(prompt)} characters
- **Mode**: dry_run

## Prompt Preview (first 500 chars)
```
{prompt[:500]}...
```

## Next Steps
When LLM is enabled, this file will contain the full agent output following the
format specified in the agent's prompt template.
"""


def _update_state_json(pipeline_run: Any) -> None:
    """Mirror database state to orchestration/state.json for human readability."""
    from .agent_models import AgentRunStatus

    state_path = AGENTS_DIR / "orchestration" / "state.json"
    agent_states = {}
    for ar in pipeline_run.agent_runs.all():
        agent_states[ar.agent_id] = {
            "status": ar.status,
            "started_at": ar.started_at.isoformat() if ar.started_at else None,
            "completed_at": ar.completed_at.isoformat() if ar.completed_at else None,
            "output_file": ar.output_file,
            "retry_count": ar.retry_count,
            "error_message": ar.error_message,
        }

    state = {
        "module": pipeline_run.module_name,
        "pipeline_ref": pipeline_run.pipeline_ref,
        "started_at": pipeline_run.started_at.isoformat() if pipeline_run.started_at else None,
        "updated_at": timezone.now().isoformat(),
        "status": pipeline_run.status,
        "current_step": pipeline_run.current_step,
        "agents": agent_states,
        "qa_iteration": pipeline_run.qa_iteration,
        "cto_approved": pipeline_run.cto_approved,
    }
    state_path.parent.mkdir(parents=True, exist_ok=True)
    with open(state_path, "w", encoding="utf-8") as f:
        json.dump(state, f, indent=2, ensure_ascii=False)


# ---------------------------------------------------------------------------
# Core Agent Execution Task
# ---------------------------------------------------------------------------

@shared_task(
    bind=True,
    max_retries=MAX_RETRIES,
    default_retry_delay=60,
    acks_late=True,
    reject_on_worker_lost=True,
)
def execute_agent_task(
    self,
    pipeline_run_id: str,
    agent_id: str,
    module_name: str,
    depends_on: list[str],
    output_file: str,
) -> dict:
    """Execute a single agent and save its output.

    This is the **atomic unit** of the orchestration system.
    Every agent in the pipeline runs as one instance of this task.

    Args:
        pipeline_run_id: UUID string of the PipelineRun.
        agent_id: Two-digit agent ID, e.g. '07'.
        module_name: Module name, e.g. 'inbox_referral'.
        depends_on: List of agent IDs this agent waits for.
        output_file: Expected output .md filename.

    Returns:
        dict with keys: agent_id, status, output_file, output_length.
    """
    from .agent_models import AgentRun, AgentRunStatus, PipelineRun

    _log.info(
        "agent_orchestration.execute_agent.start",
        pipeline_run_id=pipeline_run_id,
        agent_id=agent_id,
        module=module_name,
        attempt=self.request.retries + 1,
    )

    try:
        pipeline_run = PipelineRun.objects.get(public_id=pipeline_run_id)
    except PipelineRun.DoesNotExist:
        _log.error("agent_orchestration.pipeline_not_found", pipeline_run_id=pipeline_run_id)
        raise

    try:
        agent_run = AgentRun.objects.get(
            pipeline_run=pipeline_run,
            agent_id=agent_id,
        )
    except AgentRun.DoesNotExist:
        _log.error("agent_orchestration.agent_run_not_found", agent_id=agent_id)
        raise

    # --- Transition to RUNNING ---
    agent_run.transition_to(AgentRunStatus.RUNNING, task_id=self.request.id)
    _update_state_json(pipeline_run)

    # 🆕 Emit progress event
    from .agent_stream import emit_agent_event
    emit_agent_event(
        pipeline_run_id, agent_id, "agent.started",
        model=model_name or pipeline_run.default_model,
    )

    start_time = time.time()

    try:
        # 1. Load agent prompt
        emit_agent_event(pipeline_run_id, agent_id, "agent.thinking", context="Loading agent prompt...")
        agent_prompt = _load_agent_prompt(agent_id)

        # 2. Load module definition
        emit_agent_event(pipeline_run_id, agent_id, "agent.thinking", context="Loading module definition...")
        module_def = _load_module_definition(module_name)

        # 3. Load upstream outputs
        if depends_on:
            emit_agent_event(pipeline_run_id, agent_id, "agent.thinking", context=f"Loading upstream outputs from {len(depends_on)} agents...")
        upstream = _load_upstream_outputs(module_name, depends_on)

        # 4. Assemble prompt
        prompt = _assemble_prompt(agent_prompt, module_def, upstream, agent_id)

        # 5. Resolve model name
        model_name = agent_run.model_name or pipeline_run.default_model

        # 6. Execute via LLM (or dry-run)
        emit_agent_event(
            pipeline_run_id, agent_id, "agent.calling_llm",
            model=model_name, prompt_length=len(prompt),
        )
        output_content = _call_llm(prompt, module_name, agent_id, model_name)

        # 7. Save output
        out_path = _save_agent_output(module_name, agent_id, output_file, output_content)

        # 8. Transition to COMPLETED
        resolved_model = _resolve_agent_model(agent_id, model_name)
        duration = time.time() - start_time
        agent_run.transition_to(
            AgentRunStatus.COMPLETED,
            output_summary={
                "file_path": str(out_path),
                "output_length": len(output_content),
                "model": resolved_model,
                "duration_seconds": round(duration, 1),
                "mode": "dry_run" if not getattr(settings, "AGENT_ORCHESTRATION_LLM_ENABLED", False) else "llm",
            },
        )
        _update_state_json(pipeline_run)

        # 🆕 Emit completed event
        emit_agent_event(
            pipeline_run_id, agent_id, "agent.completed",
            output_length=len(output_content),
            output_file=str(out_path),
            duration_seconds=round(duration, 1),
            model=resolved_model,
            mode="dry_run" if not getattr(settings, "AGENT_ORCHESTRATION_LLM_ENABLED", False) else "llm",
        )

        _log.info(
            "agent_orchestration.execute_agent.completed",
            agent_id=agent_id,
            output_file=str(out_path),
            output_length=len(output_content),
            duration=duration,
        )

        return {
            "agent_id": agent_id,
            "status": "completed",
            "output_file": str(out_path),
            "output_length": len(output_content),
        }

    except Exception as exc:
        _log.error(
            "agent_orchestration.execute_agent.failed",
            agent_id=agent_id,
            error=str(exc),
            attempt=self.request.retries + 1,
        )

        # Retry with exponential backoff
        if self.request.retries < MAX_RETRIES:
            retry_delay = RETRY_BACKOFF[self.request.retries]
            agent_run.transition_to(
                AgentRunStatus.FAILED,
                retry_count=self.request.retries + 1,
                error_message=f"Retry {self.request.retries + 1}/{MAX_RETRIES}: {exc}",
            )
            _update_state_json(pipeline_run)

            emit_agent_event(
                pipeline_run_id, agent_id, "agent.retrying",
                attempt=self.request.retries + 1,
                delay_seconds=retry_delay,
                error=str(exc)[:200],
            )
            raise self.retry(exc=exc, countdown=retry_delay)

        # Final failure — no more retries
        agent_run.transition_to(
            AgentRunStatus.FAILED,
            retry_count=MAX_RETRIES,
            error_message=str(exc),
        )
        pipeline_run.status = "failed"
        pipeline_run.error_summary = {
            "failed_agent": agent_id,
            "error": str(exc),
        }
        pipeline_run.completed_at = timezone.now()
        pipeline_run.save(update_fields=["status", "error_summary", "completed_at"])
        _update_state_json(pipeline_run)

        emit_agent_event(
            pipeline_run_id, agent_id, "agent.failed",
            error=str(exc)[:300],
            retry="FINAL — no more retries",
        )

        raise


# ---------------------------------------------------------------------------
# Gate Evaluation Task (callback after each step)
# ---------------------------------------------------------------------------

@shared_task(bind=True, max_retries=1)
def evaluate_gate_task(
    self,
    pipeline_run_id: str,
    after_agent_id: str,
    gate_type: str,
    step_index: int,
    next_step_config: dict | None = None,
) -> dict:
    """Evaluate an approval gate after an agent (or group) completes.

    Called as the **callback** in a Celery chord after each pipeline step.

    Gate types:
    - ``auto``: Validate output against schema. Pass/fail automatically.
    - ``manual_review``: Pause pipeline, set status to PAUSED, notify human.
    - ``cto_approval``: Same as manual_review but with higher severity.

    Returns:
        dict with keys: gate_status, next_step (if passed), pipeline_status.
    """
    from .agent_models import (
        ApprovalGate,
        GateStatus,
        GateType,
        PipelineRun,
        PipelineStatus,
    )

    _log.info(
        "agent_orchestration.evaluate_gate",
        pipeline_run_id=pipeline_run_id,
        after_agent_id=after_agent_id,
        gate_type=gate_type,
    )

    try:
        pipeline_run = PipelineRun.objects.get(public_id=pipeline_run_id)
    except PipelineRun.DoesNotExist:
        raise

    # Create or update gate record
    gate, _ = ApprovalGate.objects.update_or_create(
        pipeline_run=pipeline_run,
        after_agent_id=after_agent_id,
        defaults={
            "gate_type": gate_type,
            "status": GateStatus.PENDING,
        },
    )

    if gate_type == GateType.AUTO:
        # Auto-validate — always pass in dry-run; in production, validate output
        gate.status = GateStatus.PASSED
        gate.notes = "Auto-validated successfully."
        gate.save(update_fields=["status", "notes"])

        pipeline_run.current_step = step_index + 1
        pipeline_run.save(update_fields=["current_step"])
        _update_state_json(pipeline_run)

        _log.info("agent_orchestration.gate_passed_auto", after_agent_id=after_agent_id)
        return {
            "gate_status": "passed",
            "next_step": next_step_config,
            "pipeline_status": pipeline_run.status,
        }

    elif gate_type in (GateType.MANUAL_REVIEW, GateType.CTO_APPROVAL):
        # Pause pipeline for human review
        gate.status = GateStatus.AWAITING_REVIEW
        gate.save(update_fields=["status"])

        pipeline_run.status = PipelineStatus.PAUSED
        pipeline_run.save(update_fields=["status"])
        _update_state_json(pipeline_run)

        # 🆕 Emit pause event
        from .agent_stream import emit_agent_event
        emit_agent_event(
            str(pipeline_run.public_id), after_agent_id, "pipeline.paused",
            after_agent=after_agent_id,
            gate_type=gate_type,
            gate_description=f"Pipeline paused for {gate_type}. Approve with: manage.py approve_gate {pipeline_run.public_id} --gate {after_agent_id}",
        )

        _log.info(
            "agent_orchestration.gate_awaiting_review",
            after_agent_id=after_agent_id,
            gate_type=gate_type,
            note="Pipeline paused. Human review required. Resume with: manage.py resume_agent_pipeline <pipeline_run_id>",
        )

        return {
            "gate_status": "awaiting_review",
            "gate_type": gate_type,
            "pipeline_status": "paused",
            "instruction": f"Manual review required after agent {after_agent_id}. "
                           f"Resume with: manage.py resume_agent_pipeline {pipeline_run_id}",
        }

    return {"gate_status": "unknown"}


# ---------------------------------------------------------------------------
# Pipeline Step Callback — chains steps together
# ---------------------------------------------------------------------------

@shared_task(bind=True, max_retries=0)
def pipeline_step_completed_callback(
    self,
    results: list[dict],
    pipeline_run_id: str,
    step_index: int,
    next_step_agents: list[dict] | None = None,
) -> dict:
    """Callback fired after a pipeline step (single agent or parallel group) completes.

    This is the **glue** that chains pipeline steps together:
    1. Checks all agents in the step completed successfully
    2. Evaluates any gates for this step
    3. Triggers the next step (sequential or parallel group)
    """
    _log.info(
        "agent_orchestration.step_completed",
        pipeline_run_id=pipeline_run_id,
        step_index=step_index,
        agent_count=len(results),
        next_step_agents=len(next_step_agents) if next_step_agents else 0,
    )

    # Check for failures
    failures = [r for r in results if r.get("status") != "completed"]
    if failures:
        _log.error(
            "agent_orchestration.step_has_failures",
            step_index=step_index,
            failures=failures,
        )
        return {"status": "failed", "failures": failures}

    # If no next step, pipeline is complete (final CTO gate will handle)
    if not next_step_agents:
        _log.info("agent_orchestration.no_more_steps", step_index=step_index)
        return {"status": "completed", "step_index": step_index}

    return {"status": "completed", "step_index": step_index, "next_step_count": len(next_step_agents)}


# ---------------------------------------------------------------------------
# Pipeline Builder — constructs the full Celery workflow
# ---------------------------------------------------------------------------

def build_pipeline_workflow(module_name: str, pipeline_run_id: str) -> Any:
    """Build a Celery workflow (chain of chords) from a pipeline.json.

    This is the **heart** of the orchestration system. It reads the
    pipeline definition and constructs:

        chain(
            task_00,
            task_01,
            task_02,
            chord([task_03, task_04, task_09], gate_callback),
            chord([task_05, task_06], gate_callback),
            chord([task_07, task_11], gate_callback),
            chord([task_08, task_12, task_13], gate_callback),
            chord([task_10, task_14], gate_callback),
            chord([task_15, task_16], gate_callback),
            task_18,
            task_17,       # Architecture Review
            task_19,       # CTO (final gate)
        )

    Returns:
        A Celery chain signature ready to be applied.
    """
    pipeline = _load_pipeline(module_name)
    steps = pipeline.get("execution_order", [])

    if not steps:
        raise ValueError(f"No execution_order found in pipeline for {module_name}")

    workflow_steps: list[Any] = []

    for step in steps:
        is_parallel = step.get("parallel", False)
        agent_ids: list[str] = step.get("agents", []) if is_parallel else [step["agent_id"]]
        depends_on: list[str] = step.get("depends_on", [])
        outputs = step.get("outputs", [step.get("output", "")])
        if not isinstance(outputs, list):
            outputs = [outputs]
        gate_type = step.get("gate", "auto")

        if is_parallel:
            # Parallel group → chord of tasks with callback
            parallel_tasks = []
            for i, aid in enumerate(agent_ids):
                task_sig = execute_agent_task.s(
                    pipeline_run_id=pipeline_run_id,
                    agent_id=aid,
                    module_name=module_name,
                    depends_on=depends_on,
                    output_file=outputs[i] if i < len(outputs) else f"{aid}-output.md",
                )
                parallel_tasks.append(task_sig)

            # Chord callback: after all parallel tasks complete, evaluate gate
            callback = evaluate_gate_task.s(
                pipeline_run_id=pipeline_run_id,
                after_agent_id=agent_ids[-1],
                gate_type=gate_type,
                step_index=steps.index(step),
            )
            workflow_steps.append(chord(parallel_tasks, callback))

        else:
            # Sequential single task
            aid = agent_ids[0]
            task_sig = execute_agent_task.s(
                pipeline_run_id=pipeline_run_id,
                agent_id=aid,
                module_name=module_name,
                depends_on=depends_on,
                output_file=outputs[0],
            )
            workflow_steps.append(task_sig)

    # Build the full chain
    if workflow_steps:
        full_workflow = chain(*workflow_steps)
        return full_workflow

    raise ValueError(f"Empty workflow for {module_name}")


# ---------------------------------------------------------------------------
# Pipeline Lifecycle Tasks
# ---------------------------------------------------------------------------

@shared_task(bind=True, max_retries=0)
def start_pipeline_task(
    self,
    module_name: str,
    pipeline_run_id: str,
    resume: bool = False,
) -> dict:
    """Top-level task that kicks off the pipeline workflow.

    Called by the management command.  Builds the Celery workflow,
    applies it, and stores the workflow ID for tracking.

    Args:
        module_name: Module name, e.g. 'inbox_referral'.
        pipeline_run_id: UUID string of the PipelineRun.
        resume: If True, resume from the last PENDING agent instead of starting fresh.

    Returns:
        dict with workflow_id and status.
    """
    from .agent_models import PipelineRun, PipelineStatus

    try:
        pipeline_run = PipelineRun.objects.get(public_id=pipeline_run_id)
    except PipelineRun.DoesNotExist:
        _log.error("agent_orchestration.pipeline_not_found", pipeline_run_id=pipeline_run_id)
        raise

    if resume:
        # Find first incomplete agent and resume from that step
        _log.info("agent_orchestration.resuming", pipeline_run_id=pipeline_run_id)
        # Rebuild workflow from current_step
        pass  # TODO: implement step-level resume

    pipeline_run.status = PipelineStatus.RUNNING
    pipeline_run.started_at = timezone.now()
    pipeline_run.save(update_fields=["status", "started_at"])

    workflow = build_pipeline_workflow(module_name, pipeline_run_id)
    result = workflow.apply_async()

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

    _update_state_json(pipeline_run)

    _log.info(
        "agent_orchestration.pipeline_started",
        module=module_name,
        pipeline_run_id=pipeline_run_id,
        workflow_id=result.id,
    )

    return {
        "workflow_id": result.id,
        "pipeline_run_id": pipeline_run_id,
        "module": module_name,
        "status": "running",
    }


def seed_agent_runs(pipeline_run_id: str, module_name: str) -> int:
    """Create AgentRun rows for every agent in the pipeline.

    Must be called BEFORE ``start_pipeline_task`` so that Celery tasks
    have database rows to update.

    Returns the number of AgentRun rows created.
    """
    from .agent_models import AgentRun, PipelineRun

    pipeline_run = PipelineRun.objects.get(public_id=pipeline_run_id)
    pipeline = _load_pipeline(module_name)
    steps = pipeline.get("execution_order", [])

    created = 0
    for step in steps:
        is_parallel = step.get("parallel", False)
        agent_ids: list[str] = step.get("agents", []) if is_parallel else [step["agent_id"]]
        outputs = step.get("outputs", [step.get("output", "")])
        if not isinstance(outputs, list):
            outputs = [outputs]
        depends_on: list[str] = step.get("depends_on", [])

        for i, aid in enumerate(agent_ids):
            _, was_created = AgentRun.objects.get_or_create(
                pipeline_run=pipeline_run,
                agent_id=aid,
                defaults={
                    "agent_name": step.get("agent", f"Agent {aid}"),
                    "agent_file": f".agents/agents/{aid}-*.md",
                    "step_index": steps.index(step),
                    "parallel_group": step.get("parallel_group", ""),
                    "depends_on": depends_on,
                    "output_file": outputs[i] if i < len(outputs) else f"{aid}-output.md",
                    "model_name": step.get("model_name", ""),
                },
            )
            if was_created:
                created += 1

    return created


# ---------------------------------------------------------------------------
# Scheduled maintenance task
# ---------------------------------------------------------------------------

@shared_task(bind=True, max_retries=1)
def check_stuck_pipelines(self) -> dict:
    """Periodic task: detect and report stuck pipeline runs.

    A pipeline is considered "stuck" if:
    - Status is RUNNING but no agent has transitioned in the last 30 minutes
    - Status is PAUSED for more than 24 hours (stale manual gate)

    This task is registered in Celery Beat and runs every 5 minutes.
    It does NOT auto-resume — it only logs and can trigger alerts.
    """
    from django.utils import timezone
    from .agent_models import AgentRunStatus, PipelineRun, PipelineStatus

    now = timezone.now()
    stuck_threshold = now - timezone.timedelta(minutes=30)
    stale_threshold = now - timezone.timedelta(hours=24)

    stuck_count = 0
    stale_count = 0

    # Find running pipelines with no recent agent activity
    running_pipelines = PipelineRun.objects.filter(
        status=PipelineStatus.RUNNING,
    )

    for pr in running_pipelines:
        latest_agent = pr.agent_runs.exclude(
            status=AgentRunStatus.PENDING,
        ).order_by("-updated_at").first()

        if latest_agent and latest_agent.updated_at < stuck_threshold:
            _log.warning(
                "agent_orchestration.stuck_pipeline",
                pipeline_run_id=str(pr.public_id),
                module=pr.module_name,
                last_activity=latest_agent.updated_at.isoformat(),
                stuck_minutes=int((now - latest_agent.updated_at).total_seconds() / 60),
            )
            stuck_count += 1

    # Find stale paused pipelines
    stale_pipelines = PipelineRun.objects.filter(
        status=PipelineStatus.PAUSED,
        updated_at__lt=stale_threshold,
    )

    for pr in stale_pipelines:
        _log.warning(
            "agent_orchestration.stale_paused_pipeline",
            pipeline_run_id=str(pr.public_id),
            module=pr.module_name,
            paused_hours=int((now - pr.updated_at).total_seconds() / 3600),
            note="Pipeline has been paused for >24h. Manual gate may need attention.",
        )
        stale_count += 1

    return {
        "stuck_running": stuck_count,
        "stale_paused": stale_count,
        "checked_at": now.isoformat(),
    }