"""
Decorator API for declarative saga definitions.

Provides a user-friendly way to define sagas using decorators rather than
manual step registration. This makes saga definitions more readable and
easier to maintain.

Quick Start:
    >>> from sagaz import Saga, step, compensate
    >>> 
    >>> class OrderSaga(Saga):
    ...     @step(name="create_order")
    ...     async def create_order(self, ctx):
    ...         return await OrderService.create(ctx["order_data"])
    ...     
    ...     @compensate("create_order") 
    ...     async def cancel_order(self, ctx):
    ...         await OrderService.delete(ctx["order_id"])
    ...     
    ...     @step(name="charge_payment", depends_on=["create_order"])
    ...     async def charge(self, ctx):
    ...         return await PaymentService.charge(ctx["amount"])
    ...     
    ...     @compensate("charge_payment", depends_on=["create_order"])
    ...     async def refund(self, ctx):
    ...         await PaymentService.refund(ctx["charge_id"])
"""

import asyncio
import inspect
import logging
from collections.abc import Awaitable, Callable
from dataclasses import dataclass, field
from typing import Any, TypeVar

from sagaz.compensation_graph import CompensationType, SagaCompensationGraph

logger = logging.getLogger(__name__)

# Type for saga step functions
F = TypeVar("F", bound=Callable[..., Awaitable[Any]])

# Hook type definitions
OnEnterHook = Callable[[dict[str, Any], str], Awaitable[None] | None]
OnSuccessHook = Callable[[dict[str, Any], str, Any], Awaitable[None] | None]
OnFailureHook = Callable[[dict[str, Any], str, Exception], Awaitable[None] | None]
OnCompensateHook = Callable[[dict[str, Any], str], Awaitable[None] | None]


@dataclass
class StepMetadata:
    """Metadata attached to step functions via decorator."""
    name: str
    depends_on: list[str] = field(default_factory=list)
    aggregate_type: str | None = None
    event_type: str | None = None
    timeout_seconds: float = 60.0
    max_retries: int = 3
    description: str | None = None
    # Lifecycle hooks
    on_enter: OnEnterHook | None = None
    on_success: OnSuccessHook | None = None
    on_failure: OnFailureHook | None = None


@dataclass
class CompensationMetadata:
    """Metadata attached to compensation functions via decorator."""
    for_step: str
    depends_on: list[str] = field(default_factory=list)
    compensation_type: CompensationType = CompensationType.MECHANICAL
    timeout_seconds: float = 30.0
    max_retries: int = 3
    description: str | None = None
    # Lifecycle hook for compensation
    on_compensate: OnCompensateHook | None = None


def step(
    name: str,
    depends_on: list[str] | None = None,
    aggregate_type: str | None = None,
    event_type: str | None = None,
    timeout_seconds: float = 60.0,
    max_retries: int = 3,
    description: str | None = None,
    on_enter: OnEnterHook | None = None,
    on_success: OnSuccessHook | None = None,
    on_failure: OnFailureHook | None = None,
) -> Callable[[F], F]:
    """
    Decorator to mark a method as a saga step.
    
    Args:
        name: Unique identifier for this step
        depends_on: List of step names that must complete before this step
        aggregate_type: Event aggregate type (for outbox pattern)
        event_type: Event type (for outbox pattern)
        timeout_seconds: Step execution timeout (default: 60s)
        max_retries: Maximum retry attempts (default: 3)
        description: Human-readable description
        on_enter: Hook called before step execution (ctx, step_name) -> None
        on_success: Hook called after success (ctx, step_name, result) -> None
        on_failure: Hook called on failure (ctx, step_name, error) -> None
    
    Example:
        >>> class OrderSaga(Saga):
        ...     @step(name="create_order", on_success=publish_order_created)
        ...     async def create_order(self, ctx):
        ...         order = await OrderService.create(ctx["order_data"])
        ...         return {"order_id": order.id}
    """
    def decorator(func: F) -> F:
        # Store metadata on the function
        func._saga_step_meta = StepMetadata(
            name=name,
            depends_on=depends_on or [],
            aggregate_type=aggregate_type,
            event_type=event_type,
            timeout_seconds=timeout_seconds,
            max_retries=max_retries,
            description=description or f"Execute {name}",
            on_enter=on_enter,
            on_success=on_success,
            on_failure=on_failure,
        )
        return func
    return decorator


