"""
Pipeline Context - The State Bag.

Manages execution state with checkpointing for crash recovery.
All state is serializable - secrets are NEVER stored directly, only references.
"""

from __future__ import annotations

import json
import uuid
from datetime import datetime, timezone
from enum import Enum
from typing import Any, Protocol, runtime_checkable

from pydantic import BaseModel, Field

from aperion_flow.definitions.pipeline import PipelineStatus, StageStatus


class CheckpointType(str, Enum):
    """Types of checkpoints."""

    PIPELINE_START = "pipeline_start"
    STAGE_START = "stage_start"
    STAGE_COMPLETE = "stage_complete"
    STAGE_FAILED = "stage_failed"
    PIPELINE_COMPLETE = "pipeline_complete"
    PIPELINE_FAILED = "pipeline_failed"
    RECOVERY = "recovery"


class Checkpoint(BaseModel):
    """
    A checkpoint representing a point in pipeline execution.

    Used for crash recovery - pipelines resume from last successful checkpoint.
    """

    id: str = Field(default_factory=lambda: str(uuid.uuid4()))
    pipeline_id: str
    execution_id: str
    checkpoint_type: CheckpointType
    stage_id: str | None = None
    stage_name: str | None = None

    # State at checkpoint
    status: PipelineStatus | StageStatus
    data_snapshot: dict[str, Any] = Field(default_factory=dict)
    completed_stages: list[str] = Field(default_factory=list)
    failed_stages: list[str] = Field(default_factory=list)

    # Timing
    created_at: datetime = Field(default_factory=lambda: datetime.now(timezone.utc))

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

    model_config = {"frozen": True}

    def to_json(self) -> str:
        """Serialize checkpoint to JSON for persistence."""
        return self.model_dump_json()

    @classmethod
    def from_json(cls, data: str) -> "Checkpoint":
        """Deserialize checkpoint from JSON."""
        return cls.model_validate_json(data)


@runtime_checkable
class CheckpointStore(Protocol):
    """Protocol for checkpoint persistence backends."""

    async def save(self, checkpoint: Checkpoint) -> None:
        """Save a checkpoint."""
        ...

    async def load(self, execution_id: str) -> Checkpoint | None:
        """Load the latest checkpoint for an execution."""
        ...

    async def list_checkpoints(self, execution_id: str) -> list[Checkpoint]:
        """List all checkpoints for an execution."""
        ...

    async def delete(self, execution_id: str) -> None:
        """Delete all checkpoints for an execution."""
        ...


class StageResult(BaseModel):
    """Result of a stage execution."""

    stage_id: str
    stage_name: str
    status: StageStatus
    data: dict[str, Any] = Field(default_factory=dict)
    error: str | None = None
    error_type: str | None = None

    # Timing
    started_at: datetime | None = None
    completed_at: datetime | None = None

    # LLM Assistance results
    llm_analysis: dict[str, Any] | None = None

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

    @property
    def duration_seconds(self) -> float | None:
        """Calculate execution duration in seconds."""
        if self.started_at and self.completed_at:
            return (self.completed_at - self.started_at).total_seconds()
        return None


