"""
Saga Pattern Implementation

Provides distributed transaction coordination with compensation logic
for microservices architecture.
"""

import asyncio
import uuid
from typing import Callable, List, Dict, Any, Optional, Tuple
from dataclasses import dataclass, field
from datetime import datetime
from enum import Enum
import logging

logger = logging.getLogger(__name__)


class SagaStepStatus(str, Enum):
    """Status of a saga step"""
    PENDING = "pending"
    IN_PROGRESS = "in_progress"
    COMPLETED = "completed"
    FAILED = "failed"
    COMPENSATING = "compensating"
    COMPENSATED = "compensated"


class SagaStatus(str, Enum):
    """Overall saga status"""
    PENDING = "pending"
    IN_PROGRESS = "in_progress"
    COMPLETED = "completed"
    FAILED = "failed"
    COMPENSATING = "compensating"
    COMPENSATED = "compensated"


@dataclass
class SagaStep:
    """
    A single step in a saga with compensation logic.

    Args:
        name: Human-readable step name
        action: Async function to execute
        compensation: Async function to undo the action
        action_args: Arguments for the action function
        action_kwargs: Keyword arguments for the action function
    """
    name: str
    action: Callable
    compensation: Optional[Callable] = None
    action_args: tuple = field(default_factory=tuple)
    action_kwargs: dict = field(default_factory=dict)

    # Runtime state
    status: SagaStepStatus = SagaStepStatus.PENDING
    result: Any = None
    error: Optional[Exception] = None
    started_at: Optional[datetime] = None
    completed_at: Optional[datetime] = None


@dataclass
class SagaExecution:
    """Tracks execution state of a saga"""
    saga_id: str
    status: SagaStatus = SagaStatus.PENDING
    steps: List[SagaStep] = field(default_factory=list)
    current_step: int = 0
    started_at: Optional[datetime] = None
    completed_at: Optional[datetime] = None
    error: Optional[Exception] = None


