"""
Saga Orchestrator - Coordinates distributed transactions.

Executes saga steps in order, handles failures with compensation,
supports retry with exponential backoff, and ensures idempotency.
"""
import importlib
import logging
import time
import traceback
from typing import Any, Callable, Dict, List, Optional, Tuple
from uuid import UUID

from django.db import transaction
from django.utils import timezone

from .models import (
    SagaExecution,
    SagaStep,
    SagaLog,
    SagaStatus,
    StepStatus,
    SagaLogEventType,
)

logger = logging.getLogger(__name__)


class SagaStepDefinition:
    """
    Definition of a saga step (used to build a saga before execution).
    
    Example:
        step = SagaStepDefinition(
            name='create_part',
            action='apps.plm.services.PartService.create_part',
            compensation='apps.plm.services.PartService.delete_part',
            is_critical=True,
            max_retries=3,
            timeout_seconds=60,
        )
    """
    
    def __init__(
        self,
        name: str,
        action: str,
        compensation: Optional[str] = None,
        is_critical: bool = True,
        max_retries: int = 3,
        timeout_seconds: int = 300,
        input_mapping: Optional[Callable[[Dict], Dict]] = None,
    ):
        self.name = name
        self.action = action
        self.compensation = compensation
        self.is_critical = is_critical
        self.max_retries = max_retries
        self.timeout_seconds = timeout_seconds
        self.input_mapping = input_mapping


class SagaDefinitionBuilder:
    """
    Builder for constructing saga definitions.
    
    Example:
        saga = (SagaDefinitionBuilder('create_product_workflow')
            .add_step(
                name='create_part',
                action='apps.plm.services.PartService.create_part',
                compensation='apps.plm.services.PartService.delete_part',
            )
            .add_step(
                name='register_manufacturing',
                action='apps.manufacturing.services.ManufacturingService.create_work_order',
                compensation='apps.manufacturing.services.ManufacturingService.cancel_work_order',
            )
            .build())
    """

    def __init__(self, name: str):
        self.name = name
        self.steps: List[SagaStepDefinition] = []

    def add_step(
        self,
        name: str,
        action: str,
        compensation: Optional[str] = None,
        is_critical: bool = True,
        max_retries: int = 3,
        timeout_seconds: int = 300,
        input_mapping: Optional[Callable[[Dict], Dict]] = None,
    ) -> 'SagaDefinitionBuilder':
        """Add a step to the saga definition."""
        self.steps.append(SagaStepDefinition(
            name=name,
            action=action,
            compensation=compensation,
            is_critical=is_critical,
            max_retries=max_retries,
            timeout_seconds=timeout_seconds,
            input_mapping=input_mapping,
        ))
        return self

    def build(self) -> Tuple[str, List[SagaStepDefinition]]:
        """Build and return the saga definition."""
        if not self.steps:
            raise ValueError("Saga must have at least one step")
        return self.name, self.steps


