"""
Docker sandbox for testing automations with production-like isolation.

This module provides a Docker-based sandbox environment that runs automations
in isolated containers, simulating the production runtime environment.
"""

from __future__ import annotations

import asyncio
import json
import os
import tempfile
import time
from dataclasses import dataclass, field
from pathlib import Path
from typing import Any

import docker
from docker.errors import ContainerError, DockerException, ImageNotFound


@dataclass
class DockerSandboxConfig:
    """Configuration for Docker sandbox execution.

    Attributes:
        memory_limit: Memory limit for the container (e.g., "512m", "1g").
        cpu_count: Number of CPUs to allocate to the container.
        timeout: Maximum execution time in seconds.
        network_mode: Docker network mode ("bridge", "none", "host").
        mock_credentials: Whether to use mock credentials.
        image: Docker image to use for the sandbox.
        auto_pull: Whether to automatically pull the image if not present.
        environment: Additional environment variables.
        remove_container: Whether to remove container after execution.
    """

    memory_limit: str = "512m"
    cpu_count: float = 1.0
    timeout: int = 300
    network_mode: str = "bridge"
    mock_credentials: bool = True
    image: str = "torivers/automation-sandbox:latest"
    auto_pull: bool = True
    environment: dict[str, str] = field(default_factory=dict)
    remove_container: bool = True


@dataclass
class SandboxResult:
    """Result from Docker sandbox execution.

    Attributes:
        success: Whether the execution completed successfully.
        output: Output data from the automation.
        logs: Container logs from execution.
        exit_code: Container exit code.
        duration_seconds: Total execution duration.
        memory_peak_mb: Peak memory usage in megabytes.
        tokens_used: Number of tokens used (if applicable).
        container_id: Docker container ID.
        error: Error message if execution failed.
    """

    success: bool
    output: dict[str, Any]
    logs: str
    exit_code: int
    duration_seconds: float = 0.0
    memory_peak_mb: float = 0.0
    tokens_used: int = 0
    container_id: str = ""
    error: str = ""

    def get_output(self, key: str, default: Any = None) -> Any:
        """Get a value from the output data.

        Args:
            key: Key to retrieve from output.
            default: Default value if key not found.

        Returns:
            Value from output or default.
        """
        return self.output.get(key, default)


class DockerSandboxError(Exception):
    """Exception raised for Docker sandbox errors."""

    pass


