"""
Saga Pattern - Distributed Transactions with Compensation (Production-Ready)

Enables multi-step business processes with automatic rollback capability.
Unlike traditional ACID transactions, Sagas use compensating transactions to undo failed steps.

Uses python-statemachine for robust state management with full async support.

Features:
- ✅ Full async/await support with proper state machine integration
- ✅ Idempotent operations with deduplication keys
- ✅ Retry logic with exponential backoff
- ✅ Timeout protection per step
- ✅ Comprehensive error handling with partial compensation recovery
- ✅ Context passing between steps
- ✅ Guard conditions for state transitions
- ✅ Detailed observability and logging
- ✅ Thread-safe execution with locks
- ✅ Saga versioning and migration support

Example - Trade Execution Saga:
    Step 1: Reserve funds         → Compensation: Unreserve funds
    Step 2: Send order to Binance → Compensation: Cancel order
    Step 3: Update position       → Compensation: Revert position
    Step 4: Log trade             → Compensation: Delete trade log

If any step fails, all previous steps are compensated (undone) in reverse order.

State Machine:
    PENDING → EXECUTING → COMPLETED
                     ↘ COMPENSATING → ROLLED_BACK
                     ↘ FAILED (unrecoverable compensation failure)

    Step states:
    PENDING → EXECUTING → COMPLETED
                     ↘ COMPENSATING → COMPENSATED
                     ↘ FAILED (unrecoverable)

Usage:
    saga = TradeSaga()
    await saga.add_step(
        name="reserve_funds",
        action=reserve_funds_action,
        compensation=unreserve_funds_compensation,
        timeout=10.0,
        retry_attempts=3
    )
    result = await saga.execute()  # Returns SagaResult with full details
"""

import asyncio
import logging
from abc import ABC
from collections.abc import Callable
from dataclasses import dataclass, field
from datetime import datetime
from typing import Any
from uuid import uuid4

from statemachine.exceptions import TransitionNotAllowed

from sagaz.exceptions import (
    SagaCompensationError,
    SagaExecutionError,
    SagaStepError,
    SagaTimeoutError,
)
from sagaz.state_machine import SagaStateMachine
from sagaz.types import ParallelFailureStrategy, SagaResult, SagaStatus, SagaStepStatus

# Configure logging
logger = logging.getLogger(__name__)


class _StepExecutor:
    """Helper class to wrap SagaStep for strategy execution"""
    def __init__(self, step: "SagaStep", saga_context: "SagaContext"):
        self.step = step
        self.saga_context = saga_context
        self.result = None
        self.error: Exception | None = None
        self.completed = False

    async def execute(self) -> Any:
        """Execute the step action and store result or error"""
        try:
            result = await self.step.action(self.saga_context)
            self.result = result
            self.step.result = result
            # Store result in context for dependent steps
            self.saga_context.set(self.step.name, result)
            self.completed = True
            return result
        except Exception as e:
            self.error = e
            self.step.error = e
            raise

    async def compensate(self) -> None:
        """Execute the step compensation if available"""
        if self.step.compensation:
            await self.step.compensation(self.result, self.saga_context)

    @property
    def name(self) -> str:
        return self.step.name


