"""
Pipeline definition and management for Aperion Legendary agent system.
"""

import time
import uuid
from dataclasses import dataclass, field
from enum import Enum
from typing import Any, Callable, Optional


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

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


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

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


class LLMAssistanceMode(Enum):
    """LLM assistance modes for pipeline stages."""

    PRE_VALIDATE = "pre_validate"  # Validate inputs before execution
    ENRICH_CONTEXT = "enrich_context"  # Add LLM insights to context
    ANALYZE_OUTPUT = "analyze_output"  # Validate/explain results
    ERROR_DIAGNOSIS = "error_diagnosis"  # Diagnose failures


@dataclass
class LLMAssistanceConfig:
    """
    Configuration for LLM assistance in pipeline stages.

    Enables optional LLM-powered features without changing existing behavior.
    All LLM operations are non-blocking and use intelligent routing.
    """

    mode: LLMAssistanceMode
    enabled: bool = True
    provider_preference: str = "gemini"  # Prefer free tier
    timeout: float = 10.0  # Max time for LLM operation
    fallback_on_error: bool = True  # Continue on LLM failure
    cache_results: bool = True  # Cache LLM responses
    validation_prompt: Optional[str] = None  # Custom prompt template
    metadata: dict[str, Any] = field(default_factory=dict)


@dataclass
class LLMAssistanceResult:
    """Result from LLM assistance operation."""

    success: bool
    mode: LLMAssistanceMode
    result: dict[str, Any]  # LLM response data
    provider_used: str
    elapsed_time: float
    cached: bool = False
    error: Optional[str] = None


@dataclass
class PipelineContext:
    """Context passed between pipeline stages."""

    data: dict[str, Any] = field(default_factory=dict)
    metadata: dict[str, Any] = field(default_factory=dict)
    llm_insights: dict[str, Any] = field(default_factory=dict)  # LLM-generated insights

    def get(self, key: str, default: Any = None) -> Any:
        """Get value from context data."""
        return self.data.get(key, default)

    def set(self, key: str, value: Any) -> None:
        """Set value in context data."""
        self.data[key] = value

    def update(self, data: dict[str, Any]) -> None:
        """Update context with new data."""
        self.data.update(data)

    def add_llm_insight(self, stage_id: str, insight: dict[str, Any]) -> None:
        """Add LLM-generated insight for a stage."""
        self.llm_insights[stage_id] = insight


@dataclass
class PipelineStage:
    """Individual stage in a pipeline."""

    id: str
    name: str
    handler: Callable[[PipelineContext], PipelineContext]
    depends_on: list[str] = field(default_factory=list)
    timeout: Optional[int] = None
    retry_count: int = 0
    llm_assistance: Optional[LLMAssistanceConfig] = None  # Optional LLM enhancement
    recovery_config: Optional[Any] = None  # Optional error recovery configuration
    status: StageStatus = StageStatus.PENDING
    error: Optional[str] = None
    start_time: Optional[float] = None
    end_time: Optional[float] = None

    def execute(self, context: PipelineContext) -> PipelineContext:
        """Execute the stage."""
        self.status = StageStatus.RUNNING
        self.start_time = time.time()

        try:
            result = self.handler(context)
            self.status = StageStatus.COMPLETED
            self.end_time = time.time()
            return result
        except Exception as e:
            self.status = StageStatus.FAILED
            self.error = str(e)
            self.end_time = time.time()
            raise

    def duration(self) -> Optional[float]:
        """Get stage execution duration."""
        if self.start_time and self.end_time:
            return self.end_time - self.start_time
        return None


@dataclass
class Pipeline:
    """Pipeline definition and execution controller."""

    id: str = field(default_factory=lambda: str(uuid.uuid4()))
    name: str = ""
    description: str = ""
    stages: list[PipelineStage] = field(default_factory=list)
    status: PipelineStatus = PipelineStatus.PENDING
    context: PipelineContext = field(default_factory=PipelineContext)
    created_at: float = field(default_factory=time.time)
    started_at: Optional[float] = None
    completed_at: Optional[float] = None
    error: Optional[str] = None

    def add_stage(self, stage: PipelineStage) -> None:
        """Add a stage to the pipeline."""
        self.stages.append(stage)

    def create_stage(
        self,
        name: str,
        handler: Callable[[PipelineContext], PipelineContext],
        depends_on: Optional[list[str]] = None,
        timeout: Optional[int] = None,
        retry_count: int = 0,
        llm_assistance: Optional[LLMAssistanceConfig] = None,
    ) -> PipelineStage:
        """Create and add a new stage to the pipeline."""
        stage_id = f"{self.id}_{len(self.stages):03d}_{name.lower().replace(' ', '_')}"
        stage = PipelineStage(
            id=stage_id,
            name=name,
            handler=handler,
            depends_on=depends_on or [],
            timeout=timeout,
            retry_count=retry_count,
            llm_assistance=llm_assistance,
        )
        self.add_stage(stage)
        return stage

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

    def get_ready_stages(self) -> list[PipelineStage]:
        """Get stages that are ready to execute."""
        ready = []
        for stage in self.stages:
            if stage.status != StageStatus.PENDING:
                continue

            # Check if all dependencies are completed
            dependencies_met = True
            for dep_id in stage.depends_on:
                dep_stage = self.get_stage(dep_id)
                if not dep_stage or dep_stage.status != StageStatus.COMPLETED:
                    dependencies_met = False
                    break

            if dependencies_met:
                ready.append(stage)

        return ready

    def reset_for_execution(self) -> None:
        """Reset pipeline state for fresh execution."""
        # Reset pipeline status
        self.status = PipelineStatus.PENDING
        self.started_at = None
        self.completed_at = None
        self.error = None

        # Reset all stage statuses
        for stage in self.stages:
            stage.status = StageStatus.PENDING
            stage.error = None
            stage.start_time = None
            stage.end_time = None

        # Create fresh context
        self.context = PipelineContext()

    def is_complete(self) -> bool:
        """Check if pipeline is complete."""
        return all(
            stage.status
            in [StageStatus.COMPLETED, StageStatus.SKIPPED, StageStatus.FAILED]
            for stage in self.stages
        )

    def has_failed_stages(self) -> bool:
        """Check if any stages have failed."""
        return any(stage.status == StageStatus.FAILED for stage in self.stages)

    def get_status_summary(self) -> dict[str, int]:
        """Get a summary of stage statuses."""
        summary = {status.value: 0 for status in StageStatus}
        for stage in self.stages:
            summary[stage.status.value] += 1
        return summary

    def duration(self) -> Optional[float]:
        """Get pipeline execution duration."""
        if self.started_at and self.completed_at:
            return self.completed_at - self.started_at
        return None

    def to_dict(self) -> dict[str, Any]:
        """Convert pipeline to dictionary."""
        return {
            "id": self.id,
            "name": self.name,
            "description": self.description,
            "status": self.status.value,
            "created_at": self.created_at,
            "started_at": self.started_at,
            "completed_at": self.completed_at,
            "duration": self.duration(),
            "error": self.error,
            "stages": [
                {
                    "id": stage.id,
                    "name": stage.name,
                    "status": stage.status.value,
                    "depends_on": stage.depends_on,
                    "timeout": stage.timeout,
                    "retry_count": stage.retry_count,
                    "error": stage.error,
                    "duration": stage.duration(),
                }
                for stage in self.stages
            ],
            "status_summary": self.get_status_summary(),
        }
