"""
Enhanced Docker Sandbox Performance System
Provides pre-warmed container pool, faster startup, and better resource management
"""
import asyncio
import docker
import logging
from datetime import datetime
from typing import Dict, Optional, List
from dataclasses import dataclass
from enum import Enum
from pathlib import Path
import json
import psutil

logger = logging.getLogger(__name__)


class ContainerStatus(Enum):
    """Container status in pool"""
    IDLE = "idle"           # Available for use
    BUSY = "busy"           # Currently in use
    WARMING = "warming"     # Starting up
    STOPPED = "stopped"     # Stopped
    ERROR = "error"         # Error state


@dataclass
class ContainerInfo:
    """Information about a container in the pool"""
    container_id: str
    status: ContainerStatus
    created_at: str
    last_used: str
    usage_count: int = 0
    total_cpu_time: float = 0.0
    peak_memory_mb: float = 0.0
    image_name: str = "python:3.13-slim"


class ContainerPool:
    """
    Pool of pre-warmed Docker containers.
    Reduces startup time by reusing containers.
    """
    
    def __init__(
        self,
        pool_size: int = 3,
        image_name: str = "python:3.13-slim",
        memory_limit: str = "2g",
        cpu_limit: float = 2.0,
        network_disabled: bool = True
    ):
        """
        Args:
            pool_size: Number of containers to keep warm
            image_name: Docker image to use
            memory_limit: Memory limit per container
            cpu_limit: CPU limit per container
            network_disabled: Disable network (security)
        """
        self.pool_size = pool_size
        self.image_name = image_name
        self.memory_limit = memory_limit
        self.cpu_limit = cpu_limit
        self.network_disabled = network_disabled
        
        self.containers: Dict[str, ContainerInfo] = {}
        self.container_queue: asyncio.Queue = asyncio.Queue()
        
        # Docker client
        self.docker_client: Optional[docker.DockerClient] = None
        
        # Pool state
        self._initialized = False
        self._maintenance_task: Optional[asyncio.Task] = None
        
        logger.info(
            "[Container Pool] Initialized (size=%d, image=%s)",
            pool_size, image_name
        )
    
    async def initialize(self) -> None:
        """Initialize the container pool"""
        if self._initialized:
            return
        
        try:
            # Connect to Docker
            self.docker_client = docker.from_env()
            
            # Test connection
            self.docker_client.ping()
            
            logger.info("[Container Pool] Connected to Docker")
            
            # Pre-warm containers
            await self._prewarm_containers()
            
            # Start maintenance task
            self._maintenance_task = asyncio.create_task(self._maintenance_loop())
            
            self._initialized = True
            
            logger.info("[Container Pool] ✓ Pool ready with %d containers", self.pool_size)
            
        except Exception as e:
            logger.error("[Container Pool] Initialization failed: %s", e)
            raise
    
    async def _prewarm_containers(self) -> None:
        """Pre-warm containers in the pool"""
        logger.info("[Container Pool] Pre-warming %d containers...", self.pool_size)
        
        for i in range(self.pool_size):
            try:
                container = await self._create_container(f"aiiq-pool-{i}")
                
                # Add to pool
                container_info = ContainerInfo(
                    container_id=container.id,
                    status=ContainerStatus.IDLE,
                    created_at=datetime.now().isoformat(),
                    last_used=datetime.now().isoformat(),
                    image_name=self.image_name
                )
                
                self.containers[container.id] = container_info
                
                # Add to queue
                await self.container_queue.put(container.id)
                
                logger.info(
                    "[Container Pool] Pre-warmed container %s: %s",
                    i, container.id[:12]
                )
                
            except Exception as e:
                logger.error(
                    "[Container Pool] Failed to pre-warm container %d: %s",
                    i, e
                )
    
    async def _create_container(
        self,
        name: str,
        command: Optional[List[str]] = None
    ) -> docker.models.containers.Container:
        """Create a new Docker container"""
        container = self.docker_client.containers.run(
            image=self.image_name,
            name=name,
            command=command or ["sleep", "infinity"],
            detach=True,
            mem_limit=self.memory_limit,
            cpu_quota=int(self.cpu_limit * 100000),
            cpu_period=100000,
            network_disabled=self.network_disabled,
            security_opt=["no-new-privileges"],
            read_only=False,  # Need write for some operations
            auto_remove=False
        )
        
        return container
    
    async def acquire_container(
        self,
        timeout: float = 5.0
    ) -> Optional[str]:
        """
        Acquire a container from the pool.
        
        Args:
            timeout: Max time to wait for container
            
        Returns:
            Container ID or None if timeout
        """
        if not self._initialized:
            await self.initialize()
        
        try:
            # Wait for available container
            container_id = await asyncio.wait_for(
                self.container_queue.get(),
                timeout=timeout
            )
            
            # Update status
            if container_id in self.containers:
                container_info = self.containers[container_id]
                container_info.status = ContainerStatus.BUSY
                container_info.last_used = datetime.now().isoformat()
                container_info.usage_count += 1
            
            logger.debug(
                "[Container Pool] Acquired container: %s",
                container_id[:12]
            )
            
            return container_id
            
        except asyncio.TimeoutError:
            logger.warning("[Container Pool] Timeout waiting for container")
            return None
    
    async def release_container(self, container_id: str) -> None:
        """
        Release a container back to the pool.
        
        Args:
            container_id: Container to release
        """
        if container_id not in self.containers:
            logger.warning(
                "[Container Pool] Unknown container: %s",
                container_id[:12]
            )
            return
        
        container_info = self.containers[container_id]
        
        # Update stats
        container_info.status = ContainerStatus.IDLE
        container_info.last_used = datetime.now().isoformat()
        
        # Put back in queue
        await self.container_queue.put(container_id)
        
        logger.debug(
            "[Container Pool] Released container: %s",
            container_id[:12]
        )
    
    async def execute_in_container(
        self,
        container_id: str,
        command: List[str],
        workdir: Optional[str] = None,
        timeout: float = 30.0
    ) -> Dict[str, Any]:
        """
        Execute a command in a container.
        
        Args:
            container_id: Container ID
            command: Command to execute
            workdir: Working directory
            timeout: Execution timeout
            
        Returns:
            Execution result
        """
        try:
            container = self.docker_client.containers.get(container_id)
            
            # Execute command
            result = container.exec_run(
                cmd=command,
                workdir=workdir,
                demux=True,
                socket=False
            )
            
            exit_code, (stdout, stderr) = result
            
            return {
                "exit_code": exit_code,
                "stdout": stdout.decode() if stdout else "",
                "stderr": stderr.decode() if stderr else "",
                "success": exit_code == 0
            }
            
        except Exception as e:
            logger.error(
                "[Container Pool] Execution error: %s",
                e
            )
            return {
                "exit_code": -1,
                "stdout": "",
                "stderr": str(e),
                "success": False
            }
    
    async def _maintenance_loop(self) -> None:
        """Background maintenance task"""
        while True:
            try:
                await asyncio.sleep(60)  # Check every minute
                
                # Clean up unused containers
                await self._cleanup_idle_containers()
                
                # Restart old containers
                await self._refresh_old_containers()
                
                # Update resource stats
                await self._update_container_stats()
                
            except Exception as e:
                logger.error("[Container Pool] Maintenance error: %s", e)
    
    async def _cleanup_idle_containers(self) -> None:
        """Remove containers idle for too long"""
        idle_threshold_minutes = 30
        
        for container_id, info in list(self.containers.items()):
            if info.status == ContainerStatus.IDLE:
                last_used = datetime.fromisoformat(info.last_used)
                idle_minutes = (datetime.now() - last_used).total_seconds() / 60
                
                if idle_minutes > idle_threshold_minutes:
                    logger.info(
                        "[Container Pool] Removing idle container: %s (idle %d min)",
                        container_id[:12], int(idle_minutes)
                    )
                    
                    await self._remove_container(container_id)
    
    async def _refresh_old_containers(self) -> None:
        """Refresh containers that are too old"""
        max_age_hours = 4
        
        for container_id, info in list(self.containers.items()):
            created = datetime.fromisoformat(info.created_at)
            age_hours = (datetime.now() - created).total_seconds() / 3600
            
            if age_hours > max_age_hours:
                logger.info(
                    "[Container Pool] Refreshing old container: %s (age %d hr)",
                    container_id[:12], int(age_hours)
                )
                
                await self._remove_container(container_id)
                
                # Create replacement
                await self._create_replacement(container_id)
    
    async def _create_replacement(self, old_id: str) -> None:
        """Create replacement container"""
        # Extract name from old ID
        old_container = self.docker_client.containers.get(old_id)
        old_name = old_container.name
        
        # Remove old
        old_container.remove(force=True)
        del self.containers[old_id]
        
        # Create new
        new_container = await self._create_container(old_name)
        
        container_info = ContainerInfo(
            container_id=new_container.id,
            status=ContainerStatus.IDLE,
            created_at=datetime.now().isoformat(),
            last_used=datetime.now().isoformat(),
            image_name=self.image_name
        )
        
        self.containers[new_container.id] = container_info
        await self.container_queue.put(new_container.id)
    
    async def _update_container_stats(self) -> None:
        """Update resource usage statistics"""
        for container_id, info in self.containers.items():
            try:
                container = self.docker_client.containers.get(container_id)
                stats = container.stats(stream=False)
                
                if stats:
                    # Update CPU and memory stats
                    cpu_delta = stats.get("cpu_stats", {}).get("cpu_usage", 0)
                    memory_mb = stats.get("memory_stats", {}).get("usage", 0) / (1024 * 1024)
                    
                    info.total_cpu_time += cpu_delta / 1e9  # Convert to seconds
                    info.peak_memory_mb = max(info.peak_memory_mb, memory_mb)
                    
            except Exception as e:
                logger.warning(
                    "[Container Pool] Error getting stats for %s: %s",
                    container_id[:12], e
                )
    
    async def _remove_container(self, container_id: str) -> None:
        """Remove a container from the pool"""
        try:
            container = self.docker_client.containers.get(container_id)
            container.stop(timeout=5)
            container.remove()
            
            if container_id in self.containers:
                del self.containers[container_id]
            
            logger.debug("[Container Pool] Removed: %s", container_id[:12])
            
        except Exception as e:
            logger.error(
                "[Container Pool] Error removing %s: %s",
                container_id[:12], e
            )
    
    async def shutdown(self) -> None:
        """Shutdown the container pool"""
        logger.info("[Container Pool] Shutting down...")
        
        # Stop maintenance
        if self._maintenance_task:
            self._maintenance_task.cancel()
        
        # Remove all containers
        for container_id in list(self.containers.keys()):
            await self._remove_container(container_id)
        
        # Close Docker client
        if self.docker_client:
            self.docker_client.close()
        
        self._initialized = False
        
        logger.info("[Container Pool] Shutdown complete")
    
    def get_pool_stats(self) -> Dict[str, Any]:
        """Get pool statistics"""
        idle = sum(1 for c in self.containers.values() if c.status == ContainerStatus.IDLE)
        busy = sum(1 for c in self.containers.values() if c.status == ContainerStatus.BUSY)
        
        total_cpu = sum(c.total_cpu_time for c in self.containers.values())
        peak_mem = max((c.peak_memory_mb for c in self.containers.values()), default=0)
        
        return {
            "total_containers": len(self.containers),
            "idle_containers": idle,
            "busy_containers": busy,
            "target_pool_size": self.pool_size,
            "total_cpu_time_seconds": total_cpu,
            "peak_memory_mb": peak_mem,
            "container_list": [
                {
                    "id": c.container_id[:12],
                    "status": c.status.value,
                    "usage_count": c.usage_count,
                    "created_at": c.created_at
                }
                for c in self.containers.values()
            ]
        }