def compensate(
    for_step: str,
    depends_on: list[str] | None = None,
    compensation_type: CompensationType = CompensationType.MECHANICAL,
    timeout_seconds: float = 30.0,
    max_retries: int = 3,
    description: str | None = None,
    on_compensate: OnCompensateHook | None = None,
) -> Callable[[F], F]:
    """
    Decorator to mark a method as compensation for a step.
    
    Compensation functions are called when a saga fails and needs to
    undo previously completed steps.
    
    Args:
        for_step: Name of the step this compensates
        depends_on: Steps whose compensations must complete BEFORE this one
        compensation_type: Type of compensation (MECHANICAL, SEMANTIC, MANUAL)
        timeout_seconds: Compensation timeout (default: 30s)
        max_retries: Maximum retry attempts (default: 3)
        description: Human-readable description
        on_compensate: Hook called when compensation runs (ctx, step_name) -> None
    
    Example:
        >>> @compensate("charge_payment", on_compensate=publish_refund_event)
        ... async def refund(self, ctx):
        ...     await PaymentService.refund(ctx["charge_id"])
    """
    def decorator(func: F) -> F:
        func._saga_compensation_meta = CompensationMetadata(
            for_step=for_step,
            depends_on=depends_on or [],
            compensation_type=compensation_type,
            timeout_seconds=timeout_seconds,
            max_retries=max_retries,
            description=description or f"Compensate {for_step}",
            on_compensate=on_compensate,
        )
        return func
    return decorator


@dataclass
class SagaStepDefinition:
    """
    Complete definition of a saga step with its compensation.
    
    Used internally to track step and compensation pairs.
    """
    step_id: str
    forward_fn: Callable[[dict[str, Any]], Awaitable[dict[str, Any] | None]]
    compensation_fn: Callable[[dict[str, Any]], Awaitable[None]] | None = None
    depends_on: list[str] = field(default_factory=list)
    compensation_depends_on: list[str] = field(default_factory=list)
    compensation_type: CompensationType = CompensationType.MECHANICAL
    aggregate_type: str | None = None
    event_type: str | None = None
    timeout_seconds: float = 60.0
    compensation_timeout_seconds: float = 30.0
    max_retries: int = 3
    description: str | None = None
    # Lifecycle hooks
    on_enter: OnEnterHook | None = None
    on_success: OnSuccessHook | None = None
    on_failure: OnFailureHook | None = None
    on_compensate: OnCompensateHook | None = None


