"""Docker container runtime implementation."""

from __future__ import annotations

import contextlib
import io
import tarfile
from datetime import datetime
from pathlib import Path
from typing import TYPE_CHECKING

import docker
from docker.errors import APIError, ImageNotFound, NotFound

from gaol.models.container import ContainerState, ContainerStatus, NetworkMode
from gaol.profiles.applicator import MergedProfile, get_profile_applicator
from gaol.runtimes.base import ContainerRuntime

if TYPE_CHECKING:
    from gaol.checkpoint.models import CheckpointResult, RestoreResult
    from gaol.models.container import ContainerConfig, ContainerStats, Snapshot

# Label prefix for gaol-managed containers
LABEL_PREFIX = "dev.gaol"


class DockerRuntime(ContainerRuntime):
    """Docker container runtime using docker-py SDK."""

    def __init__(self) -> None:
        """Initialize the Docker runtime."""
        self._client: docker.DockerClient | None = None

    @property
    def name(self) -> str:
        """Return the name of this runtime."""
        return "docker"

    @property
    def client(self) -> docker.DockerClient:
        """Get or create the Docker client."""
        if self._client is None:
            self._client = docker.from_env()
        return self._client

    def is_available(self) -> bool:
        """Check if Docker is available on the current system."""
        try:
            self.client.ping()
            return True
        except Exception:
            return False

    def _get_image(self, config: ContainerConfig) -> str:
        """Get the Docker image to use for the container."""
        if config.image:
            return config.image

        from gaol.distro import get_docker_image

        try:
            return get_docker_image(config.distro)
        except KeyError:
            # Fall back to debian:{distro} for unknown distros
            return f"debian:{config.distro}"

    def _docker_status_to_status(self, docker_status: str) -> ContainerStatus:
        """Convert Docker status string to ContainerStatus."""
        status_map = {
            "created": ContainerStatus.CREATED,
            "running": ContainerStatus.RUNNING,
            "paused": ContainerStatus.PAUSED,
            "restarting": ContainerStatus.RUNNING,
            "removing": ContainerStatus.STOPPED,
            "exited": ContainerStatus.EXITED,
            "dead": ContainerStatus.DEAD,
        }
        return status_map.get(docker_status.lower(), ContainerStatus.UNKNOWN)

    def _parse_datetime(self, dt_str: str | None) -> datetime | None:
        """Parse Docker datetime string."""
        if not dt_str:
            return None
        # Docker uses ISO format with nanoseconds
        try:
            # Remove nanoseconds if present (Python datetime only supports microseconds)
            if "." in dt_str:
                parts = dt_str.split(".")
                # Keep only microseconds (6 digits)
                nano_part = parts[1][:6]
                # Handle timezone
                if "Z" in parts[1]:
                    dt_str = f"{parts[0]}.{nano_part}+00:00"
                elif "+" in parts[1] or "-" in parts[1][1:]:
                    # Find timezone offset
                    for sep in ["+", "-"]:
                        if sep in parts[1][1:]:
                            idx = parts[1].index(sep, 1)
                            tz = parts[1][idx:]
                            dt_str = f"{parts[0]}.{nano_part}{tz}"
                            break
                else:
                    dt_str = f"{parts[0]}.{nano_part}"
            return datetime.fromisoformat(dt_str.replace("Z", "+00:00"))
        except Exception:
            return None

    def _resolve_environment(self, env: dict[str, str]) -> dict[str, str]:
        """Resolve ${secret:NAME} references in environment variables.

        Args:
            env: Dictionary of environment variables

        Returns:
            New dictionary with secrets resolved
        """
        from gaol.secrets import has_secrets
        from gaol.secrets.resolver import get_secret_resolver

        # Check if any values have secret references
        has_refs = any(has_secrets(v) for v in env.values())
        if not has_refs:
            return env

        resolver = get_secret_resolver()
        return resolver.resolve_environment(env)

    def _container_to_state(self, container: docker.models.containers.Container) -> ContainerState:
        """Convert a Docker container to ContainerState."""
        container.reload()
        attrs = container.attrs

        # Get network info
        ip_address = None
        networks = attrs.get("NetworkSettings", {}).get("Networks", {})
        if networks:
            # Get IP from first network
            first_network = next(iter(networks.values()))
            ip_address = first_network.get("IPAddress")

        # Get port mappings
        ports: dict[int, int] = {}
        port_bindings = attrs.get("NetworkSettings", {}).get("Ports", {})
        for container_port, host_bindings in (port_bindings or {}).items():
            if host_bindings:
                # container_port is like "80/tcp"
                port_num = int(container_port.split("/")[0])
                host_port = int(host_bindings[0]["HostPort"])
                ports[host_port] = port_num

        # Get state info
        state = attrs.get("State", {})

        # Get profiles from labels
        labels = attrs.get("Config", {}).get("Labels", {})
        profiles_str = labels.get(f"{LABEL_PREFIX}.profiles", "")
        profiles = profiles_str.split(",") if profiles_str else []

        return ContainerState(
            id=container.id[:12] if container.id else "",
            name=container.name or "",
            runtime="docker",
            status=self._docker_status_to_status(state.get("Status", "unknown")),
            image=attrs.get("Config", {}).get("Image", ""),
            created_at=self._parse_datetime(attrs.get("Created")),
            started_at=self._parse_datetime(state.get("StartedAt")),
            finished_at=self._parse_datetime(state.get("FinishedAt")),
            exit_code=state.get("ExitCode"),
            error=state.get("Error") or None,
            ip_address=ip_address,
            ports=ports,
            profiles=profiles,
        )

    def create(self, config: ContainerConfig) -> str:
        """Create a new container."""
        # Apply profiles if specified
        merged_profile: MergedProfile | None = None
        if config.profiles:
            applicator = get_profile_applicator()
            merged_profile = applicator.resolve_and_merge(config.profiles)
            config = applicator.apply_to_config(config, merged_profile)

        image = self._get_image(config)

        # Pull image if not present
        try:
            self.client.images.get(image)
        except ImageNotFound:
            self.client.images.pull(image)

        # Build network mode
        network_mode: str | None = None
        if config.network.mode == NetworkMode.HOST:
            network_mode = "host"
        elif config.network.mode == NetworkMode.NONE:
            network_mode = "none"
        # SHARED and BRIDGED use default bridge network

        # Build port bindings
        ports: dict[str, int] = {}
        for host_port, container_port in config.network.ports.items():
            ports[f"{container_port}/tcp"] = host_port

        # Build volume bindings
        volumes: dict[str, dict[str, str]] = {}
        for host_path, container_path in config.volumes.items():
            volumes[host_path] = {"bind": container_path, "mode": "rw"}

        # Build labels
        labels = {
            f"{LABEL_PREFIX}.managed": "true",
            f"{LABEL_PREFIX}.runtime": "docker",
            f"{LABEL_PREFIX}.profiles": ",".join(config.profiles),
            f"{LABEL_PREFIX}.distro": config.distro,
            f"{LABEL_PREFIX}.setup_commands": "\n".join(config.setup_commands),
        }

        # Add capabilities from profiles
        cap_add: list[str] | None = None
        cap_drop: list[str] | None = None
        if merged_profile:
            if merged_profile.capabilities_add:
                cap_add = merged_profile.capabilities_add
            if merged_profile.capabilities_drop:
                cap_drop = merged_profile.capabilities_drop

        # Build resource limits
        mem_limit = f"{config.resources.memory_mb}m" if config.resources.memory_mb else None
        cpu_period = 100000 if config.resources.cpu_count else None
        cpu_quota = int(config.resources.cpu_count * 100000) if config.resources.cpu_count else None

        # Resolve secrets in environment variables
        resolved_env = self._resolve_environment(config.environment)

        # GPU configuration
        device_requests = None
        devices: list[str] = []

        if config.gpu and config.gpu.enabled:
            from gaol.gpu import GPUDetector, GPUVendor
            from gaol.gpu.nvidia import NvidiaGPU
            from gaol.gpu.amd import AmdGPU
            from gaol.gpu.intel import IntelGPU

            vendor = config.gpu.vendor
            if vendor == GPUVendor.AUTO:
                # Auto-detect vendor
                if GPUDetector.detect_nvidia():
                    vendor = GPUVendor.NVIDIA
                elif GPUDetector.detect_amd():
                    vendor = GPUVendor.AMD
                elif GPUDetector.detect_intel():
                    vendor = GPUVendor.INTEL

            if vendor == GPUVendor.NVIDIA:
                gpu_config = NvidiaGPU.get_docker_config(config.gpu)
                device_requests = gpu_config["device_requests"]
                resolved_env.update(gpu_config["environment"])
            elif vendor == GPUVendor.AMD:
                gpu_config = AmdGPU.get_docker_config(config.gpu)
                devices.extend(gpu_config["devices"])
                resolved_env.update(gpu_config["environment"])
            elif vendor == GPUVendor.INTEL:
                gpu_config = IntelGPU.get_docker_config(config.gpu)
                devices.extend(gpu_config["devices"])
                resolved_env.update(gpu_config["environment"])

        # Add custom device mappings
        for device in config.devices:
            container_path = device.container_path or device.host_path
            devices.append(f"{device.host_path}:{container_path}:{device.permissions}")

        try:
            container = self.client.containers.create(
                image=image,
                name=config.name,
                command=config.command,
                environment=resolved_env,
                ports=ports if ports else None,
                volumes=volumes if volumes else None,
                network_mode=network_mode,
                hostname=config.network.hostname,
                dns=config.network.dns if config.network.dns else None,
                working_dir=config.working_dir,
                privileged=config.privileged,
                read_only=config.read_only,
                labels=labels,
                mem_limit=mem_limit,
                cpu_shares=config.resources.cpu_shares,
                cpu_period=cpu_period,
                cpu_quota=cpu_quota,
                cap_add=cap_add,
                cap_drop=cap_drop,
                stdin_open=True,
                tty=True,
                device_requests=device_requests,
                devices=devices if devices else None,
            )
            return container.id[:12]
        except APIError as e:
            raise RuntimeError(f"Failed to create container: {e}") from e

    def setup(self, container_id: str) -> tuple[int, str, str]:
        """Run profile setup script in a container.

        This should be called after start() to install packages and run setup commands.

        Args:
            container_id: ID of the container

        Returns:
            Tuple of (exit_code, stdout, stderr)
        """
        try:
            container = self.client.containers.get(container_id)
            labels = container.labels

            # Get profiles and distro from labels
            profiles_str = labels.get(f"{LABEL_PREFIX}.profiles", "")
            distro = labels.get(f"{LABEL_PREFIX}.distro", "trixie")

            setup_commands_str = labels.get(f"{LABEL_PREFIX}.setup_commands", "")
            if not profiles_str and not setup_commands_str:
                return 0, "No profiles to set up\n", ""

            profiles = profiles_str.split(",") if profiles_str else []

            # Generate setup script
            if profiles:
                applicator = get_profile_applicator()
                merged = applicator.resolve_and_merge(profiles)
            else:
                merged = MergedProfile()

            # Merge spec-level setup commands (run before profile commands)
            if setup_commands_str:
                spec_commands = setup_commands_str.split("\n")
                merged.setup_commands = spec_commands + merged.setup_commands

            setup_script = merged.generate_setup_script(distro)

            # Run setup script in container
            exit_code, output = container.exec_run(
                cmd=["bash", "-c", setup_script],
                user="root",
                tty=True,
            )

            stdout = output.decode() if output else ""
            return exit_code, stdout, ""

        except NotFound as e:
            raise RuntimeError(f"Container not found: {container_id}") from e
        except APIError as e:
            raise RuntimeError(f"Failed to run setup: {e}") from e

    def start(self, container_id: str) -> None:
        """Start a stopped container."""
        try:
            container = self.client.containers.get(container_id)
            container.start()
        except NotFound as e:
            raise RuntimeError(f"Container not found: {container_id}") from e
        except APIError as e:
            raise RuntimeError(f"Failed to start container: {e}") from e

    def stop(self, container_id: str, timeout: int = 10) -> None:
        """Stop a running container."""
        try:
            container = self.client.containers.get(container_id)
            container.stop(timeout=timeout)
        except NotFound as e:
            raise RuntimeError(f"Container not found: {container_id}") from e
        except APIError as e:
            raise RuntimeError(f"Failed to stop container: {e}") from e

    def destroy(self, container_id: str, force: bool = False) -> None:
        """Remove a container."""
        try:
            container = self.client.containers.get(container_id)
            container.remove(force=force)
        except NotFound as e:
            raise RuntimeError(f"Container not found: {container_id}") from e
        except APIError as e:
            raise RuntimeError(f"Failed to remove container: {e}") from e

    def exec(
        self,
        container_id: str,
        command: list[str],
        *,
        interactive: bool = False,
        tty: bool = False,
        user: str | None = None,
        workdir: str | None = None,
        environment: dict[str, str] | None = None,
    ) -> tuple[int, str, str]:
        """Execute a command in a running container."""
        try:
            container = self.client.containers.get(container_id)
            result = container.exec_run(
                cmd=command,
                stdin=interactive,
                tty=tty,
                user=user or "",
                workdir=workdir or "",
                environment=environment,
                demux=True,
            )

            stdout = ""
            stderr = ""
            if result.output:
                if isinstance(result.output, tuple):
                    stdout = result.output[0].decode() if result.output[0] else ""
                    stderr = result.output[1].decode() if result.output[1] else ""
                else:
                    stdout = result.output.decode() if result.output else ""

            return result.exit_code, stdout, stderr
        except NotFound as e:
            raise RuntimeError(f"Container not found: {container_id}") from e
        except APIError as e:
            raise RuntimeError(f"Failed to execute command: {e}") from e

    def list(self, all: bool = False) -> list[ContainerState]:
        """List containers."""
        filters = {"label": f"{LABEL_PREFIX}.managed=true"}
        containers = self.client.containers.list(all=all, filters=filters)
        return [self._container_to_state(c) for c in containers]

    def status(self, container_id: str) -> ContainerState:
        """Get the status of a container."""
        try:
            container = self.client.containers.get(container_id)
            return self._container_to_state(container)
        except NotFound as e:
            raise RuntimeError(f"Container not found: {container_id}") from e

    def logs(
        self,
        container_id: str,
        *,
        follow: bool = False,
        tail: int | None = None,
        since: str | None = None,
    ) -> str:
        """Get container logs."""
        try:
            container = self.client.containers.get(container_id)
            logs = container.logs(
                follow=follow,
                tail=tail if tail else "all",
                since=since,
                timestamps=True,
            )
            if isinstance(logs, bytes):
                return logs.decode()
            return "".join(chunk.decode() for chunk in logs)
        except NotFound as e:
            raise RuntimeError(f"Container not found: {container_id}") from e
        except APIError as e:
            raise RuntimeError(f"Failed to get logs: {e}") from e

    def pause(self, container_id: str) -> None:
        """Pause a running container."""
        try:
            container = self.client.containers.get(container_id)
            container.pause()
        except NotFound as e:
            raise RuntimeError(f"Container not found: {container_id}") from e
        except APIError as e:
            raise RuntimeError(f"Failed to pause container: {e}") from e

    def unpause(self, container_id: str) -> None:
        """Unpause a paused container."""
        try:
            container = self.client.containers.get(container_id)
            container.unpause()
        except NotFound as e:
            raise RuntimeError(f"Container not found: {container_id}") from e
        except APIError as e:
            raise RuntimeError(f"Failed to unpause container: {e}") from e

    def copy_to(self, container_id: str, src: str, dest: str) -> None:
        """Copy a file or directory into a container."""
        try:
            container = self.client.containers.get(container_id)
            src_path = Path(src)

            # Create tar archive of the source
            tar_stream = io.BytesIO()
            with tarfile.open(fileobj=tar_stream, mode="w") as tar:
                tar.add(src_path, arcname=src_path.name)
            tar_stream.seek(0)

            container.put_archive(dest, tar_stream)
        except NotFound as e:
            raise RuntimeError(f"Container not found: {container_id}") from e
        except APIError as e:
            raise RuntimeError(f"Failed to copy to container: {e}") from e

    def copy_from(self, container_id: str, src: str, dest: str) -> None:
        """Copy a file or directory from a container."""
        try:
            container = self.client.containers.get(container_id)
            bits, _ = container.get_archive(src)

            # Extract tar archive to destination
            tar_stream = io.BytesIO()
            for chunk in bits:
                tar_stream.write(chunk)
            tar_stream.seek(0)

            dest_path = Path(dest)
            dest_path.parent.mkdir(parents=True, exist_ok=True)

            with tarfile.open(fileobj=tar_stream, mode="r") as tar:
                tar.extractall(dest_path.parent)
        except NotFound as e:
            raise RuntimeError(f"Container not found: {container_id}") from e
        except APIError as e:
            raise RuntimeError(f"Failed to copy from container: {e}") from e

    def _snapshot_image_name(self, container_name: str, snapshot_name: str) -> str:
        """Generate image name for a snapshot."""
        return f"gaol-snapshot/{container_name}:{snapshot_name}"

    def snapshot_create(
        self, container_id: str, name: str, description: str | None = None
    ) -> Snapshot:
        """Create a snapshot of a container using docker commit."""
        from gaol.models.container import Snapshot

        try:
            container = self.client.containers.get(container_id)
            container_name = container.name or container_id

            # Create image from container
            message = description or f"Snapshot {name} of {container_name}"

            image = container.commit(
                repository=f"gaol-snapshot/{container_name}",
                tag=name,
                message=message,
            )

            # Get image size
            image.reload()
            size = image.attrs.get("Size", 0)

            return Snapshot(
                id=image.id[:12] if image.id else "",
                name=name,
                container_id=container.id[:12] if container.id else "",
                container_name=container_name,
                created_at=datetime.now(),
                size_bytes=size,
                description=description,
            )
        except NotFound as e:
            raise RuntimeError(f"Container not found: {container_id}") from e
        except APIError as e:
            raise RuntimeError(f"Failed to create snapshot: {e}") from e

    def snapshot_list(self, container_id: str) -> list[Snapshot]:
        """List snapshots for a container."""
        from gaol.models.container import Snapshot

        try:
            container = self.client.containers.get(container_id)
            container_name = container.name or container_id

            # Find images matching our snapshot naming convention
            repo_name = f"gaol-snapshot/{container_name}"
            images = self.client.images.list(name=repo_name)

            snapshots = []
            for image in images:
                # Get tag from RepoTags
                for tag in image.tags or []:
                    if tag.startswith(f"{repo_name}:"):
                        snapshot_name = tag.split(":")[-1]
                        created_at = self._parse_datetime(image.attrs.get("Created"))
                        snapshots.append(
                            Snapshot(
                                id=image.id[:12] if image.id else "",
                                name=snapshot_name,
                                container_id=container.id[:12] if container.id else "",
                                container_name=container_name,
                                created_at=created_at or datetime.now(),
                                size_bytes=image.attrs.get("Size"),
                            )
                        )

            return sorted(snapshots, key=lambda s: s.created_at, reverse=True)
        except NotFound as e:
            raise RuntimeError(f"Container not found: {container_id}") from e
        except APIError as e:
            raise RuntimeError(f"Failed to list snapshots: {e}") from e

    def snapshot_restore(self, container_id: str, snapshot_name: str) -> None:
        """Restore a container from a snapshot.

        This stops the container, removes it, and creates a new one from the snapshot.
        """
        try:
            container = self.client.containers.get(container_id)
            container_name = container.name or container_id

            # Get the snapshot image
            image_name = self._snapshot_image_name(container_name, snapshot_name)
            try:
                self.client.images.get(image_name)
            except NotFound:
                raise RuntimeError(f"Snapshot not found: {snapshot_name}") from None

            # Get container config before removing
            container.reload()
            attrs = container.attrs
            config = attrs.get("Config", {})
            host_config = attrs.get("HostConfig", {})

            # Stop and remove the container
            with contextlib.suppress(APIError):
                container.stop(timeout=10)

            container.remove(force=True)

            # Create new container from snapshot image
            new_container = self.client.containers.create(
                image=image_name,
                name=container_name,
                command=config.get("Cmd"),
                environment=config.get("Env"),
                ports=host_config.get("PortBindings"),
                volumes=host_config.get("Binds"),
                network_mode=host_config.get("NetworkMode"),
                privileged=host_config.get("Privileged", False),
                labels=config.get("Labels"),
                stdin_open=config.get("OpenStdin", True),
                tty=config.get("Tty", True),
            )

            # Start the new container
            new_container.start()

        except NotFound as e:
            raise RuntimeError(f"Container not found: {container_id}") from e
        except APIError as e:
            raise RuntimeError(f"Failed to restore snapshot: {e}") from e

    def snapshot_delete(self, container_id: str, snapshot_name: str) -> None:
        """Delete a snapshot image."""
        try:
            container = self.client.containers.get(container_id)
            container_name = container.name or container_id

            image_name = self._snapshot_image_name(container_name, snapshot_name)
            try:
                self.client.images.remove(image_name)
            except NotFound:
                raise RuntimeError(f"Snapshot not found: {snapshot_name}") from None

        except NotFound as e:
            raise RuntimeError(f"Container not found: {container_id}") from e
        except APIError as e:
            raise RuntimeError(f"Failed to delete snapshot: {e}") from e

    def stats(self, container_id: str) -> ContainerStats:
        """Get resource usage statistics for a container."""
        from gaol.models.container import ContainerStats

        try:
            container = self.client.containers.get(container_id)
            container_name = container.name or container_id

            # Get stats (non-streaming, single snapshot)
            stats_data = container.stats(stream=False)

            # Calculate CPU percentage
            cpu_percent = 0.0
            cpu_delta = (
                stats_data["cpu_stats"]["cpu_usage"]["total_usage"]
                - stats_data["precpu_stats"]["cpu_usage"]["total_usage"]
            )
            system_delta = (
                stats_data["cpu_stats"]["system_cpu_usage"]
                - stats_data["precpu_stats"]["system_cpu_usage"]
            )
            cpu_count = stats_data["cpu_stats"].get("online_cpus") or len(
                stats_data["cpu_stats"]["cpu_usage"].get("percpu_usage", [1])
            )

            if system_delta > 0 and cpu_delta > 0:
                cpu_percent = (cpu_delta / system_delta) * cpu_count * 100.0

            # Memory stats
            memory_stats = stats_data.get("memory_stats", {})
            memory_used = memory_stats.get("usage", 0)
            # Subtract cache for more accurate "used" memory
            cache = memory_stats.get("stats", {}).get("cache", 0)
            memory_used = max(0, memory_used - cache)
            memory_limit = memory_stats.get("limit", 0)
            memory_percent = (memory_used / memory_limit * 100.0) if memory_limit > 0 else 0.0

            # Network stats (sum all interfaces)
            network_rx = 0
            network_tx = 0
            networks = stats_data.get("networks", {})
            for iface_stats in networks.values():
                network_rx += iface_stats.get("rx_bytes", 0)
                network_tx += iface_stats.get("tx_bytes", 0)

            # Block I/O stats
            block_read = 0
            block_write = 0
            blkio_stats = stats_data.get("blkio_stats", {})
            for entry in blkio_stats.get("io_service_bytes_recursive", []) or []:
                if entry.get("op") == "read":
                    block_read += entry.get("value", 0)
                elif entry.get("op") == "write":
                    block_write += entry.get("value", 0)

            # Process count
            pids = stats_data.get("pids_stats", {}).get("current", 0)

            return ContainerStats(
                container_id=container.id,
                container_name=container_name,
                cpu_percent=round(cpu_percent, 2),
                cpu_count=cpu_count,
                memory_used_bytes=memory_used,
                memory_limit_bytes=memory_limit if memory_limit > 0 else None,
                memory_percent=round(memory_percent, 2),
                network_rx_bytes=network_rx,
                network_tx_bytes=network_tx,
                block_read_bytes=block_read,
                block_write_bytes=block_write,
                pids=pids,
            )

        except NotFound as e:
            raise RuntimeError(f"Container not found: {container_id}") from e
        except APIError as e:
            raise RuntimeError(f"Failed to get stats: {e}") from e

    def export_filesystem(
        self,
        container_id: str,
        dest_path: str,
        *,
        compression: str | None = "gzip",
    ) -> str:
        """Export container filesystem to a tarball."""
        from pathlib import Path

        try:
            container = self.client.containers.get(container_id)
            container_name = container.name or container_id

            # Determine output filename
            if compression == "gzip":
                ext = ".tar.gz"
            elif compression == "xz":
                ext = ".tar.xz"
            else:
                ext = ".tar"

            output_path = Path(dest_path) / f"{container_name}-export{ext}"

            # Export container to tar stream
            bits, _ = container.export()

            # Write to file with optional compression
            if compression == "gzip":
                import gzip

                with gzip.open(output_path, "wb") as f:
                    for chunk in bits:
                        f.write(chunk)
            elif compression == "xz":
                import lzma

                with lzma.open(output_path, "wb") as f:
                    for chunk in bits:
                        f.write(chunk)
            else:
                with open(output_path, "wb") as f:
                    for chunk in bits:
                        f.write(chunk)

            return str(output_path)

        except NotFound as e:
            raise RuntimeError(f"Container not found: {container_id}") from e
        except APIError as e:
            raise RuntimeError(f"Failed to export filesystem: {e}") from e

    def import_filesystem(
        self,
        config: ContainerConfig,
        source_path: str,
    ) -> str:
        """Create container by importing filesystem from tarball."""
        from pathlib import Path

        try:
            # Import tarball as a new image
            source = Path(source_path)
            if not source.exists():
                raise RuntimeError(f"Source file not found: {source_path}")

            # Detect compression and open appropriately
            if str(source).endswith(".gz"):
                import gzip

                with gzip.open(source, "rb") as f:
                    image = self.client.images.import_image(
                        f,
                        repository=f"gaol-migrated/{config.name}",
                        tag="latest",
                    )
            elif str(source).endswith(".xz"):
                import lzma

                with lzma.open(source, "rb") as f:
                    image = self.client.images.import_image(
                        f,
                        repository=f"gaol-migrated/{config.name}",
                        tag="latest",
                    )
            else:
                with open(source, "rb") as f:
                    image = self.client.images.import_image(
                        f,
                        repository=f"gaol-migrated/{config.name}",
                        tag="latest",
                    )

            # Update config to use the imported image
            migrated_config = config.model_copy(
                update={"image": f"gaol-migrated/{config.name}:latest"}
            )

            # Create container using existing create() method
            return self.create(migrated_config)

        except APIError as e:
            raise RuntimeError(f"Failed to import filesystem: {e}") from e

    def get_filesystem_size(self, container_id: str) -> int:
        """Get approximate filesystem size in bytes."""
        try:
            container = self.client.containers.get(container_id)
            # Get container info which includes size
            info = container.attrs
            # SizeRw is the size of files modified/added in the container
            # SizeRootFs is the total size including base image
            return info.get("SizeRootFs", 0) or info.get("SizeRw", 0)
        except NotFound as e:
            raise RuntimeError(f"Container not found: {container_id}") from e
        except APIError as e:
            raise RuntimeError(f"Failed to get filesystem size: {e}") from e

    # -------------------------------------------------------------------------
    # Checkpoint/Restore (CRIU) methods
    # -------------------------------------------------------------------------

    def checkpoint_create(
        self,
        container_id: str,
        checkpoint_name: str,
        image_dir: str | None = None,
        *,
        leave_running: bool = False,
        tcp_established: bool = False,
    ) -> "CheckpointResult":
        """Create a CRIU checkpoint of a running Docker container.

        Requires Docker experimental mode to be enabled.
        Add {"experimental": true} to /etc/docker/daemon.json.

        Args:
            container_id: ID of the container to checkpoint
            checkpoint_name: Name for the checkpoint
            image_dir: Directory to store checkpoint images (optional)
            leave_running: Keep container running after checkpoint
            tcp_established: Not used for Docker (handled internally)

        Returns:
            CheckpointResult with operation status
        """
        from gaol.checkpoint.executor import DockerCheckpointExecutor
        from gaol.checkpoint.models import CheckpointResult

        try:
            container = self.client.containers.get(container_id)
            container_name = container.name or container_id

            executor = DockerCheckpointExecutor()
            checkpoint_dir = Path(image_dir) if image_dir else None

            return executor.checkpoint_create(
                container_id=container.id,
                checkpoint_name=checkpoint_name,
                container_name=container_name,
                checkpoint_dir=checkpoint_dir,
                leave_running=leave_running,
            )

        except NotFound as e:
            raise RuntimeError(f"Container not found: {container_id}") from e
        except APIError as e:
            raise RuntimeError(f"Failed to create checkpoint: {e}") from e

    def checkpoint_restore(
        self,
        container_id: str,
        checkpoint_name: str,
        image_dir: str | None = None,
        *,
        tcp_established: bool = False,
    ) -> "RestoreResult":
        """Restore a Docker container from a CRIU checkpoint.

        Args:
            container_id: ID of the container to restore
            checkpoint_name: Name of the checkpoint to restore from
            image_dir: Directory containing checkpoint images
            tcp_established: Not used for Docker

        Returns:
            RestoreResult with operation status
        """
        from gaol.checkpoint.executor import DockerCheckpointExecutor
        from gaol.checkpoint.models import RestoreResult

        try:
            container = self.client.containers.get(container_id)
            container_name = container.name or container_id

            executor = DockerCheckpointExecutor()
            checkpoint_dir = Path(image_dir) if image_dir else None

            return executor.restore(
                container_id=container.id,
                checkpoint_name=checkpoint_name,
                container_name=container_name,
                checkpoint_dir=checkpoint_dir,
            )

        except NotFound as e:
            raise RuntimeError(f"Container not found: {container_id}") from e
        except APIError as e:
            raise RuntimeError(f"Failed to restore checkpoint: {e}") from e

    def checkpoint_list(self, container_id: str) -> list[str]:
        """List available checkpoints for a Docker container.

        Args:
            container_id: ID of the container

        Returns:
            List of checkpoint names
        """
        from gaol.checkpoint.executor import DockerCheckpointExecutor

        try:
            container = self.client.containers.get(container_id)
            executor = DockerCheckpointExecutor()
            return executor.checkpoint_list(container.id)

        except NotFound as e:
            raise RuntimeError(f"Container not found: {container_id}") from e
        except APIError as e:
            raise RuntimeError(f"Failed to list checkpoints: {e}") from e

    def checkpoint_delete(
        self,
        container_id: str,
        checkpoint_name: str,
        image_dir: str | None = None,
    ) -> bool:
        """Delete a Docker checkpoint.

        Args:
            container_id: ID of the container
            checkpoint_name: Name of the checkpoint to delete
            image_dir: Directory containing checkpoint images

        Returns:
            True if checkpoint was deleted
        """
        from gaol.checkpoint.executor import DockerCheckpointExecutor

        try:
            container = self.client.containers.get(container_id)
            executor = DockerCheckpointExecutor()
            checkpoint_dir = Path(image_dir) if image_dir else None
            return executor.checkpoint_delete(
                container.id, checkpoint_name, checkpoint_dir=checkpoint_dir
            )

        except NotFound as e:
            raise RuntimeError(f"Container not found: {container_id}") from e
        except APIError as e:
            raise RuntimeError(f"Failed to delete checkpoint: {e}") from e

    def checkpoint_capabilities(self) -> dict[str, bool]:
        """Get Docker checkpoint capabilities.

        Returns:
            Dictionary with capability names and availability
        """
        from gaol.checkpoint.executor import DockerCheckpointExecutor

        caps = DockerCheckpointExecutor.check_capabilities()
        return {
            "available": caps.available,
            "experimental_enabled": caps.experimental_enabled,
            "checkpoint_supported": caps.checkpoint_supported,
            "tcp_established": caps.checkpoint_supported,  # Handled by Docker
            "leave_running": caps.checkpoint_supported,
        }