class DockerSandbox:
    """
    Docker sandbox for testing automations with production-like isolation.

    Uses Docker to provide isolated execution similar to production:
    - Resource limits (memory, CPU)
    - Network restrictions
    - Credential mocking
    - Timeout enforcement

    Example:
        from torivers_sdk.testing import DockerSandbox, DockerSandboxConfig

        # Create sandbox with custom config
        config = DockerSandboxConfig(
            memory_limit="1g",
            timeout=60,
            mock_credentials=True
        )
        sandbox = DockerSandbox(config)

        # Run automation
        result = await sandbox.run(
            automation_path="./my-automation",
            input_data={"query": "test data"}
        )

        # Check results
        if result.success:
            print(f"Output: {result.output}")
        else:
            print(f"Error: {result.error}")
    """

    # Output file path inside container
    OUTPUT_FILE = "/tmp/automation_output.json"

    def __init__(self, config: DockerSandboxConfig | None = None) -> None:
        """
        Initialize the Docker sandbox.

        Args:
            config: Optional sandbox configuration.

        Raises:
            DockerSandboxError: If Docker is not available.
        """
        self._config = config or DockerSandboxConfig()
        self._client: docker.DockerClient | None = None

    def _get_client(self) -> docker.DockerClient:
        """Get or create Docker client.

        Returns:
            Docker client instance.

        Raises:
            DockerSandboxError: If Docker is not available.
        """
        if self._client is None:
            try:
                self._client = docker.from_env()
                # Test connection
                self._client.ping()
            except DockerException as e:
                raise DockerSandboxError(
                    f"Failed to connect to Docker daemon: {e}. "
                    "Ensure Docker is installed and running."
                ) from e
        return self._client

    def _ensure_image(self) -> None:
        """Ensure the sandbox image is available.

        Raises:
            DockerSandboxError: If image is not available and cannot be pulled.
        """
        client = self._get_client()
        try:
            client.images.get(self._config.image)
        except ImageNotFound:
            if self._config.auto_pull:
                try:
                    client.images.pull(self._config.image)
                except DockerException as e:
                    raise DockerSandboxError(
                        f"Failed to pull image {self._config.image}: {e}"
                    ) from e
            else:
                raise DockerSandboxError(
                    f"Image {self._config.image} not found. "
                    "Set auto_pull=True or pull the image manually."
                )

    async def run(
        self,
        automation_path: str,
        input_data: dict[str, Any],
        config: DockerSandboxConfig | None = None,
    ) -> SandboxResult:
        """
        Run automation in Docker sandbox.

        Args:
            automation_path: Path to the automation directory.
            input_data: Input data for the automation.
            config: Optional override configuration for this run.

        Returns:
            SandboxResult with output, logs, and metrics.

        Raises:
            DockerSandboxError: If execution fails due to Docker issues.
        """
        cfg = config or self._config
        start_time = time.time()

        # Resolve absolute path
        abs_path = os.path.abspath(automation_path)
        if not os.path.isdir(abs_path):
            return SandboxResult(
                success=False,
                output={},
                logs="",
                exit_code=-1,
                duration_seconds=time.time() - start_time,
                error=f"Automation path does not exist: {abs_path}",
            )

        # Run in thread pool to avoid blocking
        loop = asyncio.get_event_loop()
        return await loop.run_in_executor(
            None,
            self._run_container,
            abs_path,
            input_data,
            cfg,
            start_time,
        )

    def _run_container(
        self,
        automation_path: str,
        input_data: dict[str, Any],
        config: DockerSandboxConfig,
        start_time: float,
    ) -> SandboxResult:
        """
        Run container synchronously.

        Args:
            automation_path: Absolute path to automation directory.
            input_data: Input data for automation.
            config: Sandbox configuration.
            start_time: Execution start time.

        Returns:
            SandboxResult with execution results.
        """
        container = None
        temp_output_dir = None

        try:
            client = self._get_client()
            self._ensure_image()

            # Create temp directory for output
            temp_output_dir = tempfile.mkdtemp(prefix="torivers_sandbox_")

            # Prepare environment variables
            environment = {
                "INPUT_DATA": json.dumps(input_data),
                "MOCK_CREDENTIALS": str(config.mock_credentials).lower(),
                "OUTPUT_FILE": self.OUTPUT_FILE,
                **config.environment,
            }

            # Prepare volumes
            volumes = {
                automation_path: {"bind": "/automation", "mode": "ro"},
                temp_output_dir: {"bind": "/tmp", "mode": "rw"},
            }

            # Run container
            container = client.containers.run(
                image=config.image,
                environment=environment,
                volumes=volumes,
                mem_limit=config.memory_limit,
                cpu_count=config.cpu_count,
                network_mode=config.network_mode,
                detach=True,
                working_dir="/automation",
            )

            # Wait for completion with timeout
            try:
                result = container.wait(timeout=config.timeout)
                exit_code = result.get("StatusCode", -1)
            except Exception as e:
                # Timeout or other error
                container.kill()
                return SandboxResult(
                    success=False,
                    output={},
                    logs=self._get_logs(container),
                    exit_code=-1,
                    duration_seconds=time.time() - start_time,
                    container_id=container.id,
                    error=f"Container execution failed: {e}",
                )

            # Get logs
            logs = self._get_logs(container)

            # Get stats for memory usage
            memory_peak_mb = self._get_memory_peak(container)

            # Read output file
            output = self._read_output(temp_output_dir)

            duration = time.time() - start_time

            return SandboxResult(
                success=exit_code == 0,
                output=output.get("output_data", {}),
                logs=logs,
                exit_code=exit_code,
                duration_seconds=duration,
                memory_peak_mb=memory_peak_mb,
                tokens_used=output.get("tokens_used", 0),
                container_id=container.id,
                error=output.get("error", "") if exit_code != 0 else "",
            )

        except DockerSandboxError:
            raise
        except ContainerError as e:
            return SandboxResult(
                success=False,
                output={},
                logs=e.stderr.decode() if e.stderr else "",
                exit_code=e.exit_status,
                duration_seconds=time.time() - start_time,
                error=str(e),
            )
        except Exception as e:
            return SandboxResult(
                success=False,
                output={},
                logs="",
                exit_code=-1,
                duration_seconds=time.time() - start_time,
                error=f"Unexpected error: {e}",
            )
        finally:
            # Cleanup
            if container and config.remove_container:
                try:
                    container.remove(force=True)
                except Exception:
                    pass
            if temp_output_dir:
                try:
                    import shutil

                    shutil.rmtree(temp_output_dir, ignore_errors=True)
                except Exception:
                    pass

    def _get_logs(self, container: Any) -> str:
        """Get logs from container.

        Args:
            container: Docker container object.

        Returns:
            Container logs as string.
        """
        try:
            return container.logs(stdout=True, stderr=True).decode("utf-8")
        except Exception:
            return ""

    def _get_memory_peak(self, container: Any) -> float:
        """Get peak memory usage from container.

        Args:
            container: Docker container object.

        Returns:
            Peak memory usage in megabytes.
        """
        try:
            stats = container.stats(stream=False)
            memory_stats = stats.get("memory_stats", {})
            max_usage = memory_stats.get("max_usage", 0)
            return max_usage / (1024 * 1024)  # Convert to MB
        except Exception:
            return 0.0

    def _read_output(self, output_dir: str) -> dict[str, Any]:
        """Read output from the output directory.

        Args:
            output_dir: Path to output directory.

        Returns:
            Output dictionary.
        """
        output_file = Path(output_dir) / "automation_output.json"
        if output_file.exists():
            try:
                return json.loads(output_file.read_text())
            except (json.JSONDecodeError, IOError):
                pass
        return {}

    def is_available(self) -> bool:
        """Check if Docker sandbox is available.

        Returns:
            True if Docker is available and running.
        """
        try:
            self._get_client()
            return True
        except DockerSandboxError:
            return False

    def build_image(
        self,
        dockerfile_path: str | None = None,
        tag: str | None = None,
    ) -> bool:
        """Build the sandbox Docker image.

        Args:
            dockerfile_path: Path to Dockerfile directory.
            tag: Tag for the image.

        Returns:
            True if build succeeded.

        Raises:
            DockerSandboxError: If build fails.
        """
        client = self._get_client()
        tag = tag or self._config.image

        # Use default Dockerfile if not specified
        if dockerfile_path is None:
            dockerfile_path = str(Path(__file__).parent / "docker" / "sandbox")

        if not os.path.isdir(dockerfile_path):
            raise DockerSandboxError(
                f"Dockerfile directory not found: {dockerfile_path}"
            )

        try:
            client.images.build(
                path=dockerfile_path,
                tag=tag,
                rm=True,
            )
            return True
        except DockerException as e:
            raise DockerSandboxError(f"Failed to build image: {e}") from e
