"""State recovery manager and utilities."""

import asyncio
import json
import logging
from contextlib import asynccontextmanager
from datetime import datetime, timezone
from typing import Any, Callable, Dict, List, Optional

from ...errors import TraxError, ErrorCode, ErrorCategory, ErrorSeverity
from ...logging import get_logger
from .models import StateStorage, OperationState

logger = get_logger(__name__)


class StateRecoveryManager:
    """Manages state recovery for operations."""
    
    def __init__(self, storage: StateStorage):
        self.storage = storage
        self.active_operations: Dict[str, OperationState] = {}
        self.logger = get_logger(f"{__name__}.StateRecoveryManager")
    
    async def start_operation(
        self,
        operation_id: str,
        correlation_id: str,
        operation_type: str,
        metadata: Optional[Dict[str, Any]] = None
    ) -> OperationState:
        """Start tracking an operation."""
        state = OperationState(
            operation_id=operation_id,
            correlation_id=correlation_id,
            operation_type=operation_type,
            status='running',
            start_time=datetime.now(timezone.utc),
            last_update=datetime.now(timezone.utc),
            metadata=metadata or {}
        )
        
        self.active_operations[operation_id] = state
        await self.storage.save_state(state)
        
        self.logger.info(f"Started tracking operation: {operation_id}")
        return state
    
    async def update_operation(
        self,
        operation_id: str,
        progress: Optional[float] = None,
        checkpoint_data: Optional[Dict[str, Any]] = None,
        metadata: Optional[Dict[str, Any]] = None
    ) -> bool:
        """Update operation state."""
        if operation_id not in self.active_operations:
            self.logger.warning(f"Operation not found: {operation_id}")
            return False
        
        state = self.active_operations[operation_id]
        state.last_update = datetime.now(timezone.utc)
        
        if progress is not None:
            state.progress = max(0.0, min(1.0, progress))
        
        if checkpoint_data is not None:
            state.checkpoint_data.update(checkpoint_data)
        
        if metadata is not None:
            state.metadata.update(metadata)
        
        await self.storage.save_state(state)
        return True
    
    async def complete_operation(
        self,
        operation_id: str,
        result: Any = None,
        metadata: Optional[Dict[str, Any]] = None
    ) -> bool:
        """Mark operation as completed."""
        if operation_id not in self.active_operations:
            self.logger.warning(f"Operation not found: {operation_id}")
            return False
        
        state = self.active_operations[operation_id]
        state.status = 'completed'
        state.progress = 1.0
        state.result = result
        state.last_update = datetime.now(timezone.utc)
        
        if metadata is not None:
            state.metadata.update(metadata)
        
        await self.storage.save_state(state)
        del self.active_operations[operation_id]
        
        self.logger.info(f"Completed operation: {operation_id}")
        return True
    
    async def fail_operation(
        self,
        operation_id: str,
        error: Exception,
        metadata: Optional[Dict[str, Any]] = None
    ) -> bool:
        """Mark operation as failed."""
        if operation_id not in self.active_operations:
            self.logger.warning(f"Operation not found: {operation_id}")
            return False
        
        state = self.active_operations[operation_id]
        state.status = 'failed'
        state.last_update = datetime.now(timezone.utc)
        state.error_info = {
            'error_type': type(error).__name__,
            'error_message': str(error),
            'timestamp': datetime.now(timezone.utc).isoformat()
        }
        
        if metadata is not None:
            state.metadata.update(metadata)
        
        await self.storage.save_state(state)
        del self.active_operations[operation_id]
        
        self.logger.info(f"Failed operation: {operation_id}")
        return True
    
    async def interrupt_operation(
        self,
        operation_id: str,
        checkpoint_data: Optional[Dict[str, Any]] = None
    ) -> bool:
        """Mark operation as interrupted (for recovery)."""
        if operation_id not in self.active_operations:
            self.logger.warning(f"Operation not found: {operation_id}")
            return False
        
        state = self.active_operations[operation_id]
        state.status = 'interrupted'
        state.last_update = datetime.now(timezone.utc)
        
        if checkpoint_data is not None:
            state.checkpoint_data.update(checkpoint_data)
        
        await self.storage.save_state(state)
        del self.active_operations[operation_id]
        
        self.logger.info(f"Interrupted operation: {operation_id}")
        return True
    
    async def recover_operation(self, operation_id: str) -> Optional[OperationState]:
        """Recover an interrupted operation."""
        state = await self.storage.load_state(operation_id)
        if state is None or state.status != 'interrupted':
            return None
        
        # Mark as running again
        state.status = 'running'
        state.last_update = datetime.now(timezone.utc)
        self.active_operations[operation_id] = state
        await self.storage.save_state(state)
        
        self.logger.info(f"Recovered operation: {operation_id}")
        return state
    
    async def list_interrupted_operations(self) -> List[OperationState]:
        """List all interrupted operations."""
        operation_ids = await self.storage.list_states()
        interrupted_states = []
        
        for operation_id in operation_ids:
            state = await self.storage.load_state(operation_id)
            if state and state.status == 'interrupted':
                interrupted_states.append(state)
        
        return interrupted_states
    
    async def cleanup_completed_operations(self, max_age_hours: int = 24) -> int:
        """Clean up completed operations older than specified age."""
        return await self.storage.cleanup_expired_states(max_age_hours)


# Context manager for state tracking
@asynccontextmanager
async def operation_state_context(
    recovery_manager: StateRecoveryManager,
    operation_id: str,
    correlation_id: str,
    operation_type: str,
    metadata: Optional[Dict[str, Any]] = None
):
    """Context manager for tracking operation state."""
    state = await recovery_manager.start_operation(
        operation_id, correlation_id, operation_type, metadata
    )
    
    try:
        yield state
        await recovery_manager.complete_operation(operation_id)
    except Exception as e:
        await recovery_manager.fail_operation(operation_id, e)
        raise


# Decorator for automatic state tracking
def with_state_tracking(
    recovery_manager: StateRecoveryManager,
    operation_type: str,
    metadata: Optional[Dict[str, Any]] = None
):
    """Decorator to add state tracking to functions."""
    def decorator(func: Callable) -> Callable:
        async def wrapper(*args, **kwargs):
            # Generate operation ID from function name and args
            operation_id = f"{func.__name__}_{hash(str(args) + str(kwargs))}"
            correlation_id = kwargs.get('correlation_id', 'default')
            
            async with operation_state_context(
                recovery_manager, operation_id, correlation_id, operation_type, metadata
            ) as state:
                # Update progress periodically
                result = await func(*args, **kwargs)
                await recovery_manager.update_operation(operation_id, progress=1.0)
                return result
        
        return wrapper
    return decorator


# Utility functions for common state recovery scenarios
async def create_state_recovery_manager(storage: StateStorage) -> StateRecoveryManager:
    """Create a state recovery manager."""
    return StateRecoveryManager(storage)


async def recover_interrupted_operations(recovery_manager: StateRecoveryManager) -> List[OperationState]:
    """Recover all interrupted operations."""
    interrupted_states = await recovery_manager.list_interrupted_operations()
    recovered_states = []
    
    for state in interrupted_states:
        try:
            recovered_state = await recovery_manager.recover_operation(state.operation_id)
            if recovered_state:
                recovered_states.append(recovered_state)
        except Exception as e:
            logger.error(f"Failed to recover operation {state.operation_id}: {e}")
    
    return recovered_states
