"""Data models and storage base class for state recovery."""

import asyncio
import json
import logging
import os
import pickle
from dataclasses import dataclass, field, asdict
from datetime import datetime, timezone
from pathlib import Path
from typing import Any, Callable, Dict, List, Optional, TypeVar, Union
from contextlib import asynccontextmanager
import hashlib

from ...errors import TraxError, ErrorCode, ErrorCategory, ErrorSeverity
from ...logging import get_logger

logger = get_logger(__name__)

T = TypeVar('T')


@dataclass
class OperationState:
    """State information for an operation."""
    operation_id: str
    correlation_id: str
    operation_type: str
    status: str  # 'running', 'completed', 'failed', 'interrupted'
    start_time: datetime
    last_update: datetime
    progress: float = 0.0  # 0.0 to 1.0
    checkpoint_data: Dict[str, Any] = field(default_factory=dict)
    metadata: Dict[str, Any] = field(default_factory=dict)
    error_info: Optional[Dict[str, Any]] = None
    result: Optional[Any] = None
    
    def to_dict(self) -> Dict[str, Any]:
        """Convert to dictionary for serialization."""
        data = asdict(self)
        data['start_time'] = self.start_time.isoformat()
        data['last_update'] = self.last_update.isoformat()
        return data
    
    @classmethod
    def from_dict(cls, data: Dict[str, Any]) -> 'OperationState':
        """Create from dictionary."""
        data = data.copy()
        data['start_time'] = datetime.fromisoformat(data['start_time'])
        data['last_update'] = datetime.fromisoformat(data['last_update'])
        return cls(**data)


class StateStorage:
    """Base class for state storage implementations."""
    
    def __init__(self, storage_id: str):
        self.storage_id = storage_id
        self.logger = get_logger(f"{__name__}.{storage_id}")
    
    async def save_state(self, state: OperationState) -> bool:
        """Save operation state."""
        raise NotImplementedError
    
    async def load_state(self, operation_id: str) -> Optional[OperationState]:
        """Load operation state."""
        raise NotImplementedError
    
    async def delete_state(self, operation_id: str) -> bool:
        """Delete operation state."""
        raise NotImplementedError
    
    async def list_states(self, operation_type: Optional[str] = None) -> List[str]:
        """List available operation IDs."""
        raise NotImplementedError
    
    async def cleanup_expired_states(self, max_age_hours: int = 24) -> int:
        """Clean up expired state files."""
        raise NotImplementedError
