"""
Pipeline and Stage Definitions.

Pydantic models for defining workflows with LLM assistance integration.
Strictly typed, immutable, and serializable for crash recovery.
"""

from __future__ import annotations

import uuid
from datetime import datetime, timezone
from enum import Enum
from typing import Any, Callable, Coroutine

from pydantic import BaseModel, Field, model_validator


class PipelineStatus(str, Enum):
    """Pipeline execution status."""

    PENDING = "pending"
    RUNNING = "running"
    PAUSED = "paused"
    COMPLETED = "completed"
    FAILED = "failed"
    CANCELLED = "cancelled"


class StageStatus(str, Enum):
    """Stage execution status."""

    PENDING = "pending"
    RUNNING = "running"
    SKIPPED = "skipped"
    COMPLETED = "completed"
    FAILED = "failed"


class LLMAssistanceMode(str, Enum):
    """
    How LLM assists during pipeline execution.

    Each mode represents a specific point where LLM can provide value.
    """

    PRE_VALIDATE = "pre_validate"  # Validate inputs before stage execution
    ENRICH_CONTEXT = "enrich_context"  # Add context/suggestions to inputs
    ANALYZE_OUTPUT = "analyze_output"  # Interpret/explain stage outputs
    ERROR_DIAGNOSIS = "error_diagnosis"  # Diagnose failures with actionable fixes


class LLMAssistanceConfig(BaseModel):
    """Configuration for LLM assistance on a pipeline stage."""

    mode: LLMAssistanceMode
    enabled: bool = True
    timeout: float = Field(default=10.0, gt=0.0, le=300.0)
    cache_results: bool = True
    fallback_on_error: bool = True  # If LLM fails, continue without it
    prompt_template: str | None = None  # Custom prompt for this stage
    metadata: dict[str, Any] = Field(default_factory=dict)

    model_config = {"frozen": True}


class PipelineStage(BaseModel):
    """
    A single stage in a pipeline.

    Stages are immutable definitions. Runtime state is tracked in PipelineContext.
    """

    id: str = Field(default_factory=lambda: str(uuid.uuid4()))
    name: str
    description: str = ""

    # Execution
    handler: str  # Fully qualified function name for serialization
    timeout: float = Field(default=30.0, gt=0.0)
    retry_count: int = Field(default=0, ge=0, le=10)
    retry_delay: float = Field(default=1.0, ge=0.0)

    # Dependencies (DAG edges - stage IDs that must complete first)
    depends_on: list[str] = Field(default_factory=list)

    # Conditional execution
    condition: str | None = None  # Expression to evaluate (e.g., "ctx.get('skip') != True")
    skip_on_failure: bool = False  # Skip this stage if any dependency failed

    # LLM Assistance
    llm_assistance: list[LLMAssistanceConfig] = Field(default_factory=list)

    # Metadata
    tags: list[str] = Field(default_factory=list)
    metadata: dict[str, Any] = Field(default_factory=dict)

    model_config = {"frozen": True}

    def has_llm_mode(self, mode: LLMAssistanceMode) -> bool:
        """Check if this stage has a specific LLM assistance mode configured."""
        return any(cfg.mode == mode and cfg.enabled for cfg in self.llm_assistance)

    def get_llm_config(self, mode: LLMAssistanceMode) -> LLMAssistanceConfig | None:
        """Get LLM config for a specific mode."""
        for cfg in self.llm_assistance:
            if cfg.mode == mode and cfg.enabled:
                return cfg
        return None


