"""
Docker-based secure code execution sandbox.
"""

try:
    import docker
    DOCKER_AVAILABLE = True
except ImportError:
    DOCKER_AVAILABLE = False
    docker = None

import json
import time
from typing import Dict, Any, Optional
from dataclasses import dataclass


@dataclass
class ExecutionResult:
    """Result of code execution in sandbox."""
    success: bool
    output: str
    error: Optional[str]
    execution_time: float
    memory_used: int  # bytes
    exit_code: int


class DockerExecutor:
    """
    Execute Python code in isolated Docker containers.
    
    Features:
    - Ephemeral containers (auto-cleanup)
    - Resource limits (CPU, memory, timeout)
    - No network access by default
    - No file system access by default
    """
    
    def __init__(
        self,
        image: str = "python:3.11-slim",
        cpu_limit: float = 1.0,  # CPU cores
        memory_limit: str = "256m",  # Memory limit
        timeout: int = 30,  # Execution timeout in seconds
        network_disabled: bool = True,
        enable_pip: bool = False,  # Allow pip installs
    ):
        """
        Initialize Docker executor.
        
        Args:
            image: Docker image to use
            cpu_limit: CPU cores limit
            memory_limit: Memory limit (e.g., "256m", "1g")
            timeout: Execution timeout in seconds
            network_disabled: Disable network access
            enable_pip: Allow pip package installation
        """
        if not DOCKER_AVAILABLE:
            raise RuntimeError(
                "Docker package not installed. "
                "Install with: pip install ybagent[sandbox] or use RestrictedPython executor."
            )
        
        try:
            self.client = docker.from_env()
            # Test connection
            self.client.ping()
        except Exception as e:
            raise RuntimeError(
                f"Docker not available: {e}. "
                "Please install Docker or use RestrictedPython executor."
            )
        
        self.image = image
        self.cpu_limit = cpu_limit
        self.memory_limit = memory_limit
        self.timeout = timeout
        self.network_disabled = network_disabled
        self.enable_pip = enable_pip
        
        # Ensure image is available
        self._ensure_image()
    
    def _ensure_image(self):
        """Pull Docker image if not available."""
        try:
            self.client.images.get(self.image)
        except docker.errors.ImageNotFound:
            print(f"Pulling Docker image: {self.image}...")
            self.client.images.pull(self.image)
    
    def execute(
        self,
        code: str,
        context: Optional[Dict[str, Any]] = None
    ) -> ExecutionResult:
        """
        Execute Python code in Docker sandbox.
        
        Args:
            code: Python code to execute
            context: Optional context variables (serialized as JSON)
        
        Returns:
            ExecutionResult with output, errors, and metrics
        """
        start_time = time.time()
        
        # Prepare execution script
        script = self._prepare_script(code, context)
        
        # Container configuration
        container_config = {
            "image": self.image,
            "command": ["python", "-c", script],
            "detach": True,
            "mem_limit": self.memory_limit,
            "nano_cpus": int(self.cpu_limit * 1e9),
            "network_disabled": self.network_disabled,
            "remove": True,  # Auto-remove after execution
        }
        
        container = None
        try:
            # Create and start container
            container = self.client.containers.run(**container_config)
            
            # Wait for completion with timeout
            result = container.wait(timeout=self.timeout)
            exit_code = result.get("StatusCode", -1)
            
            # Get output
            logs = container.logs(stdout=True, stderr=True).decode("utf-8")
            
            # Get memory stats
            stats = container.stats(stream=False)
            memory_used = stats.get("memory_stats", {}).get("usage", 0)
            
            execution_time = time.time() - start_time
            
            # Parse output
            if exit_code == 0:
                return ExecutionResult(
                    success=True,
                    output=logs,
                    error=None,
                    execution_time=execution_time,
                    memory_used=memory_used,
                    exit_code=exit_code
                )
            else:
                return ExecutionResult(
                    success=False,
                    output="",
                    error=logs,
                    execution_time=execution_time,
                    memory_used=memory_used,
                    exit_code=exit_code
                )
        
        except docker.errors.ContainerError as e:
            return ExecutionResult(
                success=False,
                output="",
                error=f"Container error: {e}",
                execution_time=time.time() - start_time,
                memory_used=0,
                exit_code=-1
            )
        
        except Exception as e:
            return ExecutionResult(
                success=False,
                output="",
                error=f"Execution error: {e}",
                execution_time=time.time() - start_time,
                memory_used=0,
                exit_code=-1
            )
        
        finally:
            # Cleanup (container auto-removes, but force if needed)
            if container:
                try:
                    container.stop(timeout=1)
                except:
                    pass
    
    def _prepare_script(
        self,
        code: str,
        context: Optional[Dict[str, Any]] = None
    ) -> str:
        """
        Prepare execution script with context injection.
        
        Args:
            code: User code
            context: Context variables
        
        Returns:
            Complete Python script
        """
        script_parts = []
        
        # Inject context if provided
        if context:
            script_parts.append("import json")
            script_parts.append(f"__context__ = {json.dumps(context)}")
            script_parts.append("globals().update(__context__)")
        
        # Add user code
        script_parts.append(code)
        
        return "\n".join(script_parts)
    
    def install_packages(self, packages: list) -> bool:
        """
        Pre-install packages in a custom image.
        
        Args:
            packages: List of pip packages
        
        Returns:
            True if successful
        """
        if not self.enable_pip:
            raise RuntimeError("Package installation disabled")
        
        # Create Dockerfile
        dockerfile = f"""
FROM {self.image}
RUN pip install --no-cache-dir {' '.join(packages)}
"""
        
        # Build custom image
        tag = f"{self.image}-custom"
        try:
            self.client.images.build(
                fileobj=dockerfile.encode(),
                tag=tag,
                rm=True
            )
            self.image = tag
            return True
        except Exception as e:
            print(f"Failed to install packages: {e}")
            return False