class SagaCoordinator:
    """
    Coordinates saga execution with automatic compensation on failure.

    Example:
        ```python
        saga = SagaCoordinator()

        saga.add_step(
            "create_order",
            action=create_order,
            compensation=delete_order,
            action_args=(order_data,)
        )

        saga.add_step(
            "reserve_inventory",
            action=reserve_inventory,
            compensation=release_inventory,
            action_args=(order_id, items)
        )

        result = await saga.execute()
        ```
    """

    def __init__(self, saga_id: Optional[str] = None):
        self.execution = SagaExecution(
            saga_id=saga_id or str(uuid.uuid4())
        )
        self._completed_steps: List[SagaStep] = []

    def add_step(
        self,
        name: str,
        action: Callable,
        compensation: Optional[Callable] = None,
        action_args: tuple = (),
        action_kwargs: dict = None
    ):
        """
        Add a step to the saga.

        Args:
            name: Human-readable step name
            action: Async function to execute
            compensation: Async function to undo the action (optional)
            action_args: Arguments for the action function
            action_kwargs: Keyword arguments for the action function
        """
        step = SagaStep(
            name=name,
            action=action,
            compensation=compensation,
            action_args=action_args,
            action_kwargs=action_kwargs or {}
        )
        self.execution.steps.append(step)
        logger.debug(f"Added step '{name}' to saga {self.execution.saga_id}")

    async def execute(self) -> Tuple[bool, Optional[Any], Optional[Exception]]:
        """
        Execute all saga steps in sequence.

        Returns:
            Tuple of (success: bool, final_result: Any, error: Optional[Exception])
        """
        self.execution.status = SagaStatus.IN_PROGRESS
        self.execution.started_at = datetime.now()

        logger.info(
            f"Starting saga {self.execution.saga_id} with {len(self.execution.steps)} steps"
        )

        try:
            # Execute each step
            for idx, step in enumerate(self.execution.steps):
                self.execution.current_step = idx

                success = await self._execute_step(step)

                if not success:
                    # Step failed, trigger compensation
                    logger.error(
                        f"Saga {self.execution.saga_id} failed at step '{step.name}': {step.error}"
                    )
                    await self._compensate()

                    self.execution.status = SagaStatus.COMPENSATED
                    self.execution.completed_at = datetime.now()
                    self.execution.error = step.error

                    return False, None, step.error

                # Step succeeded
                self._completed_steps.append(step)

            # All steps completed successfully
            self.execution.status = SagaStatus.COMPLETED
            self.execution.completed_at = datetime.now()

            # Return the result of the last step
            final_result = self.execution.steps[-1].result if self.execution.steps else None

            logger.info(f"Saga {self.execution.saga_id} completed successfully")
            return True, final_result, None

        except Exception as e:
            logger.exception(f"Unexpected error in saga {self.execution.saga_id}: {e}")
            await self._compensate()

            self.execution.status = SagaStatus.FAILED
            self.execution.completed_at = datetime.now()
            self.execution.error = e

            return False, None, e

    async def _execute_step(self, step: SagaStep) -> bool:
        """
        Execute a single saga step.

        Returns:
            True if step succeeded, False otherwise
        """
        step.status = SagaStepStatus.IN_PROGRESS
        step.started_at = datetime.now()

        logger.info(f"Executing saga step '{step.name}'")

        try:
            # Execute the action
            if asyncio.iscoroutinefunction(step.action):
                result = await step.action(*step.action_args, **step.action_kwargs)
            else:
                result = step.action(*step.action_args, **step.action_kwargs)

            step.result = result
            step.status = SagaStepStatus.COMPLETED
            step.completed_at = datetime.now()

            logger.info(f"Saga step '{step.name}' completed successfully")
            return True

        except Exception as e:
            step.error = e
            step.status = SagaStepStatus.FAILED
            step.completed_at = datetime.now()

            logger.error(f"Saga step '{step.name}' failed: {e}")
            return False

    async def _compensate(self):
        """
        Execute compensation logic for all completed steps in reverse order.
        """
        if not self._completed_steps:
            logger.info(f"No steps to compensate for saga {self.execution.saga_id}")
            return

        self.execution.status = SagaStatus.COMPENSATING

        logger.info(
            f"Starting compensation for saga {self.execution.saga_id} "
            f"({len(self._completed_steps)} steps to compensate)"
        )

        # Compensate in reverse order
        for step in reversed(self._completed_steps):
            if step.compensation is None:
                logger.warning(
                    f"Step '{step.name}' has no compensation function, skipping"
                )
                continue

            step.status = SagaStepStatus.COMPENSATING

            try:
                logger.info(f"Compensating step '{step.name}'")

                # Execute compensation with the result from the original action
                compensation_args = (step.result,) if step.result is not None else ()

                if asyncio.iscoroutinefunction(step.compensation):
                    await step.compensation(*compensation_args)
                else:
                    step.compensation(*compensation_args)

                step.status = SagaStepStatus.COMPENSATED
                logger.info(f"Step '{step.name}' compensated successfully")

            except Exception as e:
                logger.error(f"Failed to compensate step '{step.name}': {e}")
                # Continue compensating other steps even if one fails

        logger.info(f"Compensation completed for saga {self.execution.saga_id}")

    def get_execution_summary(self) -> Dict[str, Any]:
        """Get summary of saga execution"""
        return {
            "saga_id": self.execution.saga_id,
            "status": self.execution.status.value,
            "total_steps": len(self.execution.steps),
            "current_step": self.execution.current_step,
            "completed_steps": len(self._completed_steps),
            "started_at": self.execution.started_at.isoformat() if self.execution.started_at else None,
            "completed_at": self.execution.completed_at.isoformat() if self.execution.completed_at else None,
            "error": str(self.execution.error) if self.execution.error else None,
            "steps": [
                {
                    "name": step.name,
                    "status": step.status.value,
                    "has_compensation": step.compensation is not None,
                    "error": str(step.error) if step.error else None
                }
                for step in self.execution.steps
            ]
        }