class Pipeline(BaseModel):
    """
    A complete pipeline definition.

    Pipelines are immutable blueprints. They define WHAT to execute, not HOW.
    Execution state is managed separately in PipelineContext.
    """

    id: str = Field(default_factory=lambda: str(uuid.uuid4()))
    name: str
    description: str = ""
    version: str = "1.0.0"

    # Stages form a DAG
    stages: list[PipelineStage]

    # Global settings
    max_parallel: int = Field(default=4, ge=1, le=32)
    timeout: float = Field(default=300.0, gt=0.0)  # Total pipeline timeout
    fail_fast: bool = True  # Stop on first failure

    # Global LLM assistance (applied to all stages unless overridden)
    default_llm_config: LLMAssistanceConfig | None = None

    # Metadata
    category: str = ""
    tags: list[str] = Field(default_factory=list)
    metadata: dict[str, Any] = Field(default_factory=dict)
    created_at: datetime = Field(default_factory=lambda: datetime.now(timezone.utc))

    model_config = {"frozen": True}

    @model_validator(mode="after")
    def validate_dag(self) -> "Pipeline":
        """Validate that stages form a valid DAG (no cycles)."""
        stage_ids = {stage.id for stage in self.stages}

        for stage in self.stages:
            for dep_id in stage.depends_on:
                if dep_id not in stage_ids:
                    raise ValueError(
                        f"Stage '{stage.name}' depends on unknown stage ID: {dep_id}"
                    )

        # Check for cycles using topological sort
        visited: set[str] = set()
        rec_stack: set[str] = set()
        adj: dict[str, list[str]] = {s.id: s.depends_on for s in self.stages}

        def has_cycle(node: str) -> bool:
            visited.add(node)
            rec_stack.add(node)
            for neighbor in adj.get(node, []):
                if neighbor not in visited:
                    if has_cycle(neighbor):
                        return True
                elif neighbor in rec_stack:
                    return True
            rec_stack.discard(node)
            return False

        for stage in self.stages:
            if stage.id not in visited:
                if has_cycle(stage.id):
                    raise ValueError(f"Pipeline contains a cycle involving stage: {stage.name}")

        return self

    def get_stage(self, stage_id: str) -> PipelineStage | None:
        """Get a stage by ID."""
        for stage in self.stages:
            if stage.id == stage_id:
                return stage
        return None

    def get_stage_by_name(self, name: str) -> PipelineStage | None:
        """Get a stage by name."""
        for stage in self.stages:
            if stage.name == name:
                return stage
        return None

    def get_root_stages(self) -> list[PipelineStage]:
        """Get stages with no dependencies (entry points)."""
        return [stage for stage in self.stages if not stage.depends_on]

    def get_dependents(self, stage_id: str) -> list[PipelineStage]:
        """Get stages that depend on a given stage."""
        return [stage for stage in self.stages if stage_id in stage.depends_on]

    def topological_order(self) -> list[PipelineStage]:
        """Return stages in topological order (dependencies first)."""
        in_degree: dict[str, int] = {s.id: len(s.depends_on) for s in self.stages}
        stage_map = {s.id: s for s in self.stages}
        queue = [sid for sid, deg in in_degree.items() if deg == 0]
        result: list[PipelineStage] = []

        while queue:
            sid = queue.pop(0)
            result.append(stage_map[sid])
            for dependent in self.get_dependents(sid):
                in_degree[dependent.id] -= 1
                if in_degree[dependent.id] == 0:
                    queue.append(dependent.id)

        return result


# Type alias for stage handlers
StageHandler = Callable[
    [dict[str, Any]],
    Coroutine[Any, Any, dict[str, Any]]
]


class PipelineBuilder:
    """Fluent builder for creating pipelines."""

    def __init__(self, name: str, description: str = ""):
        self._name = name
        self._description = description
        self._stages: list[PipelineStage] = []
        self._max_parallel = 4
        self._timeout = 300.0
        self._fail_fast = True
        self._default_llm_config: LLMAssistanceConfig | None = None
        self._category = ""
        self._tags: list[str] = []
        self._metadata: dict[str, Any] = {}

    def add_stage(
        self,
        name: str,
        handler: str,
        *,
        description: str = "",
        depends_on: list[str] | None = None,
        timeout: float = 30.0,
        retry_count: int = 0,
        llm_modes: list[LLMAssistanceMode] | None = None,
    ) -> "PipelineBuilder":
        """Add a stage to the pipeline."""
        llm_assistance = []
        if llm_modes:
            llm_assistance = [LLMAssistanceConfig(mode=mode) for mode in llm_modes]

        stage = PipelineStage(
            name=name,
            description=description,
            handler=handler,
            depends_on=depends_on or [],
            timeout=timeout,
            retry_count=retry_count,
            llm_assistance=llm_assistance,
        )
        self._stages.append(stage)
        return self

    def with_max_parallel(self, max_parallel: int) -> "PipelineBuilder":
        """Set maximum parallel stage execution."""
        self._max_parallel = max_parallel
        return self

    def with_timeout(self, timeout: float) -> "PipelineBuilder":
        """Set total pipeline timeout."""
        self._timeout = timeout
        return self

    def with_fail_fast(self, fail_fast: bool) -> "PipelineBuilder":
        """Set fail-fast behavior."""
        self._fail_fast = fail_fast
        return self

    def with_default_llm_config(self, config: LLMAssistanceConfig) -> "PipelineBuilder":
        """Set default LLM assistance config for all stages."""
        self._default_llm_config = config
        return self

    def with_category(self, category: str) -> "PipelineBuilder":
        """Set pipeline category."""
        self._category = category
        return self

    def with_tags(self, *tags: str) -> "PipelineBuilder":
        """Add tags to the pipeline."""
        self._tags.extend(tags)
        return self

    def with_metadata(self, **kwargs: Any) -> "PipelineBuilder":
        """Add metadata to the pipeline."""
        self._metadata.update(kwargs)
        return self

    def build(self) -> Pipeline:
        """Build the pipeline."""
        return Pipeline(
            name=self._name,
            description=self._description,
            stages=self._stages,
            max_parallel=self._max_parallel,
            timeout=self._timeout,
            fail_fast=self._fail_fast,
            default_llm_config=self._default_llm_config,
            category=self._category,
            tags=self._tags,
            metadata=self._metadata,
        )