class Saga(ABC):  # noqa: B024
    """
    Production-ready base class for saga implementations with state machine

    Concrete sagas should:
    1. Inherit from Saga
    2. Define steps in constructor or add_step calls
    3. Call execute() to run the saga
    4. Handle SagaResult appropriately
    """



    def __init__(
        self,
        name: str = "Saga",
        version: str = "1.0",
        failure_strategy: ParallelFailureStrategy = ParallelFailureStrategy.FAIL_FAST_WITH_GRACE,
        retry_backoff_base: float = 0.01,  # Base timeout for exponential backoff (seconds)
    ):
        self.name = name
        self.version = version
        self.saga_id = str(uuid4())
        self.status = SagaStatus.PENDING
        self.steps: list[SagaStep] = []
        self.completed_steps: list[SagaStep] = []
        self.context = SagaContext()
        self.started_at: datetime | None = None
        self.completed_at: datetime | None = None
        self.error: Exception | None = None
        self.compensation_errors: list[Exception] = []
        self._state_machine = SagaStateMachine(self)
        self._executing = False
        self._execution_lock = asyncio.Lock()
        self._executed_step_keys: set[str] = set()  # For idempotency
        self.retry_backoff_base = retry_backoff_base  # Configurable retry backoff base

        # DAG/Parallel execution support
        self.step_dependencies: dict[str, set[str]] = {}  # step_name -> set of dependency names
        self.execution_batches: list[set[str]] = []
        self.failure_strategy = failure_strategy
        self._has_dependencies = False  # Track if any step has explicit dependencies

    async def _on_enter_executing(self) -> None:
        """Callback: entering EXECUTING state"""
        self.status = SagaStatus.EXECUTING
        self.started_at = datetime.now()
        self.completed_steps = []
        logger.info(f"Saga {self.name} [{self.saga_id}] entering EXECUTING state")

    async def _on_enter_compensating(self) -> None:
        """Callback: entering COMPENSATING state"""
        self.status = SagaStatus.COMPENSATING
        logger.warning(
            f"Saga {self.name} [{self.saga_id}] entering COMPENSATING state - "
            f"rolling back {len(self.completed_steps)} steps"
        )

    async def _on_enter_completed(self) -> None:
        """Callback: entering COMPLETED state"""
        self.status = SagaStatus.COMPLETED
        self.completed_at = datetime.now()
        execution_time = (
            (self.completed_at - self.started_at).total_seconds() if self.started_at else 0
        )
        logger.info(
            f"Saga {self.name} [{self.saga_id}] COMPLETED successfully in {execution_time:.2f}s"
        )

    async def _on_enter_rolled_back(self) -> None:
        """Callback: entering ROLLED_BACK state"""
        self.status = SagaStatus.ROLLED_BACK
        self.completed_at = datetime.now()
        execution_time = (
            (self.completed_at - self.started_at).total_seconds() if self.started_at else 0
        )
        logger.info(
            f"Saga {self.name} [{self.saga_id}] ROLLED_BACK after {execution_time:.2f}s"
        )

    async def _on_enter_failed(self) -> None:
        """Callback: entering FAILED state (unrecoverable)"""
        self.status = SagaStatus.FAILED
        self.completed_at = datetime.now()
        logger.error(
            f"Saga {self.name} [{self.saga_id}] FAILED - unrecoverable error during compensation"
        )

    async def add_step(
        self,
        name: str,
        action: Callable[..., Any],
        compensation: Callable[..., Any] | None = None,
        timeout: float = 30.0,
        compensation_timeout: float = 30.0,
        max_retries: int = 3,
        idempotency_key: str | None = None,
        dependencies: set[str] | None = None,
    ) -> None:
        """
        Add a step to the saga

        Args:
            name: Step name
            action: Forward action to execute (can be sync or async)
            compensation: Rollback action (can be sync or async)
            timeout: Timeout in seconds for action execution
            compensation_timeout: Timeout in seconds for compensation
            max_retries: Maximum retry attempts for action
            idempotency_key: Custom idempotency key (auto-generated if None)
            dependencies: Optional set of step names this step depends on.
                         If None: Sequential execution (depends on previous step)
                         If empty set or specified: Parallel DAG execution
        """
        if self._executing:
            raise SagaExecutionError("Cannot add steps while saga is executing")

        # Validate no duplicate step names
        if any(s.name == name for s in self.steps):
            raise ValueError(f"Step '{name}' already exists")

        step = SagaStep(
            name=name,
            action=action,
            compensation=compensation,
            timeout=timeout,
            compensation_timeout=compensation_timeout,
            max_retries=max_retries,
            idempotency_key=idempotency_key or str(uuid4()),
        )
        self.steps.append(step)

        # Track dependencies for DAG execution
        if dependencies is not None:
            self._has_dependencies = True
            self.step_dependencies[name] = dependencies
        else:
            self.step_dependencies[name] = set()

        logger.debug(f"Added step '{name}' to saga {self.name}")

    def set_failure_strategy(self, strategy: ParallelFailureStrategy) -> None:
        """Set the failure strategy for parallel execution"""
        self.failure_strategy = strategy
        logger.info(f"Saga {self.name} failure strategy set to {strategy.value}")

    def _build_execution_batches(self) -> list[set[str]]:
        """
        Build execution batches for DAG parallel execution
        
        Returns list of sets, where each set contains steps that can run in parallel
        """
        if not self._has_dependencies:
            # Pure sequential mode - each step is its own batch
            return [{step.name} for step in self.steps]

        batches = []
        executed = set()
        step_names = {step.name for step in self.steps}
        remaining = step_names.copy()

        while remaining:
            # Find all steps that can execute (dependencies satisfied)
            ready = {
                name for name in remaining
                if self.step_dependencies[name].issubset(executed)
            }

            if not ready:
                # Circular dependency or missing dependency
                missing_deps = []
                for name in remaining:
                    missing = self.step_dependencies[name] - executed
                    if missing:
                        missing_deps.append(f"{name} needs {missing}")

                raise ValueError(f"Circular or missing dependencies detected: {missing_deps}")

            batches.append(ready)
            executed.update(ready)
            remaining -= ready

        return batches

    async def _execute_dag(self) -> SagaResult:
        """Execute saga using DAG parallel execution"""
        start_time = datetime.now()
        step_map = {step.name: step for step in self.steps}
        strategy = self._get_parallel_strategy()

        try:
            # Execute each batch in sequence
            for batch_idx, batch in enumerate(self.execution_batches):
                batch_error = await self._execute_batch(
                    batch_idx, batch, step_map, strategy
                )

                if batch_error:
                    await self._compensate_all()
                    return self._build_dag_result(
                        start_time, success=False,
                        status=SagaStatus.ROLLED_BACK, error=batch_error
                    )

            return self._build_dag_result(start_time, success=True, status=SagaStatus.COMPLETED)

        except Exception as e:
            logger.error(f"DAG execution failed for saga {self.name}: {e}")
            return self._build_dag_result(
                start_time, success=False, status=SagaStatus.FAILED, error=e
            )

    def _get_parallel_strategy(self):
        """Get the parallel execution strategy implementation."""
        from sagaz.strategies.fail_fast import FailFastStrategy
        from sagaz.strategies.fail_fast_grace import FailFastWithGraceStrategy
        from sagaz.strategies.wait_all import WaitAllStrategy

        if self.failure_strategy == ParallelFailureStrategy.FAIL_FAST:
            return FailFastStrategy()
        if self.failure_strategy == ParallelFailureStrategy.WAIT_ALL:
            return WaitAllStrategy()
        return FailFastWithGraceStrategy()

    async def _execute_batch(self, batch_idx: int, batch: set, step_map: dict, strategy) -> Exception | None:
        """Execute a single batch. Returns exception if failed, None if success."""
        logger.info(f"Executing batch {batch_idx + 1}/{len(self.execution_batches)}: {batch}")

        batch_executors = [
            _StepExecutor(step_map[step_name], self.context)
            for step_name in batch
        ]

        try:
            await strategy.execute_parallel_steps(batch_executors)
            self._mark_batch_completed(batch, step_map)
            return None

        except Exception as batch_error:
            logger.error(f"Batch {batch_idx + 1} failed: {batch_error}")
            self._mark_completed_executors(batch_executors)
            return batch_error

    def _mark_batch_completed(self, batch: set, step_map: dict):
        """Mark all steps in batch as completed."""
        for step_name in batch:
            step = step_map[step_name]
            self.completed_steps.append(step)
            self._executed_step_keys.add(step.idempotency_key)
            logger.info(f"Step '{step_name}' completed successfully")

    def _mark_completed_executors(self, executors: list):
        """Mark any executors that completed before failure."""
        for executor in executors:
            if executor.completed and executor.step not in self.completed_steps:
                step = executor.step
                self.completed_steps.append(step)
                self._executed_step_keys.add(step.idempotency_key)
                logger.info(f"Step '{step.name}' completed before batch failure")

    def _build_dag_result(self, start_time, success: bool, status: SagaStatus, error: Exception = None) -> SagaResult:
        """Build a SagaResult for DAG execution."""
        execution_time = (datetime.now() - start_time).total_seconds()
        return SagaResult(
            success=success,
            saga_name=self.name,
            status=status,
            completed_steps=len(self.completed_steps),
            total_steps=len(self.steps),
            error=error,
            execution_time=execution_time,
            context=self.context if success else self.context,
            compensation_errors=self.compensation_errors if not success else [],
        )

    def set_failure_strategy(self, strategy: ParallelFailureStrategy) -> None:
        """Set the failure strategy for parallel execution"""
        self.failure_strategy = strategy
        logger.info(f"Saga {self.name} failure strategy set to {strategy.value}")

    async def execute(self) -> SagaResult:
        """
        Execute the saga with full error handling and compensation
        
        Automatically detects execution mode:
        - Sequential: When no dependencies specified (traditional saga)
        - Parallel DAG: When dependencies are specified

        Returns:
            SagaResult: Detailed result of saga execution

        Raises:
            SagaExecutionError: If saga is already executing or in invalid state
        """
        async with self._execution_lock:
            if self._executing:
                raise SagaExecutionError("Saga is already executing")

            self._executing = True
            start_time = datetime.now()

            try:
                return await self._execute_inner(start_time)

            except SagaExecutionError:
                raise

            except Exception as e:
                return await self._handle_execution_failure(e, start_time)

            finally:
                self._executing = False

    async def _execute_inner(self, start_time) -> SagaResult:
        """Inner execution logic."""
        # Handle empty saga
        if not self.steps:
            return self._empty_saga_result(start_time)

        # Build execution plan
        plan_error = self._build_plan()
        if plan_error:
            return self._planning_failure_result(plan_error, start_time)

        # Start state machine
        try:
            await self._state_machine.start()
        except TransitionNotAllowed as e:
            raise SagaExecutionError(f"Cannot start saga: {e}")

        # Execute based on mode
        if self._has_dependencies:
            return await self._execute_dag_mode(start_time)
        return await self._execute_sequential_mode(start_time)

    def _empty_saga_result(self, start_time) -> SagaResult:
        """Result for saga with no steps."""
        logger.info(f"Saga {self.name} has no steps - completing immediately")
        return SagaResult(
            success=True, saga_name=self.name, status=SagaStatus.COMPLETED,
            completed_steps=0, total_steps=0,
            execution_time=(datetime.now() - start_time).total_seconds(),
            context=self.context,
        )

    def _build_plan(self) -> Exception | None:
        """Build execution batches. Returns error if failed."""
        try:
            self.execution_batches = self._build_execution_batches()
            logger.info(f"Saga {self.name} execution plan: {len(self.execution_batches)} batches")
            return None
        except ValueError as ve:
            self.error = ve
            logger.error(f"Saga {self.name} failed during planning: {ve}")
            return ve

    def _planning_failure_result(self, error, start_time) -> SagaResult:
        """Result for planning failure."""
        return SagaResult(
            success=False, saga_name=self.name, status=SagaStatus.FAILED,
            completed_steps=0, total_steps=len(self.steps), error=error,
            execution_time=(datetime.now() - start_time).total_seconds(),
        )

    async def _execute_dag_mode(self, start_time) -> SagaResult:
        """Execute in DAG/parallel mode."""
        logger.info(f"Executing saga {self.name} in DAG mode")
        result = await self._execute_dag()

        if result.success:
            await self._state_machine.succeed()
        else:
            await self._finalize_dag_failure()

        return result

    async def _finalize_dag_failure(self):
        """Finalize state machine after DAG failure."""
        try:
            await self._state_machine.fail()
            await self._state_machine.finish_compensation()
        except TransitionNotAllowed:
            await self._state_machine.fail_unrecoverable()

    async def _execute_sequential_mode(self, start_time) -> SagaResult:
        """Execute in sequential mode."""
        logger.info(f"Executing saga {self.name} in sequential mode")

        for step in self.steps:
            if step.idempotency_key in self._executed_step_keys:
                logger.info(f"Skipping step '{step.name}' - already executed (idempotent)")
                continue

            await self._execute_step_with_retry(step)
            self.completed_steps.append(step)
            self._executed_step_keys.add(step.idempotency_key)

        await self._state_machine.succeed()

        return SagaResult(
            success=True, saga_name=self.name, status=SagaStatus.COMPLETED,
            completed_steps=len(self.completed_steps), total_steps=len(self.steps),
            execution_time=(datetime.now() - start_time).total_seconds(),
            context=self.context,
        )

    async def _handle_execution_failure(self, error: Exception, start_time) -> SagaResult:
        """Handle execution failure with compensation."""
        self.error = error
        logger.error(f"Saga {self.name} failed: {error}", exc_info=True)

        try:
            await self._state_machine.fail()
        except TransitionNotAllowed:
            return self._no_compensation_result(error, start_time)

        return await self._attempt_compensation(error, start_time)

    def _no_compensation_result(self, error, start_time) -> SagaResult:
        """Result when no compensation needed."""
        self.status = SagaStatus.ROLLED_BACK
        return SagaResult(
            success=False, saga_name=self.name, status=SagaStatus.ROLLED_BACK,
            completed_steps=0, total_steps=len(self.steps), error=error,
            execution_time=(datetime.now() - start_time).total_seconds(),
            context=self.context,
        )

    async def _attempt_compensation(self, error, start_time) -> SagaResult:
        """Attempt to compensate and return result."""
        try:
            await self._compensate_all()
            await self._state_machine.finish_compensation()

            return SagaResult(
                success=False, saga_name=self.name, status=SagaStatus.ROLLED_BACK,
                completed_steps=len(self.completed_steps), total_steps=len(self.steps),
                error=error, execution_time=(datetime.now() - start_time).total_seconds(),
                context=self.context, compensation_errors=self.compensation_errors,
            )

        except SagaCompensationError:
            await self._state_machine.compensation_failed()

            return SagaResult(
                success=False, saga_name=self.name, status=SagaStatus.FAILED,
                completed_steps=len(self.completed_steps), total_steps=len(self.steps),
                error=error, execution_time=(datetime.now() - start_time).total_seconds(),
                compensation_errors=self.compensation_errors,
            )

    async def _execute_step_with_retry(self, step: "SagaStep") -> None:
        """Execute a step with retry logic and exponential backoff"""
        last_error = None
        total_attempts = step.max_retries + 1  # Initial attempt + retries

        for attempt in range(total_attempts):
            try:
                step.retry_count = attempt
                await self._execute_step(step)
                return  # Success!

            except SagaTimeoutError as e:
                last_error = e
                logger.warning(
                    f"Step '{step.name}' timed out (attempt {attempt + 1}/{total_attempts})"
                )

            except SagaStepError as e:
                last_error = e
                logger.warning(
                    f"Step '{step.name}' failed (attempt {attempt + 1}/{total_attempts}): {e}"
                )

            # Exponential backoff before retry (configurable base timeout)
            if attempt < total_attempts - 1:
                backoff_time = self.retry_backoff_base * (2**attempt)
                logger.info(f"Retrying step '{step.name}' in {backoff_time}s...")
                await asyncio.sleep(backoff_time)

        # All retries exhausted
        step.error = last_error
        raise last_error

    async def _execute_step(self, step: "SagaStep") -> None:
        """Execute a single step with timeout"""
        try:
            step.status = SagaStepStatus.EXECUTING
            logger.info(f"Executing step: {step.name}")

            # Execute with timeout
            step.result = await asyncio.wait_for(
                self._invoke(step.action, self.context), timeout=step.timeout
            )

            # Store result in context for next steps
            self.context.set(step.name, step.result)

            step.status = SagaStepStatus.COMPLETED
            step.executed_at = datetime.now()
            logger.info(f"Step '{step.name}' completed successfully")

        except TimeoutError:
            step.status = SagaStepStatus.FAILED
            error = SagaTimeoutError(f"Step '{step.name}' timed out after {step.timeout}s")
            step.error = error
            raise error

        except Exception as e:
            step.status = SagaStepStatus.FAILED
            step.error = e
            msg = f"Step '{step.name}' failed: {e!s}"
            raise SagaStepError(msg)

    async def _compensate_all(self) -> None:
        """
        Compensate all completed steps in reverse order
        Continues even if some compensations fail, collecting all errors
        """
        compensation_errors = []

        for step in reversed(self.completed_steps):
            if step.compensation:
                try:
                    await self._compensate_step_with_retry(step)
                except SagaCompensationError as e:
                    compensation_errors.append(e)
                    self.compensation_errors.append(e)
                    logger.error(f"Compensation failed for step '{step.name}': {e}")
                    # Continue compensating other steps

        if compensation_errors:
            # Raise aggregate error with all compensation failures
            error_summary = "; ".join(str(e) for e in compensation_errors)
            raise SagaCompensationError(
                f"Failed to compensate {len(compensation_errors)} steps: {error_summary}"
            )

    async def _compensate_step_with_retry(self, step: "SagaStep") -> None:
        """Compensate a step with retry logic"""
        last_error = None
        max_comp_retries = 3

        for attempt in range(max_comp_retries):
            try:
                await self._compensate_step(step)
                return  # Success!

            except SagaCompensationError as e:
                last_error = e
                logger.warning(
                    f"Compensation for '{step.name}' failed "
                    f"(attempt {attempt + 1}/{max_comp_retries})"
                )

            # Exponential backoff (uses configurable base timeout)
            if attempt < max_comp_retries - 1:
                backoff_time = self.retry_backoff_base * (2**attempt)
                await asyncio.sleep(backoff_time)

        # All retries exhausted
        raise last_error

    async def _compensate_step(self, step: "SagaStep") -> None:
        """Compensate a single step with timeout"""
        try:
            step.status = SagaStepStatus.COMPENSATING
            logger.info(f"Compensating step: {step.name}")

            # Pass the step result to compensation for context
            await asyncio.wait_for(
                self._invoke(step.compensation, step.result, self.context),
                timeout=step.compensation_timeout,
            )

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

        except TimeoutError:
            step.status = SagaStepStatus.FAILED
            msg = f"Compensation for '{step.name}' timed out after {step.compensation_timeout}s"
            raise SagaCompensationError(msg)

        except Exception as e:
            step.status = SagaStepStatus.FAILED
            step.error = e
            msg = f"Compensation for step '{step.name}' failed: {e!s}"
            raise SagaCompensationError(msg)

    @staticmethod
    async def _invoke(func: Callable[..., Any], *args, **kwargs) -> Any:
        """Invoke function (handle both sync and async)"""
        if asyncio.iscoroutinefunction(func):
            return await func(*args, **kwargs)
        return func(*args, **kwargs)

    @property
    def current_state(self) -> str:
        """Get current state name"""
        return self._state_machine.current_state.name

    def get_status(self) -> dict[str, Any]:
        """Get detailed saga status"""
        return {
            "saga_id": self.saga_id,
            "name": self.name,
            "version": self.version,
            "status": self.status.value,
            "current_state": self.current_state,
            "total_steps": len(self.steps),
            "completed_steps": len(self.completed_steps),
            "steps": [
                {
                    "name": step.name,
                    "status": step.status.value,
                    "retry_count": step.retry_count,
                    "error": str(step.error) if step.error else None,
                    "executed_at": step.executed_at.isoformat() if step.executed_at else None,
                    "compensated_at": (
                        step.compensated_at.isoformat() if step.compensated_at else None
                    ),
                }
                for step in self.steps
            ],
            "started_at": self.started_at.isoformat() if self.started_at else None,
            "completed_at": self.completed_at.isoformat() if self.completed_at else None,
            "error": str(self.error) if self.error else None,
            "compensation_errors": [str(e) for e in self.compensation_errors],
        }


@dataclass
class SagaContext:
    """Context passed between saga steps for data sharing"""

    data: dict[str, Any] = field(default_factory=dict)
    metadata: dict[str, Any] = field(default_factory=dict)

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

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

    def has(self, key: str) -> bool:
        """Check if a key exists in the context"""
        return key in self.data


@dataclass
class SagaStep:
    """Represents a single step in a saga with full metadata"""

    name: str
    action: Callable[..., Any]  # Forward action
    compensation: Callable[..., Any] | None = None  # Rollback action
    status: SagaStepStatus = field(default=SagaStepStatus.PENDING)
    result: Any | None = None
    error: Exception | None = None
    executed_at: datetime | None = None
    compensated_at: datetime | None = None
    idempotency_key: str = field(default_factory=lambda: str(uuid4()))
    retry_attempts: int = 0
    max_retries: int = 3
    timeout: float = 30.0  # seconds
    compensation_timeout: float = 30.0
    retry_count: int = 0

    def __hash__(self):
        return hash(self.idempotency_key)

