"""
Automatic state management features.
Handles automatic syncing, cleanup, retries, and health checks.
"""

from typing import Optional, Callable, Dict, Any
from datetime import datetime, timedelta
import time
import threading
from functools import wraps

from .client import UASSClient
from .backends import StateBackend


class AutoManager:
    """
    Automatic state management wrapper around UASSClient.
    
    Adds automatic features:
    - Auto-retry on failures
    - TTL/expiration for states
    - Health checks and reconnection
    - State validation
    - Background cleanup
    """
    
    def __init__(
        self,
        client: UASSClient,
        auto_retry: bool = True,
        max_retries: int = 3,
        retry_delay: float = 1.0,
        ttl_hours: Optional[int] = None,
        cleanup_interval_hours: int = 24,
        enable_health_check: bool = True,
    ):
        """
        Initialize auto manager.
        
        Args:
            client: UASSClient instance
            auto_retry: Enable automatic retries
            max_retries: Maximum retry attempts
            retry_delay: Delay between retries (seconds)
            ttl_hours: Auto-delete states older than this (None = no expiration)
            cleanup_interval_hours: How often to run cleanup (hours)
            enable_health_check: Enable connection health checks
        """
        self.client = client
        self.auto_retry = auto_retry
        self.max_retries = max_retries
        self.retry_delay = retry_delay
        self.ttl_hours = ttl_hours
        self.cleanup_interval = timedelta(hours=cleanup_interval_hours)
        self.enable_health_check = enable_health_check
        
        self._cleanup_thread: Optional[threading.Thread] = None
        self._stop_cleanup = threading.Event()
        
        if ttl_hours:
            self._start_cleanup_thread()
    
    def _retry_on_failure(self, func: Callable, *args, **kwargs):
        """Retry function on failure."""
        if not self.auto_retry:
            return func(*args, **kwargs)
        
        last_error = None
        for attempt in range(self.max_retries):
            try:
                return func(*args, **kwargs)
            except Exception as e:
                last_error = e
                if attempt < self.max_retries - 1:
                    time.sleep(self.retry_delay * (2 ** attempt))  # Exponential backoff
                else:
                    raise last_error
        return None
    
    def write_state(
        self,
        session_id: str,
        agent_id: str,
        framework: str,
        **kwargs
    ) -> bool:
        """
        Write state with automatic retry.
        """
        def _write():
            return self.client.write_state(
                session_id=session_id,
                agent_id=agent_id,
                framework=framework,
                **kwargs
            )
        
        return self._retry_on_failure(_write)
    
    def read_state(
        self,
        session_id: str,
        agent_id: str,
    ) -> Optional[Dict[str, Any]]:
        """
        Read state with automatic retry.
        """
        def _read():
            return self.client.read_state(session_id, agent_id)
        
        return self._retry_on_failure(_read)
    
    def delete_state(self, session_id: str, agent_id: str) -> bool:
        """Delete state with automatic retry."""
        def _delete():
            return self.client.delete_state(session_id, agent_id)
        
        return self._retry_on_failure(_delete)
    
    def _start_cleanup_thread(self):
        """Start background cleanup thread."""
        if self._cleanup_thread and self._cleanup_thread.is_alive():
            return
        
        def cleanup_loop():
            while not self._stop_cleanup.is_set():
                try:
                    self._cleanup_expired_states()
                except Exception as e:
                    print(f"Error in cleanup thread: {e}")
                
                # Wait for next cleanup interval
                self._stop_cleanup.wait(self.cleanup_interval.total_seconds())
        
        self._cleanup_thread = threading.Thread(target=cleanup_loop, daemon=True)
        self._cleanup_thread.start()
    
    def _cleanup_expired_states(self):
        """Clean up expired states (if TTL is set)."""
        if not self.ttl_hours:
            return
        
        # This is a simplified version - would need backend-specific implementation
        # For Redis, we could use TTL
        # For PostgreSQL, we'd query by updated_at
        print(f"Cleanup: Removing states older than {self.ttl_hours} hours")
        # Implementation would depend on backend type
    
    def health_check(self) -> bool:
        """
        Check backend health.
        
        Returns:
            True if backend is healthy, False otherwise
        """
        if not self.enable_health_check:
            return True
        
        try:
            # Try a simple read operation
            # This depends on backend type
            if hasattr(self.client.backend, 'client'):
                # Redis backend
                if hasattr(self.client.backend.client, 'ping'):
                    return self.client.backend.client.ping()
            return True
        except Exception:
            return False
    
    def stop(self):
        """Stop background threads."""
        if self._cleanup_thread:
            self._stop_cleanup.set()
            self._cleanup_thread.join(timeout=5)


class AutoSyncHook:
    """
    Automatic state synchronization hook.
    Can be attached to framework state changes to auto-sync to UASS.
    """
    
    def __init__(self, client: UASSClient, session_id: str, agent_id: str, framework: str):
        """
        Initialize auto-sync hook.
        
        Args:
            client: UASSClient instance
            session_id: Session identifier
            agent_id: Agent identifier
            framework: Framework name
        """
        self.client = client
        self.session_id = session_id
        self.agent_id = agent_id
        self.framework = framework
    
    def sync(self, framework_state: Dict[str, Any]):
        """
        Sync framework state to UASS.
        
        Args:
            framework_state: Framework-specific state
        """
        # Extract common state from framework state
        # This would be framework-specific
        self.client.write_state(
            session_id=self.session_id,
            agent_id=self.agent_id,
            framework=self.framework,
            framework_state=framework_state,
        )