class ResourceMonitor:
    """Monitor system resource usage"""
    
    def __init__(self):
        self._history: List[Dict[str, float]] = []
        self._max_history = 100
    
    async def collect_stats(self) -> Dict[str, float]:
        """Collect current resource usage"""
        # CPU
        cpu_percent = psutil.cpu_percent(interval=0.1)
        
        # Memory
        memory = psutil.virtual_memory()
        
        # Disk
        disk = psutil.disk_usage('/')
        
        stats = {
            "cpu_percent": cpu_percent,
            "memory_available_gb": memory.available / (1024**3),
            "memory_used_gb": memory.used / (1024**3),
            "memory_percent": memory.percent,
            "disk_free_gb": disk.free / (1024**3),
            "disk_used_gb": disk.used / (1024**3),
            "disk_percent": disk.percent,
            "timestamp": datetime.now().timestamp()
        }
        
        # Add to history
        self._history.append(stats)
        if len(self._history) > self._max_history:
            self._history.pop(0)
        
        return stats
    
    def get_average_stats(self, samples: int = 10) -> Optional[Dict[str, float]]:
        """Get average resource usage over recent samples"""
        if not self._history:
            return None
        
        recent = self._history[-samples:]
        
        return {
            "avg_cpu_percent": sum(s["cpu_percent"] for s in recent) / len(recent),
            "avg_memory_percent": sum(s["memory_percent"] for s in recent) / len(recent),
            "avg_disk_percent": sum(s["disk_percent"] for s in recent) / len(recent)
        }


