"""Storage implementations for state recovery."""

import asyncio
import json
import logging
import os
import pickle
from datetime import datetime, timezone
from pathlib import Path
from typing import Any, Callable, Dict, List, Optional
import hashlib

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

logger = get_logger(__name__)


class FileStateStorage(StateStorage):
    """File-based state storage implementation."""
    
    def __init__(self, storage_dir: str = "data/state"):
        super().__init__("file_storage")
        self.storage_dir = Path(storage_dir)
        self.storage_dir.mkdir(parents=True, exist_ok=True)
    
    def _get_state_file(self, operation_id: str) -> Path:
        """Get the file path for an operation state."""
        # Use hash to avoid filesystem issues with special characters
        safe_id = hashlib.md5(operation_id.encode()).hexdigest()
        return self.storage_dir / f"{safe_id}.json"
    
    async def save_state(self, state: OperationState) -> bool:
        """Save operation state to file."""
        try:
            state_file = self._get_state_file(state.operation_id)
            state_data = state.to_dict()
            
            # Create temporary file first
            temp_file = state_file.with_suffix('.tmp')
            with open(temp_file, 'w') as f:
                json.dump(state_data, f, indent=2)
            
            # Atomic move
            temp_file.replace(state_file)
            
            self.logger.debug(f"Saved state for operation: {state.operation_id}")
            return True
        except Exception as e:
            self.logger.error(f"Failed to save state for {state.operation_id}: {e}")
            return False
    
    async def load_state(self, operation_id: str) -> Optional[OperationState]:
        """Load operation state from file."""
        try:
            state_file = self._get_state_file(operation_id)
            if not state_file.exists():
                return None
            
            with open(state_file, 'r') as f:
                state_data = json.load(f)
            
            state = OperationState.from_dict(state_data)
            self.logger.debug(f"Loaded state for operation: {operation_id}")
            return state
        except Exception as e:
            self.logger.error(f"Failed to load state for {operation_id}: {e}")
            return None
    
    async def delete_state(self, operation_id: str) -> bool:
        """Delete operation state file."""
        try:
            state_file = self._get_state_file(operation_id)
            if state_file.exists():
                state_file.unlink()
                self.logger.debug(f"Deleted state for operation: {operation_id}")
            return True
        except Exception as e:
            self.logger.error(f"Failed to delete state for {operation_id}: {e}")
            return False
    
    async def list_states(self, operation_type: Optional[str] = None) -> List[str]:
        """List available operation IDs."""
        try:
            operation_ids = []
            for state_file in self.storage_dir.glob("*.json"):
                if state_file.name.endswith('.tmp'):
                    continue
                
                try:
                    with open(state_file, 'r') as f:
                        state_data = json.load(f)
                    
                    if operation_type is None or state_data.get('operation_type') == operation_type:
                        operation_ids.append(state_data['operation_id'])
                except Exception as e:
                    self.logger.warning(f"Failed to read state file {state_file}: {e}")
            
            return operation_ids
        except Exception as e:
            self.logger.error(f"Failed to list states: {e}")
            return []
    
    async def cleanup_expired_states(self, max_age_hours: int = 24) -> int:
        """Clean up expired state files."""
        try:
            cutoff_time = datetime.now(timezone.utc).timestamp() - (max_age_hours * 3600)
            deleted_count = 0
            
            for state_file in self.storage_dir.glob("*.json"):
                if state_file.name.endswith('.tmp'):
                    continue
                
                try:
                    file_mtime = state_file.stat().st_mtime
                    if file_mtime < cutoff_time:
                        state_file.unlink()
                        deleted_count += 1
                except Exception as e:
                    self.logger.warning(f"Failed to check/delete {state_file}: {e}")
            
            self.logger.info(f"Cleaned up {deleted_count} expired state files")
            return deleted_count
        except Exception as e:
            self.logger.error(f"Failed to cleanup expired states: {e}")
            return 0


# Utility functions for common state storage scenarios
async def create_file_state_storage(storage_dir: str = "data/state") -> FileStateStorage:
    """Create a file-based state storage."""
    return FileStateStorage(storage_dir)
