"""
Docker Executor for Mai Safe Code Execution

Provides isolated container execution using Docker with comprehensive
resource limits, security restrictions, and audit logging integration.
"""

import logging
import tempfile
import time
import uuid
from dataclasses import dataclass
from pathlib import Path
from typing import Any

try:
    import docker
    from docker.errors import APIError, ContainerError, DockerException, ImageNotFound
    from docker.models.containers import Container

    DOCKER_AVAILABLE = True
except ImportError:
    docker = None
    Container = None
    DockerException = Exception
    APIError = Exception
    ContainerError = Exception
    ImageNotFound = Exception
    DOCKER_AVAILABLE = False

from .audit_logger import AuditLogger


@dataclass
class ContainerConfig:
    """Configuration for Docker container execution"""

    image: str = "python:3.10-slim"
    timeout_seconds: int = 30
    memory_limit: str = "128m"  # Docker memory limit format
    cpu_limit: str = "0.5"  # CPU quota (0.5 = 50% of one CPU)
    network_disabled: bool = True
    read_only_filesystem: bool = True
    tmpfs_size: str = "64m"  # Temporary filesystem size
    working_dir: str = "/app"
    user: str = "nobody"  # Non-root user


@dataclass
class ContainerResult:
    """Result of container execution"""

    success: bool
    container_id: str
    exit_code: int
    stdout: str | None = None
    stderr: str | None = None
    execution_time: float = 0.0
    error: str | None = None
    resource_usage: dict[str, Any] | None = None