class Saga:
    """
    Base class for declarative saga definitions.
    
    Subclass this and use @step and @compensate decorators to define
    your saga's steps declaratively. The saga will automatically:
    
    - Collect all decorated methods
    - Build the execution dependency graph
    - Build the compensation dependency graph
    - Execute steps in parallel where possible
    - Compensate in correct order on failure
    - Notify all listeners of lifecycle events
    
    Class Attributes:
        saga_name: Optional name for the saga (used in events/metrics)
        listeners: List of SagaListener instances to notify
    
    Example:
        >>> from sagaz.listeners import LoggingSagaListener, MetricsSagaListener
        >>> 
        >>> class OrderSaga(Saga):
        ...     saga_name = "order-processing"
        ...     listeners = [LoggingSagaListener(), MetricsSagaListener()]
        ...     
        ...     @step(name="create_order")
        ...     async def create_order(self, ctx):
        ...         return {"order_id": "ORD-123"}
    """
    
    # Class-level attributes (override in subclass)
    saga_name: str | None = None
    listeners: list = []  # List of SagaListener instances

    def __init__(self):
        self._steps: list[SagaStepDefinition] = []
        self._step_registry: dict[str, SagaStepDefinition] = {}
        self._compensation_graph = SagaCompensationGraph()
        self._context: dict[str, Any] = {}
        self._saga_id: str = ""
        self._collect_steps()

    def _collect_steps(self) -> None:
        """Collect decorated methods into step definitions."""
        # First pass: collect all steps
        for attr_name in dir(self):
            if attr_name.startswith("_"):
                continue

            attr = getattr(self, attr_name)

            if hasattr(attr, "_saga_step_meta"):
                meta: StepMetadata = attr._saga_step_meta
                step_def = SagaStepDefinition(
                    step_id=meta.name,
                    forward_fn=attr,
                    depends_on=meta.depends_on.copy(),
                    aggregate_type=meta.aggregate_type,
                    event_type=meta.event_type,
                    timeout_seconds=meta.timeout_seconds,
                    max_retries=meta.max_retries,
                    description=meta.description,
                    on_enter=meta.on_enter,
                    on_success=meta.on_success,
                    on_failure=meta.on_failure,
                )
                self._steps.append(step_def)
                self._step_registry[meta.name] = step_def

        # Second pass: attach compensations
        for attr_name in dir(self):
            if attr_name.startswith("_"):
                continue

            attr = getattr(self, attr_name)

            if hasattr(attr, "_saga_compensation_meta"):
                meta: CompensationMetadata = attr._saga_compensation_meta
                step_name = meta.for_step

                if step_name in self._step_registry:
                    step = self._step_registry[step_name]
                    step.compensation_fn = attr
                    step.compensation_depends_on = meta.depends_on.copy()
                    step.compensation_type = meta.compensation_type
                    step.compensation_timeout_seconds = meta.timeout_seconds
                    step.on_compensate = meta.on_compensate

    def get_steps(self) -> list[SagaStepDefinition]:
        """Get all step definitions."""
        return self._steps.copy()

    def get_step(self, name: str) -> SagaStepDefinition | None:
        """Get a specific step by name."""
        return self._step_registry.get(name)

    def get_execution_order(self) -> list[list[SagaStepDefinition]]:
        """
        Compute step execution order respecting dependencies.
        
        Returns:
            List of levels, where steps in each level can run in parallel
        
        Raises:
            ValueError: If circular dependencies detected
        """
        if not self._steps:
            return []

        step_map = {s.step_id: s for s in self._steps}
        in_degree = {s.step_id: len(s.depends_on) for s in self._steps}
        remaining = set(step_map.keys())

        return self._topological_sort_steps(step_map, in_degree, remaining)

    def _topological_sort_steps(
        self,
        step_map: dict[str, SagaStepDefinition],
        in_degree: dict[str, int],
        remaining: set
    ) -> list[list[SagaStepDefinition]]:
        """Perform topological sort on steps."""
        levels: list[list[SagaStepDefinition]] = []

        while remaining:
            current_level = [step_map[sid] for sid in remaining if in_degree[sid] == 0]

            if not current_level:
                raise ValueError("Circular dependency detected in saga steps")

            levels.append(current_level)

            for step in current_level:
                remaining.remove(step.step_id)
                for other_id in remaining:
                    if step.step_id in step_map[other_id].depends_on:
                        in_degree[other_id] -= 1

        return levels

    async def run(
        self,
        initial_context: dict[str, Any],
        saga_id: str | None = None
    ) -> dict[str, Any]:
        """
        Execute this saga.
        
        Args:
            initial_context: Initial data for the saga
            saga_id: Optional saga identifier for tracing
        
        Returns:
            Final context with all step results merged
        
        Raises:
            Exception: If saga fails (compensations will be attempted first)
        """
        import uuid

        # Use saga_id from parameter, or from context, or generate new
        self._saga_id = saga_id or initial_context.get("saga_id") or str(uuid.uuid4())
        self._context = initial_context.copy()
        self._context["saga_id"] = self._saga_id  # Ensure saga_id is in context
        self._compensation_graph.reset_execution()

        # Get saga name (class attribute or class name)
        name = self._get_saga_name()

        # Register all compensations in the graph
        for step in self._steps:
            if step.compensation_fn:
                self._compensation_graph.register_compensation(
                    step.step_id,
                    step.compensation_fn,
                    depends_on=step.compensation_depends_on,
                    compensation_type=step.compensation_type,
                    max_retries=step.max_retries,
                    timeout_seconds=step.compensation_timeout_seconds
                )

        try:
            # Notify listeners: saga starting
            await self._notify_listeners("on_saga_start", name, self._saga_id, self._context)
            
            # Execute steps level by level
            execution_levels = self.get_execution_order()

            for level in execution_levels:
                await self._execute_level(level)

            # Notify listeners: saga completed
            await self._notify_listeners("on_saga_complete", name, self._saga_id, self._context)
            
            return self._context

        except Exception as e:
            # Compensate executed steps
            await self._compensate()
            
            # Notify listeners: saga failed
            await self._notify_listeners("on_saga_failed", name, self._saga_id, self._context, e)
            raise
    
    def _get_saga_name(self) -> str:
        """Get the saga name from class attribute or class name."""
        return self.saga_name or self.__class__.__name__
    
    async def _notify_listeners(self, event_name: str, *args) -> None:
        """Notify all listeners of an event."""
        for listener in self.listeners:
            try:
                handler = getattr(listener, event_name, None)
                if handler:
                    result = handler(*args)
                    if inspect.iscoroutine(result):
                        await result
            except Exception as e:
                logger.warning(f"Listener {type(listener).__name__}.{event_name} error: {e}")

    async def _execute_level(self, level: list[SagaStepDefinition]) -> None:
        """Execute all steps in a level concurrently."""
        if not level:
            return

        tasks = [
            self._execute_step(step)
            for step in level
        ]

        # Execute in parallel
        results = await asyncio.gather(*tasks, return_exceptions=True)

        # Check for exceptions
        for result in results:
            if isinstance(result, Exception):
                raise result

    async def _execute_step(self, step: SagaStepDefinition) -> None:
        """Execute a single step with lifecycle hooks and listeners."""
        saga_name = self._get_saga_name()
        
        # Notify listeners: step entering
        await self._notify_listeners("on_step_enter", saga_name, step.step_id, self._context)
        
        # Call on_enter hook
        await self._call_hook(step.on_enter, self._context, step.step_id)
        
        try:
            # Apply timeout
            result = await asyncio.wait_for(
                step.forward_fn(self._context),
                timeout=step.timeout_seconds
            )

            # Merge result into context
            if result and isinstance(result, dict):
                self._context.update(result)

            # Mark step as executed for compensation tracking
            self._context[f"__{step.step_id}_completed"] = True
            self._compensation_graph.mark_step_executed(step.step_id)
            
            # Call on_success hook
            await self._call_hook(step.on_success, self._context, step.step_id, result)
            
            # Notify listeners: step success
            await self._notify_listeners("on_step_success", saga_name, step.step_id, self._context, result)

        except TimeoutError as e:
            # Call on_failure hook
            await self._call_hook(step.on_failure, self._context, step.step_id, e)
            # Notify listeners: step failure
            await self._notify_listeners("on_step_failure", saga_name, step.step_id, self._context, e)
            raise TimeoutError(f"Step '{step.step_id}' timed out after {step.timeout_seconds}s")
        except Exception as e:
            # Call on_failure hook
            await self._call_hook(step.on_failure, self._context, step.step_id, e)
            # Notify listeners: step failure
            await self._notify_listeners("on_step_failure", saga_name, step.step_id, self._context, e)
            raise
    
    async def _call_hook(self, hook: Callable | None, *args) -> None:
        """Call a hook function, handling both sync and async hooks."""
        if hook is None:
            return
        
        try:
            result = hook(*args)
            # If it's a coroutine, await it
            if inspect.iscoroutine(result):
                await result
        except Exception as e:
            # Log but don't fail - hooks should not break saga execution
            logger.warning(f"Hook error (non-fatal): {e}")

    async def _compensate(self) -> None:
        """Execute compensations in dependency order."""
        comp_levels = self._compensation_graph.get_compensation_order()

        for level in comp_levels:
            tasks = [
                self._execute_compensation(step_id)
                for step_id in level
            ]
            # Execute compensations in parallel, don't fail on individual errors
            await asyncio.gather(*tasks, return_exceptions=True)

    async def _execute_compensation(self, step_id: str) -> None:
        """Execute a single compensation with lifecycle hook and listeners."""
        node = self._compensation_graph.get_compensation_info(step_id)
        if not node:
            return
        
        saga_name = self._get_saga_name()
        
        # Get the on_compensate hook from the step definition
        step = self._step_registry.get(step_id)
        on_compensate = step.on_compensate if step else None

        try:
            # Notify listeners: compensation starting
            await self._notify_listeners("on_compensation_start", saga_name, step_id, self._context)
            
            # Call on_compensate hook before compensation
            await self._call_hook(on_compensate, self._context, step_id)
            
            await asyncio.wait_for(
                node.compensation_fn(self._context),
                timeout=node.timeout_seconds
            )
            self._context[f"__{step_id}_compensated"] = True
            
            # Notify listeners: compensation complete
            await self._notify_listeners("on_compensation_complete", saga_name, step_id, self._context)
        except Exception as e:
            # Log but don't fail - we want all compensations to attempt
            self._context[f"__{step_id}_compensation_error"] = str(e)

    def __repr__(self) -> str:
        return f"{self.__class__.__name__}(steps={len(self._steps)})"


# Backward compatibility alias
DeclarativeSaga = Saga

# Terminology alias: @action is preferred over @step for clarity
# (A "step" conceptually contains both action AND compensation)
action = step