class EnhancedSandboxManager:
    """
    Enhanced sandbox manager with container pooling and resource monitoring.
    """
    
    def __init__(
        self,
        pool_size: int = 3,
        image_name: str = "python:3.13-slim"
    ):
        self.container_pool = ContainerPool(
            pool_size=pool_size,
            image_name=image_name
        )
        
        self.resource_monitor = ResourceMonitor()
        
        logger.info("[Enhanced Sandbox] Initialized")
    
    async def initialize(self) -> None:
        """Initialize the enhanced sandbox"""
        await self.container_pool.initialize()
        logger.info("[Enhanced Sandbox] ✓ Ready")
    
    async def execute_code(
        self,
        code: str,
        timeout: float = 30.0
    ) -> Dict[str, Any]:
        """
        Execute code in a sandboxed container.
        
        Args:
            code: Python code to execute
            timeout: Execution timeout
            
        Returns:
            Execution result
        """
        # Acquire container
        container_id = await self.container_pool.acquire_container()
        
        if not container_id:
            return {
                "success": False,
                "error": "No available container"
            }
        
        try:
            # Execute code
            result = await self.container_pool.execute_in_container(
                container_id=container_id,
                command=["python", "-c", code],
                timeout=timeout
            )
            
            return result
            
        finally:
            # Release container
            await self.container_pool.release_container(container_id)
    
    async def run_script(
        self,
        script_path: str,
        args: Optional[List[str]] = None,
        workdir: Optional[str] = None
    ) -> Dict[str, Any]:
        """
        Run a script in the sandbox.
        
        Args:
            script_path: Path to script
            args: Script arguments
            workdir: Working directory
            
        Returns:
            Execution result
        """
        container_id = await self.container_pool.acquire_container()
        
        if not container_id:
            return {
                "success": False,
                "error": "No available container"
            }
        
        try:
            command = ["python", script_path]
            if args:
                command.extend(args)
            
            result = await self.container_pool.execute_in_container(
                container_id=container_id,
                command=command,
                workdir=workdir
            )
            
            return result
            
        finally:
            await self.container_pool.release_container(container_id)
    
    def get_stats(self) -> Dict[str, Any]:
        """Get sandbox statistics"""
        return {
            "container_pool": self.container_pool.get_pool_stats(),
            "system_resources": self.resource_monitor.get_average_stats()
        }
    
    async def shutdown(self) -> None:
        """Shutdown the enhanced sandbox"""
        await self.container_pool.shutdown()


# Global instance
_enhanced_sandbox: Optional[EnhancedSandboxManager] = None


def get_enhanced_sandbox() -> EnhancedSandboxManager:
    """Get the global enhanced sandbox"""
    global _enhanced_sandbox
    if _enhanced_sandbox is None:
        _enhanced_sandbox = EnhancedSandboxManager()
    return _enhanced_sandbox