class DockerExecutor:
    """
    Docker-based container executor for isolated code execution.

    Provides secure sandboxing using Docker containers with resource limits,
    network restrictions, and comprehensive audit logging.
    """

    def __init__(self, audit_logger: AuditLogger | None = None):
        """
        Initialize Docker executor

        Args:
            audit_logger: Optional audit logger for execution logging
        """
        self.audit_logger = audit_logger
        self.client = None
        self.available = False

        # Try to initialize Docker client
        self._initialize_docker()

        # Setup logging
        self.logger = logging.getLogger(__name__)

    def _initialize_docker(self) -> None:
        """Initialize Docker client and verify availability"""
        if not DOCKER_AVAILABLE:
            self.available = False
            return

        try:
            if docker is not None:
                self.client = docker.from_env()
                # Test Docker connection
                self.client.ping()
                self.available = True
            else:
                self.available = False
                self.client = None
        except Exception as e:
            self.logger.warning(f"Docker not available: {e}")
            self.available = False
            self.client = None

    def is_available(self) -> bool:
        """Check if Docker executor is available"""
        return self.available and self.client is not None

    def execute_code(
        self,
        code: str,
        config: ContainerConfig | None = None,
        environment: dict[str, str] | None = None,
        files: dict[str, str] | None = None,
    ) -> ContainerResult:
        """
        Execute code in isolated Docker container

        Args:
            code: Python code to execute
            config: Container configuration
            environment: Environment variables
            files: Additional files to mount in container

        Returns:
            ContainerResult with execution details
        """
        if not self.is_available() or self.client is None:
            return ContainerResult(
                success=False, container_id="", exit_code=-1, error="Docker executor not available"
            )

        config = config or ContainerConfig()
        container_id = str(uuid.uuid4())[:8]
        start_time = time.time()

        try:
            # Create temporary directory for files
            with tempfile.TemporaryDirectory() as temp_dir:
                temp_path = Path(temp_dir)

                # Write code to file
                code_file = temp_path / "execute.py"
                code_file.write_text(code)

                # Prepare volume mounts
                volumes = {
                    str(code_file): {
                        "bind": f"{config.working_dir}/execute.py",
                        "mode": "ro",  # read-only
                    }
                }

                # Add additional files if provided
                if files:
                    for filename, content in files.items():
                        file_path = temp_path / filename
                        file_path.write_text(content)
                        volumes[str(file_path)] = {
                            "bind": f"{config.working_dir}/{filename}",
                            "mode": "ro",
                        }

                # Prepare container configuration
                container_config = self._build_container_config(config, environment)

                # Create and start container
                container = self.client.containers.run(
                    image=config.image,
                    command=["python", "execute.py"],
                    volumes=volumes,
                    **container_config,
                    detach=True,
                )

                # Get container ID safely
                container_id = getattr(container, "id", container_id)

                try:
                    # Wait for completion with timeout
                    result = container.wait(timeout=config.timeout_seconds)
                    exit_code = result["StatusCode"]

                    # Get logs
                    stdout = container.logs(stdout=True, stderr=False).decode("utf-8")
                    stderr = container.logs(stdout=False, stderr=True).decode("utf-8")

                    # Get resource usage stats
                    stats = self._get_container_stats(container)

                    # Determine success
                    success = exit_code == 0 and not stderr

                    execution_result = ContainerResult(
                        success=success,
                        container_id=container_id,
                        exit_code=exit_code,
                        stdout=stdout,
                        stderr=stderr,
                        execution_time=time.time() - start_time,
                        resource_usage=stats,
                    )

                    # Log execution if audit logger available
                    if self.audit_logger:
                        self._log_container_execution(code, execution_result, config)

                    return execution_result

                finally:
                    # Always cleanup container
                    try:
                        container.remove(force=True)
                    except Exception:
                        pass  # Best effort cleanup

        except ContainerError as e:
            return ContainerResult(
                success=False,
                container_id=container_id or "unknown",
                exit_code=getattr(e, "exit_code", -1),
                stderr=str(e),
                execution_time=time.time() - start_time,
                error=f"Container execution error: {e}",
            )

        except ImageNotFound as e:
            return ContainerResult(
                success=False,
                container_id=container_id,
                exit_code=-1,
                error=f"Docker image not found: {e}",
            )

        except APIError as e:
            return ContainerResult(
                success=False,
                container_id=container_id,
                exit_code=-1,
                error=f"Docker API error: {e}",
            )

        except Exception as e:
            return ContainerResult(
                success=False,
                container_id=container_id,
                exit_code=-1,
                execution_time=time.time() - start_time,
                error=f"Unexpected error: {e}",
            )

    def _build_container_config(
        self, config: ContainerConfig, environment: dict[str, str] | None = None
    ) -> dict[str, Any]:
        """Build Docker container configuration"""
        container_config = {
            "mem_limit": config.memory_limit,
            "cpu_quota": int(float(config.cpu_limit) * 100000),  # Convert to microseconds
            "cpu_period": 100000,  # 100ms period
            "network_disabled": config.network_disabled,
            "read_only": config.read_only_filesystem,
            "tmpfs": {"/tmp": f"size={config.tmpfs_size},noexec,nosuid,nodev"},
            "user": config.user,
            "working_dir": config.working_dir,
            "remove": True,  # Auto-remove container
        }

        # Add environment variables
        if environment:
            container_config["environment"] = {
                **environment,
                "PYTHONPATH": config.working_dir,
                "PYTHONDONTWRITEBYTECODE": "1",
            }
        else:
            container_config["environment"] = {
                "PYTHONPATH": config.working_dir,
                "PYTHONDONTWRITEBYTECODE": "1",
            }

        # Security options
        container_config["security_opt"] = [
            "no-new-privileges:true",
            "seccomp:unconfined",  # Python needs some syscalls
        ]

        # Capabilities (drop all capabilities)
        container_config["cap_drop"] = ["ALL"]
        container_config["cap_add"] = ["CHOWN", "DAC_OVERRIDE"]  # Minimal capabilities for Python

        return container_config

    def _get_container_stats(self, container) -> dict[str, Any]:
        """Get resource usage statistics from container"""
        try:
            stats = container.stats(stream=False)

            # Calculate CPU usage
            cpu_stats = stats.get("cpu_stats", {})
            precpu_stats = stats.get("precpu_stats", {})

            cpu_usage = cpu_stats.get("cpu_usage", {}).get("total_usage", 0)
            precpu_usage = precpu_stats.get("cpu_usage", {}).get("total_usage", 0)

            system_usage = cpu_stats.get("system_cpu_usage", 0)
            presystem_usage = precpu_stats.get("system_cpu_usage", 0)

            cpu_count = cpu_stats.get("online_cpus", 1)

            cpu_percent = 0.0
            if system_usage > presystem_usage:
                cpu_delta = cpu_usage - precpu_usage
                system_delta = system_usage - presystem_usage
                cpu_percent = (cpu_delta / system_delta) * cpu_count * 100.0

            # Calculate memory usage
            memory_stats = stats.get("memory_stats", {})
            memory_usage = memory_stats.get("usage", 0)
            memory_limit = memory_stats.get("limit", 1)
            memory_percent = (memory_usage / memory_limit) * 100.0

            return {
                "cpu_percent": round(cpu_percent, 2),
                "memory_usage_bytes": memory_usage,
                "memory_limit_bytes": memory_limit,
                "memory_percent": round(memory_percent, 2),
                "memory_usage_mb": round(memory_usage / (1024 * 1024), 2),
            }

        except Exception:
            return {
                "cpu_percent": 0.0,
                "memory_usage_bytes": 0,
                "memory_limit_bytes": 0,
                "memory_percent": 0.0,
                "memory_usage_mb": 0.0,
            }

    def _log_container_execution(
        self, code: str, result: ContainerResult, config: ContainerConfig
    ) -> None:
        """Log container execution to audit logger"""
        if not self.audit_logger:
            return

        execution_data = {
            "type": "docker_container",
            "container_id": result.container_id,
            "exit_code": result.exit_code,
            "stdout": result.stdout,
            "stderr": result.stderr,
            "execution_time": result.execution_time,
            "config": {
                "image": config.image,
                "timeout": config.timeout_seconds,
                "memory_limit": config.memory_limit,
                "cpu_limit": config.cpu_limit,
                "network_disabled": config.network_disabled,
                "read_only_filesystem": config.read_only_filesystem,
            },
            "resource_usage": result.resource_usage,
        }

        # Note: execution_type parameter not available in current AuditLogger
        self.audit_logger.log_execution(code=code, execution_result=execution_data)

    def get_available_images(self) -> list[str]:
        """Get list of available Docker images"""
        if not self.is_available() or self.client is None:
            return []

        try:
            images = self.client.images.list()
            return [img.tags[0] for img in images if img.tags]
        except Exception:
            return []

    def pull_image(self, image_name: str) -> bool:
        """Pull Docker image"""
        if not self.is_available() or self.client is None:
            return False

        try:
            self.client.images.pull(image_name)
            return True
        except Exception:
            return False

    def cleanup_containers(self) -> int:
        """Clean up any dangling containers"""
        if not self.is_available() or self.client is None:
            return 0

        try:
            containers = self.client.containers.list(all=True, filters={"status": "exited"})
            count = 0
            for container in containers:
                try:
                    container.remove(force=True)
                    count += 1
                except Exception:
                    pass
            return count
        except Exception:
            return 0

    def get_system_info(self) -> dict[str, Any]:
        """Get Docker system information"""
        if not self.is_available() or self.client is None:
            return {"available": False}

        try:
            info = self.client.info()
            version = self.client.version()

            return {
                "available": True,
                "version": version.get("Version", "unknown"),
                "api_version": version.get("ApiVersion", "unknown"),
                "containers": info.get("Containers", 0),
                "containers_running": info.get("ContainersRunning", 0),
                "containers_paused": info.get("ContainersPaused", 0),
                "containers_stopped": info.get("ContainersStopped", 0),
                "images": info.get("Images", 0),
                "memory_total": info.get("MemTotal", 0),
                "ncpu": info.get("NCPU", 0),
            }
        except Exception:
            return {"available": False, "error": "Failed to get system info"}