class SagaOrchestrator:
    """
    Orchestrates saga execution with compensation and retry logic.
    
    Features:
    - Sequential step execution with forward/compensation actions
    - Exponential backoff retry mechanism
    - Idempotency key support to prevent duplicate execution
    - Saga context shared across all steps
    - Full audit logging via SagaLog
    - Pause/Resume support
    - Compatible with future migration to Temporal/Camunda
    
    Example:
        orchestrator = SagaOrchestrator()
        
        # Define the saga
        saga_name, steps = (SagaDefinitionBuilder('create_product')
            .add_step(
                name='create_part',
                action='apps.plm.services.create_part',
                compensation='apps.plm.services.delete_part',
            )
            .add_step(
                name='create_work_order',
                action='apps.manufacturing.services.create_work_order',
                compensation='apps.manufacturing.services.cancel_work_order',
            )
            .build())
        
        execution = orchestrator.execute_saga(
            tenant=tenant,
            saga_name=saga_name,
            steps=steps,
            context={'product_name': 'Widget A'},
            idempotency_key='create-product-widget-a',
        )
    """

    # Registry of action functions (dotted path → callable)
    _action_registry: Dict[str, Callable] = {}

    @classmethod
    def register_action(cls, path: str, func: Callable) -> None:
        """Register a callable action by its dotted path."""
        cls._action_registry[path] = func

    @classmethod
    def _resolve_action(cls, action_path: str) -> Callable:
        """
        Resolve a dotted path to a callable.
        First checks registry, then tries dynamic import.
        """
        if action_path in cls._action_registry:
            return cls._action_registry[action_path]

        # Dynamic import: 'module.path.ClassName.method_name'
        try:
            parts = action_path.rsplit('.', 1)
            if len(parts) == 2:
                module_path, attr_name = parts
                # Try importing as module.attribute
                try:
                    module = importlib.import_module(module_path)
                    func = getattr(module, attr_name)
                    cls._action_registry[action_path] = func
                    return func
                except (ImportError, AttributeError):
                    # Try class.method pattern
                    class_parts = module_path.rsplit('.', 1)
                    if len(class_parts) == 2:
                        mod_path, class_name = class_parts
                        module = importlib.import_module(mod_path)
                        klass = getattr(module, class_name)
                        func = getattr(klass, attr_name)
                        cls._action_registry[action_path] = func
                        return func
        except Exception as e:
            raise ImportError(
                f"Cannot resolve action '{action_path}': {e}"
            ) from e

        raise ImportError(f"Cannot resolve action '{action_path}'")

    def execute_saga(
        self,
        tenant,
        saga_name: str,
        steps: List[SagaStepDefinition],
        context: Optional[Dict[str, Any]] = None,
        idempotency_key: Optional[str] = None,
        correlation_id: Optional[str] = None,
        user=None,
        workflow_instance_id: Optional[UUID] = None,
        max_retries: int = 3,
    ) -> SagaExecution:
        """
        Execute a saga with the given steps and context.
        
        Args:
            tenant: Tenant instance
            saga_name: Name of the saga
            steps: List of SagaStepDefinition instances
            context: Shared context data for all steps
            idempotency_key: Unique key to prevent duplicate execution
            correlation_id: Correlation ID for tracing
            user: User who initiated the saga
            workflow_instance_id: Optional linked workflow instance
            max_retries: Max retries per step
            
        Returns:
            SagaExecution instance
        """
        context = context or {}

        # Idempotency check
        if idempotency_key:
            existing = SagaExecution.objects.filter(
                idempotency_key=idempotency_key
            ).first()
            if existing:
                logger.info(
                    f"Saga with idempotency_key={idempotency_key} already exists: {existing.id}"
                )
                return existing

        # Create saga execution
        execution = SagaExecution.objects.create(
            tenant=tenant,
            name=saga_name,
            status=SagaStatus.PENDING,
            context=context,
            idempotency_key=idempotency_key,
            correlation_id=correlation_id,
            workflow_instance_id=workflow_instance_id,
            started_by=user,
            max_retries=max_retries,
        )

        # Create step records
        for i, step_def in enumerate(steps):
            SagaStep.objects.create(
                saga=execution,
                name=step_def.name,
                order=i + 1,
                action_type=step_def.action,
                compensation_type=step_def.compensation,
                is_critical=step_def.is_critical,
                max_retries=step_def.max_retries,
                timeout_seconds=step_def.timeout_seconds,
                status=StepStatus.PENDING,
            )

        # Log saga start
        self._log(execution, None, SagaLogEventType.SAGA_STARTED, 'Saga execution started')

        # Execute steps
        return self._run_steps(execution)

    def compensate(
        self,
        saga_id: UUID,
        failed_step_order: Optional[int] = None,
    ) -> SagaExecution:
        """
        Trigger compensation for a saga, starting from failed_step_order
        (or the last completed step) in reverse order.
        
        Args:
            saga_id: UUID of the saga execution
            failed_step_order: Optional step order to start compensation from
            
        Returns:
            Updated SagaExecution
        """
        execution = SagaExecution.objects.get(pk=saga_id)

        if execution.status not in [SagaStatus.FAILED, SagaStatus.RUNNING]:
            raise ValueError(
                f"Cannot compensate saga in status {execution.status}"
            )

        execution.status = SagaStatus.COMPENSATING
        execution.save(update_fields=['status', 'updated_at'])

        self._log(execution, None, SagaLogEventType.SAGA_COMPENSATING, 'Starting compensation')

        return self._run_compensation(execution, failed_step_order)

    def resume_saga(self, saga_id: UUID) -> SagaExecution:
        """
        Resume a paused or failed saga from the last incomplete step.
        
        Args:
            saga_id: UUID of the saga execution
            
        Returns:
            Updated SagaExecution
        """
        execution = SagaExecution.objects.get(pk=saga_id)

        if execution.status not in [SagaStatus.PAUSED, SagaStatus.FAILED]:
            raise ValueError(
                f"Cannot resume saga in status {execution.status}"
            )

        execution.status = SagaStatus.RUNNING
        execution.save(update_fields=['status', 'updated_at'])

        self._log(execution, None, SagaLogEventType.SAGA_RESUMED, 'Saga resumed')

        return self._run_steps(execution)

    def _run_steps(self, execution: SagaExecution) -> SagaExecution:
        """Execute pending steps in order."""
        execution.status = SagaStatus.RUNNING
        execution.started_at = execution.started_at or timezone.now()
        execution.save(update_fields=['status', 'started_at', 'updated_at'])

        steps = execution.steps.filter(
            status__in=[StepStatus.PENDING, StepStatus.FAILED]
        ).order_by('order')

        for step in steps:
            success = self._execute_step(execution, step)
            if not success:
                if step.is_critical:
                    # Critical step failed → start compensation
                    execution.status = SagaStatus.COMPENSATING
                    execution.error_message = step.error_message
                    execution.save(update_fields=['status', 'error_message', 'updated_at'])
                    return self._run_compensation(execution, step.order)
                else:
                    # Non-critical step failed → skip and continue
                    step.status = StepStatus.SKIPPED
                    step.save(update_fields=['status'])
                    logger.warning(
                        f"Non-critical step {step.name} failed, skipping. "
                        f"Error: {step.error_message}"
                    )

        # All steps completed
        execution.status = SagaStatus.COMPLETED
        execution.completed_at = timezone.now()
        execution.save(update_fields=['status', 'completed_at', 'updated_at'])

        self._log(execution, None, SagaLogEventType.SAGA_COMPLETED, 'Saga completed successfully')

        logger.info(f"Saga {execution.id} ({execution.name}) completed successfully")
        return execution

    def _execute_step(self, execution: SagaExecution, step: SagaStep) -> bool:
        """
        Execute a single step with retry and exponential backoff.
        Returns True if the step completed successfully.
        """
        step.status = StepStatus.RUNNING
        step.started_at = timezone.now()
        step.save(update_fields=['status', 'started_at'])

        self._log(
            execution, step, SagaLogEventType.STEP_STARTED,
            f'Executing step: {step.name}'
        )

        while step.retry_count <= step.max_retries:
            try:
                # Resolve and execute the action
                action_func = self._resolve_action(step.action_type)
                result = action_func(execution.context, step.input_data)

                # Store output
                if isinstance(result, dict):
                    step.output_data = result
                    # Merge result into saga context for subsequent steps
                    execution.context.update(result)
                    execution.save(update_fields=['context', 'updated_at'])

                step.status = StepStatus.COMPLETED
                step.completed_at = timezone.now()
                step.save(update_fields=[
                    'status', 'completed_at', 'output_data'
                ])

                execution.current_step_order = step.order
                execution.save(update_fields=['current_step_order', 'updated_at'])

                self._log(
                    execution, step, SagaLogEventType.STEP_COMPLETED,
                    f'Step completed: {step.name}',
                    payload=step.output_data,
                )

                return True

            except Exception as e:
                step.retry_count += 1
                step.error_message = str(e)
                step.save(update_fields=['retry_count', 'error_message'])

                error_details = {
                    'error': str(e),
                    'traceback': traceback.format_exc(),
                    'retry_count': step.retry_count,
                    'max_retries': step.max_retries,
                }

                if step.retry_count <= step.max_retries:
                    # Retry with exponential backoff
                    backoff_seconds = min(2 ** step.retry_count, 60)
                    self._log(
                        execution, step, SagaLogEventType.STEP_RETRYING,
                        f'Retrying step {step.name} in {backoff_seconds}s '
                        f'(attempt {step.retry_count}/{step.max_retries})',
                        error_details=error_details,
                    )
                    time.sleep(backoff_seconds)
                else:
                    # Max retries reached → step failed
                    step.status = StepStatus.FAILED
                    step.save(update_fields=['status'])

                    self._log(
                        execution, step, SagaLogEventType.STEP_FAILED,
                        f'Step failed after {step.max_retries} retries: {step.name}',
                        error_details=error_details,
                    )

                    logger.error(
                        f"Saga step {step.name} failed after {step.max_retries} "
                        f"retries in saga {execution.id}: {e}"
                    )
                    return False

        return False

    def _run_compensation(
        self,
        execution: SagaExecution,
        from_step_order: Optional[int] = None,
    ) -> SagaExecution:
        """
        Run compensation for completed steps in reverse order.
        """
        self._log(
            execution, None, SagaLogEventType.SAGA_COMPENSATING,
            f'Starting compensation from step order {from_step_order}'
        )

        # Get completed steps in reverse order up to from_step_order
        completed_steps = execution.steps.filter(
            status=StepStatus.COMPLETED,
        ).order_by('-order')

        if from_step_order:
            completed_steps = completed_steps.filter(order__lt=from_step_order)

        all_compensated = True

        for step in completed_steps:
            if not step.compensation_type:
                logger.warning(
                    f"No compensation defined for step {step.name}, skipping"
                )
                continue

            success = self._compensate_step(execution, step)
            if not success:
                all_compensated = False
                logger.error(
                    f"Compensation failed for step {step.name} in saga {execution.id}"
                )

        if all_compensated:
            execution.status = SagaStatus.COMPENSATED
            execution.completed_at = timezone.now()
            execution.save(update_fields=['status', 'completed_at', 'updated_at'])
            self._log(
                execution, None, SagaLogEventType.SAGA_COMPENSATED,
                'All compensation steps completed successfully'
            )
        else:
            execution.status = SagaStatus.FAILED
            execution.error_message = (
                f'{execution.error_message or ""}\nCompensation partially failed'
            ).strip()
            execution.save(update_fields=['status', 'error_message', 'updated_at'])
            self._log(
                execution, None, SagaLogEventType.SAGA_FAILED,
                'Compensation partially failed — manual intervention required'
            )

        return execution

    def _compensate_step(self, execution: SagaExecution, step: SagaStep) -> bool:
        """Execute compensation for a single step."""
        step.status = StepStatus.COMPENSATING
        step.save(update_fields=['status'])

        self._log(
            execution, step, SagaLogEventType.STEP_COMPENSATING,
            f'Compensating step: {step.name}'
        )

        try:
            comp_func = self._resolve_action(step.compensation_type)
            comp_func(execution.context, step.output_data)

            step.status = StepStatus.COMPENSATED
            step.save(update_fields=['status'])

            self._log(
                execution, step, SagaLogEventType.STEP_COMPENSATED,
                f'Step compensated: {step.name}'
            )

            return True

        except Exception as e:
            step.status = StepStatus.FAILED
            step.error_message = f'Compensation error: {e}'
            step.save(update_fields=['status', 'error_message'])

            self._log(
                execution, step, SagaLogEventType.STEP_FAILED,
                f'Compensation failed for step {step.name}',
                error_details={
                    'error': str(e),
                    'traceback': traceback.format_exc(),
                },
            )

            logger.error(
                f"Compensation failed for step {step.name} in saga {execution.id}: {e}"
            )
            return False

    @staticmethod
    def _log(
        execution: SagaExecution,
        step: Optional[SagaStep],
        event_type: str,
        message: str,
        payload: Optional[Dict] = None,
        error_details: Optional[Dict] = None,
    ) -> SagaLog:
        """Create a saga log entry."""
        return SagaLog.objects.create(
            saga=execution,
            step=step,
            event_type=event_type,
            message=message,
            payload=payload or {},
            error_details=error_details or {},
        )