from pydantic import BaseModel, ConfigDict, Field
from redis.asyncio import (Redis, ConnectionPool)
from typing import Any, Optional, Dict, List, Optional, TypeGuard, Union
import json
import logging
import os
from datetime import datetime, timezone, UTC, timedelta
import asyncio
from models import (
    # User models
    Candidate, Employer, BaseUser, EvidenceDetail, Guest, Authentication, AuthResponse, SkillAssessment,
)
import backstory_traceback as traceback
from .constants import KEY_PREFIXES
from .core import RedisDatabase

logger = logging.getLogger(__name__)

# _RedisManager is a singleton class that manages the Redis connection and 
# provides methods for connecting, disconnecting, and performing health checks.
#
# It uses connection pooling for better performance and resource management.
class _RedisManager:
    def __init__(self):
        self.redis: Optional[Redis] = None
        self.redis_url = os.getenv("REDIS_URL", "redis://redis:6379")
        self.redis_db = int(os.getenv("REDIS_DB", "0"))
    
        # Append database to URL if not already present
        if not self.redis_url.endswith(f"/{self.redis_db}"):
            self.redis_url = f"{self.redis_url}/{self.redis_db}"
            
        self._connection_pool: Optional[ConnectionPool] = None
        self._is_connected = False
    
    async def connect(self):
        """Initialize Redis connection with connection pooling"""
        if self._is_connected and self.redis:
            logger.info("Redis already connected")
            return

        try:
            # Create connection pool for better resource management
            self._connection_pool = ConnectionPool.from_url(
                self.redis_url,
                encoding="utf-8",
                decode_responses=True,
                max_connections=20,
                retry_on_timeout=True,
                socket_keepalive=True,
                socket_keepalive_options={},
                health_check_interval=30
            )
            
            self.redis = Redis(
                connection_pool=self._connection_pool
            )
            
            if not self.redis:
                raise RuntimeError("Redis client not initialized")

            # Test connection
            await self.redis.ping()
            self._is_connected = True
            logger.info("Successfully connected to Redis")
            
            # Log Redis info
            info = await self.redis.info()
            logger.info(f"Redis version: {info.get('redis_version', 'unknown')}")
            
        except Exception as e:
            logger.error(f"Failed to connect to Redis: {e}")
            self._is_connected = False
            self.redis = None
            self._connection_pool = None
            raise
    
    async def disconnect(self):
        """Close Redis connection gracefully"""
        if not self._is_connected:
            logger.info("Redis already disconnected")
            return
        
        try:
            if self.redis:
                # Wait for any pending operations to complete
                await asyncio.sleep(0.1)
                
                # Close the client
                await self.redis.aclose()
                logger.info("Redis client closed")
            
            if self._connection_pool:
                # Close the connection pool
                await self._connection_pool.aclose()
                logger.info("Redis connection pool closed")
            
            self._is_connected = False
            self.redis = None
            self._connection_pool = None
            
            logger.info("Successfully disconnected from Redis")
            
        except Exception as e:
            logger.error(f"Error during Redis disconnect: {e}")
            # Force cleanup even if there's an error
            self._is_connected = False
            self.redis = None
            self._connection_pool = None
    
    def get_client(self) -> Redis:
        """Get Redis client instance"""
        if not self._is_connected or not self.redis:
            raise RuntimeError("Redis client not initialized or disconnected")
        return self.redis
    
    @property
    def is_connected(self) -> bool:
        """Check if Redis is connected"""
        return self._is_connected and self.redis is not None
    
    async def health_check(self) -> dict:
        """Perform health check on Redis connection"""
        if not self.is_connected:
            return {"status": "disconnected", "error": "Redis not connected"}
        
        if not self.redis:
            raise RuntimeError("Redis client not initialized")

        try:
            # Test basic operations
            await self.redis.ping()
            info = await self.redis.info()
            
            return {
                "status": "healthy",
                "redis_version": info.get("redis_version", "unknown"),
                "uptime_seconds": info.get("uptime_in_seconds", 0),
                "connected_clients": info.get("connected_clients", 0),
                "used_memory_human": info.get("used_memory_human", "unknown"),
                "total_commands_processed": info.get("total_commands_processed", 0)
            }
        except Exception as e:
            logger.error(f"Redis health check failed: {e}")
            return {"status": "error", "error": str(e)}
    
    async def force_save(self, background: bool = True) -> bool:
        """Force Redis to save data to disk"""
        if not self.is_connected:
            logger.warning("Cannot save: Redis not connected")
            return False
        
        try:
            if not self.redis:
                raise RuntimeError("Redis client not initialized")

            if background:
                # Non-blocking background save
                await self.redis.bgsave()
                logger.info("Background save initiated")
            else:
                # Blocking save
                await self.redis.save()
                logger.info("Synchronous save completed")
            return True
        except Exception as e:
            logger.error(f"Redis save failed: {e}")
            return False
    
    async def get_info(self) -> Optional[dict]:
        """Get Redis server information"""
        if not self.is_connected:
            return None
        
        try:
            if not self.redis:
                raise RuntimeError("Redis client not initialized")
            return await self.redis.info()
        except Exception as e:
            logger.error(f"Failed to get Redis info: {e}")
            return None