class PipelineContext(BaseModel):
    """
    Execution context for a pipeline run.

    This is the "State Bag" that flows through the pipeline.
    It is mutable during execution but serializable for checkpointing.

    CRITICAL SECURITY RULE: Never store secrets directly.
    Use SecretRef(name="KEY_NAME") patterns for secret references.
    """

    # Identity
    pipeline_id: str
    pipeline_name: str
    execution_id: str = Field(default_factory=lambda: str(uuid.uuid4()))

    # Status
    status: PipelineStatus = PipelineStatus.PENDING

    # Data flow - this is where stage outputs accumulate
    data: dict[str, Any] = Field(default_factory=dict)

    # Stage tracking
    stage_results: dict[str, StageResult] = Field(default_factory=dict)
    current_stage_id: str | None = None

    # Timing
    started_at: datetime | None = None
    completed_at: datetime | None = None

    # Error info
    error: str | None = None
    error_type: str | None = None
    error_stage: str | None = None

    # Metadata
    correlation_id: str | None = None
    source: str | None = None
    metadata: dict[str, Any] = Field(default_factory=dict)

    # Allow mutation during execution
    model_config = {"frozen": False}

    @property
    def completed_stage_ids(self) -> list[str]:
        """Get IDs of completed stages."""
        return [
            sid for sid, result in self.stage_results.items()
            if result.status == StageStatus.COMPLETED
        ]

    @property
    def failed_stage_ids(self) -> list[str]:
        """Get IDs of failed stages."""
        return [
            sid for sid, result in self.stage_results.items()
            if result.status == StageStatus.FAILED
        ]

    @property
    def pending_stage_ids(self) -> list[str]:
        """Get IDs of stages not yet started."""
        return [
            sid for sid, result in self.stage_results.items()
            if result.status == StageStatus.PENDING
        ]

    @property
    def duration_seconds(self) -> float | None:
        """Calculate total execution duration in seconds."""
        if self.started_at and self.completed_at:
            return (self.completed_at - self.started_at).total_seconds()
        return None

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

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

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

    def mark_started(self) -> None:
        """Mark pipeline as started."""
        self.status = PipelineStatus.RUNNING
        self.started_at = datetime.now(timezone.utc)

    def mark_completed(self) -> None:
        """Mark pipeline as completed."""
        self.status = PipelineStatus.COMPLETED
        self.completed_at = datetime.now(timezone.utc)

    def mark_failed(self, error: str, error_type: str | None = None, stage_id: str | None = None) -> None:
        """Mark pipeline as failed."""
        self.status = PipelineStatus.FAILED
        self.completed_at = datetime.now(timezone.utc)
        self.error = error
        self.error_type = error_type
        self.error_stage = stage_id

    def mark_cancelled(self) -> None:
        """Mark pipeline as cancelled."""
        self.status = PipelineStatus.CANCELLED
        self.completed_at = datetime.now(timezone.utc)

    def record_stage_start(self, stage_id: str, stage_name: str) -> None:
        """Record that a stage has started."""
        self.current_stage_id = stage_id
        self.stage_results[stage_id] = StageResult(
            stage_id=stage_id,
            stage_name=stage_name,
            status=StageStatus.RUNNING,
            started_at=datetime.now(timezone.utc),
        )

    def record_stage_complete(
        self,
        stage_id: str,
        data: dict[str, Any] | None = None,
        llm_analysis: dict[str, Any] | None = None,
    ) -> None:
        """Record that a stage has completed successfully."""
        if stage_id in self.stage_results:
            result = self.stage_results[stage_id]
            self.stage_results[stage_id] = StageResult(
                stage_id=result.stage_id,
                stage_name=result.stage_name,
                status=StageStatus.COMPLETED,
                data=data or {},
                started_at=result.started_at,
                completed_at=datetime.now(timezone.utc),
                llm_analysis=llm_analysis,
                metadata=result.metadata,
            )
        self.current_stage_id = None

    def record_stage_failed(
        self,
        stage_id: str,
        error: str,
        error_type: str | None = None,
    ) -> None:
        """Record that a stage has failed."""
        if stage_id in self.stage_results:
            result = self.stage_results[stage_id]
            self.stage_results[stage_id] = StageResult(
                stage_id=result.stage_id,
                stage_name=result.stage_name,
                status=StageStatus.FAILED,
                error=error,
                error_type=error_type,
                started_at=result.started_at,
                completed_at=datetime.now(timezone.utc),
                metadata=result.metadata,
            )
        self.current_stage_id = None

    def record_stage_skipped(self, stage_id: str, stage_name: str, reason: str = "") -> None:
        """Record that a stage was skipped."""
        self.stage_results[stage_id] = StageResult(
            stage_id=stage_id,
            stage_name=stage_name,
            status=StageStatus.SKIPPED,
            metadata={"skip_reason": reason} if reason else {},
        )

    def create_checkpoint(self, checkpoint_type: CheckpointType) -> Checkpoint:
        """Create a checkpoint of current state."""
        return Checkpoint(
            pipeline_id=self.pipeline_id,
            execution_id=self.execution_id,
            checkpoint_type=checkpoint_type,
            stage_id=self.current_stage_id,
            status=self.status,
            data_snapshot=self.data.copy(),
            completed_stages=self.completed_stage_ids,
            failed_stages=self.failed_stage_ids,
            metadata=self.metadata.copy(),
        )

    def restore_from_checkpoint(self, checkpoint: Checkpoint) -> None:
        """Restore context state from a checkpoint."""
        self.data = checkpoint.data_snapshot.copy()
        # Mark completed stages
        for stage_id in checkpoint.completed_stages:
            if stage_id not in self.stage_results:
                self.stage_results[stage_id] = StageResult(
                    stage_id=stage_id,
                    stage_name=f"stage_{stage_id[:8]}",  # Placeholder name
                    status=StageStatus.COMPLETED,
                )
        # Mark failed stages
        for stage_id in checkpoint.failed_stages:
            if stage_id not in self.stage_results:
                self.stage_results[stage_id] = StageResult(
                    stage_id=stage_id,
                    stage_name=f"stage_{stage_id[:8]}",
                    status=StageStatus.FAILED,
                )

    def to_dict(self) -> dict[str, Any]:
        """Convert context to dictionary for serialization."""
        return self.model_dump(mode="json")

    @classmethod
    def from_dict(cls, data: dict[str, Any]) -> "PipelineContext":
        """Create context from dictionary."""
        return cls.model_validate(data)


class SecretRef(BaseModel):
    """
    A reference to a secret.

    NEVER store actual secret values in PipelineContext.
    Use this class to reference secrets by name.
    The secret resolver will fetch actual values at runtime.
    """

    name: str
    source: str = "env"  # "env", "vault", "aws_secrets_manager", etc.
    required: bool = True
    default_ref: str | None = None  # Fallback reference name

    model_config = {"frozen": True}

    def __repr__(self) -> str:
        return f"SecretRef({self.source}:{self.name})"


class InMemoryCheckpointStore:
    """In-memory checkpoint store for testing and development."""

    def __init__(self) -> None:
        self._checkpoints: dict[str, list[Checkpoint]] = {}

    async def save(self, checkpoint: Checkpoint) -> None:
        """Save a checkpoint."""
        exec_id = checkpoint.execution_id
        if exec_id not in self._checkpoints:
            self._checkpoints[exec_id] = []
        self._checkpoints[exec_id].append(checkpoint)

    async def load(self, execution_id: str) -> Checkpoint | None:
        """Load the latest checkpoint for an execution."""
        checkpoints = self._checkpoints.get(execution_id, [])
        return checkpoints[-1] if checkpoints else None

    async def list_checkpoints(self, execution_id: str) -> list[Checkpoint]:
        """List all checkpoints for an execution."""
        return self._checkpoints.get(execution_id, [])

    async def delete(self, execution_id: str) -> None:
        """Delete all checkpoints for an execution."""
        self._checkpoints.pop(execution_id, None)

    def clear(self) -> None:
        """Clear all checkpoints (for testing)."""
        self._checkpoints.clear()
