"""
Workflow State Manager.

Manages workflow state persistence, snapshots, and replay.
"""
import logging
from typing import Any, Dict, List, Optional
from uuid import UUID

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

from ..models import ProcessInstance, ProcessVariable, Task, InstanceStatus, TaskStatus
from .models import (
    WorkflowExecutionLog,
    WorkflowSnapshot,
    ExecutionLogEventType,
    SnapshotReason,
)

logger = logging.getLogger(__name__)


class WorkflowStateManager:
    """
    Manages workflow execution state with snapshot and replay capabilities.
    
    Features:
    - Save/load workflow state
    - Event logging with sequence numbers
    - Automatic and manual snapshots
    - Restore from snapshot
    - Replay from event log
    - Export/Import for migration
    
    Example:
        manager = WorkflowStateManager()
        
        # Log an event
        manager.log_event(
            execution_id=instance.id,
            event_type=ExecutionLogEventType.STEP_COMPLETED,
            step_id='activity_1',
            payload={'result': 'approved'},
        )
        
        # Create a snapshot
        snapshot = manager.create_snapshot(
            execution_id=instance.id,
            reason=SnapshotReason.BEFORE_CRITICAL_STEP,
        )
        
        # Restore from snapshot
        manager.restore_from_snapshot(snapshot.id)
    """

    def save_state(
        self,
        execution_id: UUID,
        state: Dict[str, Any],
    ) -> ProcessInstance:
        """
        Save the current state of a workflow execution.
        
        Args:
            execution_id: Process instance UUID
            state: State dictionary containing status, current_element_ids, context etc.
            
        Returns:
            Updated ProcessInstance
        """
        instance = ProcessInstance.objects.get(pk=execution_id)
        
        if 'status' in state:
            instance.status = state['status']
        if 'current_element_ids' in state:
            instance.current_element_ids = state['current_element_ids']
        if 'context' in state:
            instance.context = state['context']
        if 'error_message' in state:
            instance.error_message = state['error_message']

        update_fields = ['status', 'current_element_ids', 'context', 'error_message']
        instance.save(update_fields=update_fields)

        logger.debug(f"State saved for execution {execution_id}")
        return instance

    def get_state(self, execution_id: UUID) -> Dict[str, Any]:
        """
        Get the current complete state of a workflow execution.
        
        Returns:
            Dictionary with full execution state
        """
        instance = ProcessInstance.objects.get(pk=execution_id)
        variables = ProcessVariable.objects.filter(instance=instance)
        tasks = Task.objects.filter(instance=instance)

        return {
            'execution_id': str(instance.id),
            'definition_id': str(instance.definition_id),
            'status': instance.status,
            'current_element_ids': instance.current_element_ids,
            'context': instance.context,
            'business_key': instance.business_key,
            'error_message': instance.error_message,
            'started_at': instance.started_at.isoformat() if instance.started_at else None,
            'variables': {
                var.name: {
                    'value': var.value,
                    'type': var.type,
                    'scope': var.scope,
                }
                for var in variables
            },
            'tasks': [
                {
                    'id': str(task.id),
                    'task_definition_key': task.task_definition_key,
                    'element_id': task.element_id,
                    'name': task.name,
                    'status': task.status,
                    'task_type': task.task_type,
                    'assignee_id': str(task.assignee_id) if task.assignee_id else None,
                    'priority': task.priority,
                    'form_data': task.form_data,
                    'input_variables': task.input_variables,
                    'output_variables': task.output_variables,
                    'created_at': task.created_at.isoformat() if task.created_at else None,
                    'completed_at': task.completed_at.isoformat() if task.completed_at else None,
                }
                for task in tasks
            ],
        }

    def log_event(
        self,
        execution_id: UUID,
        event_type: str,
        step_id: Optional[str] = None,
        payload: Optional[Dict[str, Any]] = None,
        user=None,
        duration_ms: Optional[int] = None,
        error_message: Optional[str] = None,
        metadata: Optional[Dict[str, Any]] = None,
    ) -> WorkflowExecutionLog:
        """
        Log an event in the execution log.
        
        Automatically assigns sequence number and captures state before/after.
        
        Args:
            execution_id: Process instance UUID
            event_type: Type of event (from ExecutionLogEventType)
            step_id: Optional BPMN element ID
            payload: Event-specific data
            user: User who triggered the event
            duration_ms: Duration of the operation in milliseconds
            error_message: Error message if applicable
            metadata: Additional metadata
            
        Returns:
            Created WorkflowExecutionLog
        """
        instance = ProcessInstance.objects.get(pk=execution_id)

        # Get next sequence number
        last_seq = WorkflowExecutionLog.objects.filter(
            execution=instance
        ).order_by('-sequence_number').values_list(
            'sequence_number', flat=True
        ).first() or 0
        next_seq = last_seq + 1

        # Capture current state
        current_state = {
            'status': instance.status,
            'current_element_ids': instance.current_element_ids,
        }

        log_entry = WorkflowExecutionLog.objects.create(
            execution=instance,
            event_type=event_type,
            step_id=step_id,
            sequence_number=next_seq,
            payload=payload or {},
            state_before=current_state,
            state_after={},  # Will be updated after state change if needed
            user=user,
            duration_ms=duration_ms,
            error_message=error_message,
            metadata=metadata or {},
        )

        logger.debug(
            f"Logged event #{next_seq} ({event_type}) for execution {execution_id}"
        )
        return log_entry

    @transaction.atomic
    def create_snapshot(
        self,
        execution_id: UUID,
        reason: str = SnapshotReason.MANUAL,
        description: Optional[str] = None,
        user=None,
    ) -> WorkflowSnapshot:
        """
        Create a full snapshot of the workflow execution state.
        
        Args:
            execution_id: Process instance UUID
            reason: Reason for creating the snapshot
            description: Optional description
            user: User creating the snapshot
            
        Returns:
            Created WorkflowSnapshot
        """
        state = self.get_state(execution_id)
        instance = ProcessInstance.objects.get(pk=execution_id)

        # Get last event sequence
        last_seq = WorkflowExecutionLog.objects.filter(
            execution=instance
        ).order_by('-sequence_number').values_list(
            'sequence_number', flat=True
        ).first() or 0

        snapshot = WorkflowSnapshot.objects.create(
            execution=instance,
            state_data={
                'status': state['status'],
                'current_element_ids': state['current_element_ids'],
                'context': state['context'],
                'business_key': state['business_key'],
                'error_message': state['error_message'],
                'started_at': state['started_at'],
            },
            variables_data=state['variables'],
            tasks_data=state['tasks'],
            snapshot_reason=reason,
            description=description,
            last_event_sequence=last_seq,
            created_by=user,
        )

        # Log the snapshot event
        self.log_event(
            execution_id=execution_id,
            event_type=ExecutionLogEventType.SNAPSHOT_CREATED,
            payload={
                'snapshot_id': str(snapshot.id),
                'reason': reason,
            },
            user=user,
        )

        logger.info(
            f"Snapshot created for execution {execution_id}: {snapshot.id} ({reason})"
        )
        return snapshot

    def auto_snapshot_before_critical_step(
        self,
        execution_id: UUID,
        step_id: str,
    ) -> WorkflowSnapshot:
        """
        Automatically create a snapshot before a critical step.
        
        Args:
            execution_id: Process instance UUID
            step_id: BPMN element ID of the critical step
        """
        return self.create_snapshot(
            execution_id=execution_id,
            reason=SnapshotReason.BEFORE_CRITICAL_STEP,
            description=f'Auto-snapshot before critical step: {step_id}',
        )

    @transaction.atomic
    def restore_from_snapshot(
        self,
        snapshot_id: UUID,
        user=None,
    ) -> ProcessInstance:
        """
        Restore a workflow execution to a snapshot state.
        
        Args:
            snapshot_id: Snapshot UUID
            user: User performing the restore
            
        Returns:
            Restored ProcessInstance
        """
        snapshot = WorkflowSnapshot.objects.select_related('execution').get(
            pk=snapshot_id
        )
        instance = snapshot.execution
        state_data = snapshot.state_data

        # Restore instance state
        instance.status = state_data.get('status', InstanceStatus.RUNNING)
        instance.current_element_ids = state_data.get('current_element_ids', [])
        instance.context = state_data.get('context', {})
        instance.error_message = state_data.get('error_message')
        instance.save(update_fields=[
            'status', 'current_element_ids', 'context', 'error_message'
        ])

        # Restore variables
        ProcessVariable.objects.filter(instance=instance).delete()
        for var_name, var_data in snapshot.variables_data.items():
            ProcessVariable.objects.create(
                instance=instance,
                name=var_name,
                value=var_data.get('value'),
                type=var_data.get('type', 'string'),
                scope=var_data.get('scope'),
            )

        # Log the restore event
        self.log_event(
            execution_id=instance.id,
            event_type=ExecutionLogEventType.SNAPSHOT_RESTORED,
            payload={
                'snapshot_id': str(snapshot.id),
                'restored_to_sequence': snapshot.last_event_sequence,
            },
            user=user,
        )

        logger.info(
            f"Execution {instance.id} restored from snapshot {snapshot.id}"
        )
        return instance

    def replay_from_log(
        self,
        execution_id: UUID,
        from_event_sequence: int = 0,
        to_event_sequence: Optional[int] = None,
        dry_run: bool = True,
    ) -> List[Dict[str, Any]]:
        """
        Replay execution from event log for debugging and audit.
        
        In dry_run mode, returns the sequence of events and computed states
        without actually modifying the execution. When dry_run=False,
        actually restores and replays (use with caution).
        
        Args:
            execution_id: Process instance UUID
            from_event_sequence: Starting sequence number (inclusive)
            to_event_sequence: Ending sequence number (inclusive), None for all
            dry_run: If True, only compute states without modifying
            
        Returns:
            List of event records with computed states
        """
        instance = ProcessInstance.objects.get(pk=execution_id)

        log_qs = WorkflowExecutionLog.objects.filter(
            execution=instance,
            sequence_number__gte=from_event_sequence,
        ).order_by('sequence_number')

        if to_event_sequence is not None:
            log_qs = log_qs.filter(sequence_number__lte=to_event_sequence)

        replay_results = []
        for log_entry in log_qs:
            replay_results.append({
                'sequence_number': log_entry.sequence_number,
                'event_type': log_entry.event_type,
                'step_id': log_entry.step_id,
                'payload': log_entry.payload,
                'state_before': log_entry.state_before,
                'state_after': log_entry.state_after,
                'timestamp': log_entry.timestamp.isoformat(),
                'user_id': str(log_entry.user_id) if log_entry.user_id else None,
                'duration_ms': log_entry.duration_ms,
                'error_message': log_entry.error_message,
            })

        logger.info(
            f"Replayed {len(replay_results)} events for execution {execution_id} "
            f"(from_seq={from_event_sequence}, to_seq={to_event_sequence}, dry_run={dry_run})"
        )

        return replay_results

    def get_execution_timeline(
        self,
        execution_id: UUID,
    ) -> List[Dict[str, Any]]:
        """
        Get a human-readable timeline of the execution.
        
        Returns:
            List of timeline entries
        """
        logs = WorkflowExecutionLog.objects.filter(
            execution_id=execution_id
        ).order_by('sequence_number').select_related('user')

        timeline = []
        for log in logs:
            timeline.append({
                'sequence': log.sequence_number,
                'event': log.get_event_type_display(),
                'event_type': log.event_type,
                'step_id': log.step_id,
                'timestamp': log.timestamp.isoformat(),
                'user': log.user.email if log.user else None,
                'duration_ms': log.duration_ms,
                'has_error': bool(log.error_message),
                'payload_summary': {
                    k: v for k, v in (log.payload or {}).items()
                    if k not in ['_internal']
                },
            })

        return timeline

    def export_execution(self, execution_id: UUID) -> Dict[str, Any]:
        """
        Export a complete execution (state + logs + snapshots) for migration.
        
        Returns:
            Complete execution data as a dictionary
        """
        state = self.get_state(execution_id)
        instance = ProcessInstance.objects.get(pk=execution_id)

        logs = list(
            WorkflowExecutionLog.objects.filter(
                execution=instance
            ).order_by('sequence_number').values(
                'sequence_number', 'event_type', 'step_id', 'payload',
                'state_before', 'state_after', 'duration_ms',
                'error_message', 'metadata', 'timestamp',
            )
        )
        # Serialize datetime
        for log in logs:
            if log.get('timestamp'):
                log['timestamp'] = log['timestamp'].isoformat()

        snapshots = list(
            WorkflowSnapshot.objects.filter(
                execution=instance
            ).order_by('created_at').values(
                'id', 'state_data', 'variables_data', 'tasks_data',
                'snapshot_reason', 'description', 'last_event_sequence',
                'created_at',
            )
        )
        for snap in snapshots:
            if snap.get('created_at'):
                snap['created_at'] = snap['created_at'].isoformat()
            snap['id'] = str(snap['id'])

        return {
            'version': '1.0',
            'exported_at': timezone.now().isoformat(),
            'execution': state,
            'event_log': logs,
            'snapshots': snapshots,
        }

    @transaction.atomic
    def import_execution(
        self,
        data: Dict[str, Any],
        tenant,
    ) -> ProcessInstance:
        """
        Import an execution from exported data.
        
        Args:
            data: Exported execution data (from export_execution)
            tenant: Target tenant
            
        Returns:
            Created ProcessInstance
        """
        execution_data = data['execution']

        instance = ProcessInstance.objects.get(pk=execution_data['execution_id'])

        # Import event logs
        for log_data in data.get('event_log', []):
            WorkflowExecutionLog.objects.create(
                execution=instance,
                event_type=log_data['event_type'],
                step_id=log_data.get('step_id'),
                sequence_number=log_data['sequence_number'],
                payload=log_data.get('payload', {}),
                state_before=log_data.get('state_before', {}),
                state_after=log_data.get('state_after', {}),
                duration_ms=log_data.get('duration_ms'),
                error_message=log_data.get('error_message'),
                metadata=log_data.get('metadata', {}),
            )

        logger.info(f"Imported execution data for {instance.id}")
        return instance