# Global Redis manager instance
redis_manager = _RedisManager()

# DatabaseManager is an enhanced database manager that provides graceful shutdown capabilities
# It manages the Redis connection, tracks active requests, and allows for data backup before shutdown.
class DatabaseManager:
    """Enhanced database manager with graceful shutdown capabilities"""
   
    def __init__(self):
        self.db: Optional[RedisDatabase] = None
        self._shutdown_initiated = False
        self._active_requests = 0
        self._shutdown_timeout = int(os.getenv("SHUTDOWN_TIMEOUT", "30"))  # seconds
        self._backup_on_shutdown = os.getenv("BACKUP_ON_SHUTDOWN", "false").lower() == "true"
    
    async def initialize(self):
        """Initialize database connection"""
        try:
            # Connect to Redis
            await redis_manager.connect()
            logger.info("Redis connection established")
            
            # Create database instance
            self.db = RedisDatabase(redis_manager.get_client())
            
            # Test connection and log stats
            if not redis_manager.redis:
                raise RuntimeError("Redis client not initialized")
            await redis_manager.redis.ping()
            stats = await self.db.get_stats()
            logger.info(f"Database initialized successfully. Stats: {stats}")
            
            return self.db
            
        except Exception as e:
            logger.error(f"Failed to initialize database: {e}")
            raise
    
    async def backup_data(self) -> Optional[str]:
        """Create a backup of critical data before shutdown"""
        if not self.db:
            return None
            
        try:
            backup_data = {
                "timestamp": datetime.now(UTC).isoformat(),
                "stats": await self.db.get_stats(),
                "users": await self.db.get_all_users(),
                # Add other critical data as needed
            }
            
            backup_filename = f"backup_{datetime.now(UTC).strftime('%Y%m%d_%H%M%S')}.json"
            
            # Save to local file (you might want to save to cloud storage instead)
            with open(backup_filename, 'w') as f:
                json.dump(backup_data, f, indent=2, default=str)
            
            logger.info(f"Backup created: {backup_filename}")
            return backup_filename
            
        except Exception as e:
            logger.error(f"Backup failed: {e}")
            return None
    
    async def graceful_shutdown(self):
        """Perform graceful shutdown with optional backup"""
        self._shutdown_initiated = True
        logger.info("Initiating graceful shutdown...")
        
        # Wait for active requests to complete (with timeout)
        wait_time = 0
        while self._active_requests > 0 and wait_time < self._shutdown_timeout:
            logger.info(f"Waiting for {self._active_requests} active requests to complete...")
            await asyncio.sleep(1)
            wait_time += 1
        
        if self._active_requests > 0:
            logger.warning(f"Shutdown timeout reached. {self._active_requests} requests may be interrupted.")
        
        # Create backup if configured
        if self._backup_on_shutdown:
            backup_file = await self.backup_data()
            if backup_file:
                logger.info(f"Pre-shutdown backup completed: {backup_file}")
        
        # Force Redis to save data to disk
        try:
            if redis_manager.redis:
                # Try BGSAVE first (non-blocking)
                try:
                    await redis_manager.redis.bgsave()
                    logger.info("Background save initiated")
                    
                    # Wait a bit for background save to start
                    await asyncio.sleep(0.5)
                    
                except Exception as e:
                    logger.warning(f"Background save failed, trying synchronous save: {e}")
                    try:
                        # Fallback to synchronous save
                        await redis_manager.redis.save()
                        logger.info("Synchronous save completed")
                    except Exception as e2:
                        logger.warning(f"Synchronous save also failed (Redis persistence may be disabled): {e2}")
        
        except Exception as e:
            logger.error(f"Error during Redis save: {e}")
        
        # Close Redis connection
        try:
            await redis_manager.disconnect()
            logger.info("Redis connection closed successfully")
        except Exception as e:
            logger.error(f"Error closing Redis connection: {e}")
        
        logger.info("Graceful shutdown completed")
    
    def increment_requests(self):
        """Track active requests"""
        self._active_requests += 1
    
    def decrement_requests(self):
        """Track completed requests"""
        self._active_requests = max(0, self._active_requests - 1)
    
    @property
    def is_shutting_down(self) -> bool:
        """Check if shutdown is in progress"""
        return self._shutdown_initiated
    
    def get_database(self) -> RedisDatabase:
        """Get database instance"""
        if self.db is None:
            raise RuntimeError("Database not initialized")
        if self._shutdown_initiated:
            raise RuntimeError("Application is shutting down")
        return self.db