"""Low-level CRIU and Docker checkpoint execution."""

from __future__ import annotations

import hashlib
import logging
import os
import shutil
import subprocess
from datetime import datetime
from pathlib import Path
from uuid import uuid4

from gaol.checkpoint.exceptions import (
    CRIUKernelError,
    CRIUNotAvailableError,
    CRIUPermissionError,
    DockerExperimentalError,
)
from gaol.checkpoint.models import (
    CheckpointResult,
    CheckpointStatus,
    CRIUCapabilities,
    DockerCheckpointCapabilities,
    RestoreResult,
)

logger = logging.getLogger(__name__)


def _run_cmd(
    cmd: list[str],
    *,
    check: bool = True,
    capture_output: bool = True,
) -> subprocess.CompletedProcess[str]:
    """Run a command and return the result."""
    try:
        result = subprocess.run(
            cmd,
            check=check,
            capture_output=capture_output,
            text=True,
        )
        return result
    except subprocess.CalledProcessError as e:
        raise RuntimeError(
            f"Command failed: {' '.join(cmd)}\n"
            f"stdout: {e.stdout}\n"
            f"stderr: {e.stderr}"
        ) from e


class CRIUExecutor:
    """Direct CRIU command execution for checkpoint/restore.

    CRIU (Checkpoint/Restore In Userspace) allows checkpointing running
    processes including their memory, file descriptors, and network connections.

    Requirements:
        - Linux kernel >= 3.11 with CONFIG_CHECKPOINT_RESTORE
        - CRIU package installed (apt install criu)
        - CAP_SYS_ADMIN capability or root
        - /proc/sys/kernel/ns_last_pid must exist
    """

    @staticmethod
    def is_available() -> bool:
        """Check if CRIU is installed and accessible."""
        return shutil.which("criu") is not None

    @staticmethod
    def get_version() -> str | None:
        """Get the CRIU version string."""
        try:
            result = _run_cmd(["criu", "--version"], check=False)
            if result.returncode == 0:
                # Parse "Version: X.Y.Z"
                for line in result.stdout.split("\n"):
                    if line.startswith("Version:"):
                        return line.split(":", 1)[1].strip()
                return result.stdout.strip().split("\n")[0]
        except Exception:
            pass
        return None

    @classmethod
    def check_capabilities(cls) -> CRIUCapabilities:
        """Detect available CRIU features and kernel support."""
        caps = CRIUCapabilities(available=cls.is_available())

        if not caps.available:
            return caps

        version = cls.get_version()

        # Check kernel features
        kernel_cr = Path("/proc/sys/kernel/ns_last_pid").exists()
        ns_last_pid = Path("/proc/sys/kernel/ns_last_pid").exists()

        # Check for userfaultfd support
        userfaultfd = Path("/proc/sys/vm/unprivileged_userfaultfd").exists()

        # Run CRIU check command
        try:
            result = _run_cmd(["criu", "check"], check=False)
            criu_works = result.returncode == 0
        except Exception:
            criu_works = False

        # Check CAP_SYS_ADMIN
        has_cap = os.geteuid() == 0

        # Try to detect TCP support
        tcp_supported = False
        file_locks = False
        lazy_pages = False

        try:
            # Check CRIU features via check command with specific options
            result = _run_cmd(
                ["criu", "check", "--feature", "tcp_established"],
                check=False,
            )
            tcp_supported = result.returncode == 0

            result = _run_cmd(
                ["criu", "check", "--feature", "file_locks"],
                check=False,
            )
            file_locks = result.returncode == 0

            result = _run_cmd(
                ["criu", "check", "--feature", "lazy_pages"],
                check=False,
            )
            lazy_pages = result.returncode == 0
        except Exception:
            pass

        return CRIUCapabilities(
            available=caps.available and criu_works,
            version=version,
            kernel_checkpoint_restore=kernel_cr,
            kernel_ns_last_pid=ns_last_pid,
            kernel_userfaultfd=userfaultfd,
            tcp_established=tcp_supported,
            external_unix_sockets=True,  # Usually available
            file_locks=file_locks,
            lazy_pages=lazy_pages,
            has_cap_sys_admin=has_cap,
            can_checkpoint_namespaces=criu_works and has_cap,
        )

    @classmethod
    def require_available(cls) -> None:
        """Raise an exception if CRIU is not available."""
        if not cls.is_available():
            raise CRIUNotAvailableError("CRIU binary not found")

        caps = cls.check_capabilities()
        if not caps.kernel_ns_last_pid:
            raise CRIUKernelError("ns_last_pid")
        if not caps.has_cap_sys_admin:
            raise CRIUPermissionError("checkpoint", "requires root or CAP_SYS_ADMIN")

    def dump(
        self,
        pid: int,
        image_dir: Path,
        *,
        checkpoint_name: str,
        container_name: str,
        leave_running: bool = False,
        tcp_established: bool = False,
        external_unix_sockets: bool = False,
        shell_job: bool = True,
    ) -> CheckpointResult:
        """Checkpoint a process tree.

        Args:
            pid: Leader process ID to checkpoint
            image_dir: Directory to store CRIU images
            checkpoint_name: Name for this checkpoint
            container_name: Container name
            leave_running: Keep process running after checkpoint
            tcp_established: Checkpoint established TCP connections
            external_unix_sockets: Allow external unix sockets
            shell_job: Process is a shell job

        Returns:
            CheckpointResult with operation status
        """
        self.require_available()

        checkpoint_id = str(uuid4())
        result = CheckpointResult(
            container_name=container_name,
            checkpoint_name=checkpoint_name,
            checkpoint_id=checkpoint_id,
            started_at=datetime.now(),
            status=CheckpointStatus.FREEZING,
            leave_running=leave_running,
            pid_at_checkpoint=pid,
        )

        # Ensure image directory exists
        image_dir.mkdir(parents=True, exist_ok=True)

        cmd = [
            "criu", "dump",
            "-t", str(pid),
            "-D", str(image_dir),
            "--images-dir", str(image_dir),
            "-v4",  # Verbose for debugging
            "-o", str(image_dir / "dump.log"),
        ]

        if leave_running:
            cmd.append("--leave-running")

        if tcp_established:
            cmd.append("--tcp-established")

        if external_unix_sockets:
            cmd.append("--ext-unix-sk")

        if shell_job:
            cmd.append("--shell-job")

        try:
            result.status = CheckpointStatus.CHECKPOINTING
            logger.info(f"Running CRIU dump: {' '.join(cmd)}")

            proc_result = _run_cmd(cmd, check=False)

            if proc_result.returncode != 0:
                result.status = CheckpointStatus.FAILED
                result.error_message = "CRIU dump failed"
                result.error_output = proc_result.stderr
                logger.error(f"CRIU dump failed: {proc_result.stderr}")
            else:
                result.status = CheckpointStatus.COMPLETED
                result.completed_at = datetime.now()
                result.image_dir = str(image_dir)

                # Calculate size
                result.size_bytes = sum(
                    f.stat().st_size for f in image_dir.rglob("*") if f.is_file()
                )

                # Calculate checksum of core image
                core_files = list(image_dir.glob("core-*.img"))
                if core_files:
                    result.checksum = self._calculate_checksum(core_files[0])

                logger.info(
                    f"CRIU dump completed: {result.size_bytes} bytes"
                )

        except Exception as e:
            result.status = CheckpointStatus.FAILED
            result.error_message = str(e)
            logger.error(f"CRIU dump error: {e}")

        return result

    def restore(
        self,
        image_dir: Path,
        *,
        checkpoint_name: str,
        container_name: str,
        tcp_established: bool = False,
        shell_job: bool = True,
    ) -> RestoreResult:
        """Restore a process from checkpoint.

        Args:
            image_dir: Directory containing CRIU images
            checkpoint_name: Name of the checkpoint
            container_name: Container name
            tcp_established: Restore established TCP connections
            shell_job: Process is a shell job

        Returns:
            RestoreResult with operation status
        """
        self.require_available()

        result = RestoreResult(
            container_name=container_name,
            checkpoint_name=checkpoint_name,
            checkpoint_id="",  # Will be set from metadata
            started_at=datetime.now(),
            status=CheckpointStatus.RESTORING,
        )

        cmd = [
            "criu", "restore",
            "-D", str(image_dir),
            "--images-dir", str(image_dir),
            "-v4",
            "-o", str(image_dir / "restore.log"),
            "--detach",
        ]

        if tcp_established:
            cmd.append("--tcp-established")

        if shell_job:
            cmd.append("--shell-job")

        try:
            logger.info(f"Running CRIU restore: {' '.join(cmd)}")
            proc_result = _run_cmd(cmd, check=False)

            if proc_result.returncode != 0:
                result.status = CheckpointStatus.FAILED
                result.error_message = "CRIU restore failed"
                result.error_output = proc_result.stderr
                logger.error(f"CRIU restore failed: {proc_result.stderr}")
            else:
                result.status = CheckpointStatus.COMPLETED
                result.completed_at = datetime.now()
                logger.info("CRIU restore completed")

        except Exception as e:
            result.status = CheckpointStatus.FAILED
            result.error_message = str(e)
            logger.error(f"CRIU restore error: {e}")

        return result

    @staticmethod
    def _calculate_checksum(file_path: Path) -> str:
        """Calculate SHA256 checksum of a file."""
        sha256 = hashlib.sha256()
        with open(file_path, "rb") as f:
            for chunk in iter(lambda: f.read(8192), b""):
                sha256.update(chunk)
        return sha256.hexdigest()


class DockerCheckpointExecutor:
    """Docker checkpoint API wrapper.

    Docker's experimental checkpoint feature uses CRIU internally
    to checkpoint running containers.

    Requirements:
        - Docker with experimental features enabled
        - Add {"experimental": true} to /etc/docker/daemon.json
    """

    @staticmethod
    def is_available() -> bool:
        """Check if Docker is available."""
        return shutil.which("docker") is not None

    @staticmethod
    def is_experimental_enabled() -> bool:
        """Check if Docker experimental mode is enabled."""
        if not DockerCheckpointExecutor.is_available():
            return False

        try:
            result = _run_cmd(
                ["docker", "info", "--format", "{{.ExperimentalBuild}}"],
                check=False,
            )
            return result.returncode == 0 and result.stdout.strip().lower() == "true"
        except Exception:
            return False

    @staticmethod
    def get_version() -> str | None:
        """Get Docker version string."""
        try:
            result = _run_cmd(
                ["docker", "version", "--format", "{{.Server.Version}}"],
                check=False,
            )
            if result.returncode == 0:
                return result.stdout.strip()
        except Exception:
            pass
        return None

    @classmethod
    def check_capabilities(cls) -> DockerCheckpointCapabilities:
        """Detect Docker checkpoint capabilities."""
        available = cls.is_available()
        experimental = cls.is_experimental_enabled() if available else False
        version = cls.get_version() if available else None

        return DockerCheckpointCapabilities(
            available=available,
            experimental_enabled=experimental,
            checkpoint_supported=experimental,
            version=version,
        )

    @classmethod
    def require_experimental(cls) -> None:
        """Raise an exception if Docker experimental is not enabled."""
        if not cls.is_available():
            raise CRIUNotAvailableError("Docker not found")
        if not cls.is_experimental_enabled():
            raise DockerExperimentalError()

    def checkpoint_create(
        self,
        container_id: str,
        checkpoint_name: str,
        *,
        container_name: str,
        checkpoint_dir: Path | None = None,
        leave_running: bool = False,
    ) -> CheckpointResult:
        """Create a checkpoint for a Docker container.

        Args:
            container_id: Docker container ID or name
            checkpoint_name: Name for this checkpoint
            container_name: Gaol container name
            checkpoint_dir: Custom directory for checkpoint
            leave_running: Keep container running after checkpoint

        Returns:
            CheckpointResult with operation status
        """
        self.require_experimental()

        checkpoint_id = str(uuid4())
        result = CheckpointResult(
            container_name=container_name,
            checkpoint_name=checkpoint_name,
            checkpoint_id=checkpoint_id,
            started_at=datetime.now(),
            status=CheckpointStatus.CHECKPOINTING,
            leave_running=leave_running,
        )

        cmd = [
            "docker", "checkpoint", "create",
        ]

        if checkpoint_dir:
            checkpoint_dir.mkdir(parents=True, exist_ok=True)
            cmd.extend(["--checkpoint-dir", str(checkpoint_dir)])

        if leave_running:
            cmd.append("--leave-running")

        cmd.extend([container_id, checkpoint_name])

        try:
            logger.info(f"Creating Docker checkpoint: {' '.join(cmd)}")
            proc_result = _run_cmd(cmd, check=False)

            if proc_result.returncode != 0:
                result.status = CheckpointStatus.FAILED
                result.error_message = "Docker checkpoint create failed"
                result.error_output = proc_result.stderr
                logger.error(f"Docker checkpoint failed: {proc_result.stderr}")
            else:
                result.status = CheckpointStatus.COMPLETED
                result.completed_at = datetime.now()
                if checkpoint_dir:
                    result.image_dir = str(checkpoint_dir / checkpoint_name)
                logger.info("Docker checkpoint created successfully")

        except Exception as e:
            result.status = CheckpointStatus.FAILED
            result.error_message = str(e)
            logger.error(f"Docker checkpoint error: {e}")

        return result

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

        Args:
            container_id: Docker container ID or name

        Returns:
            List of checkpoint names
        """
        self.require_experimental()

        try:
            result = _run_cmd(
                ["docker", "checkpoint", "ls", container_id, "--format", "{{.Name}}"],
                check=False,
            )
            if result.returncode == 0:
                return [
                    name.strip()
                    for name in result.stdout.strip().split("\n")
                    if name.strip()
                ]
        except Exception as e:
            logger.error(f"Error listing checkpoints: {e}")

        return []

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

        Args:
            container_id: Docker container ID or name
            checkpoint_name: Name of checkpoint to delete
            checkpoint_dir: Custom checkpoint directory

        Returns:
            True if deleted successfully
        """
        self.require_experimental()

        cmd = ["docker", "checkpoint", "rm"]

        if checkpoint_dir:
            cmd.extend(["--checkpoint-dir", str(checkpoint_dir)])

        cmd.extend([container_id, checkpoint_name])

        try:
            result = _run_cmd(cmd, check=False)
            return result.returncode == 0
        except Exception as e:
            logger.error(f"Error deleting checkpoint: {e}")
            return False

    def restore(
        self,
        container_id: str,
        checkpoint_name: str,
        *,
        container_name: str,
        checkpoint_dir: Path | None = None,
    ) -> RestoreResult:
        """Restore a container from a checkpoint.

        Note: Docker restores by starting the container with --checkpoint flag.

        Args:
            container_id: Docker container ID or name
            checkpoint_name: Name of checkpoint to restore
            container_name: Gaol container name
            checkpoint_dir: Custom checkpoint directory

        Returns:
            RestoreResult with operation status
        """
        self.require_experimental()

        result = RestoreResult(
            container_name=container_name,
            checkpoint_name=checkpoint_name,
            checkpoint_id="",
            started_at=datetime.now(),
            status=CheckpointStatus.RESTORING,
        )

        cmd = ["docker", "start", f"--checkpoint={checkpoint_name}"]

        if checkpoint_dir:
            cmd.append(f"--checkpoint-dir={checkpoint_dir}")

        cmd.append(container_id)

        try:
            logger.info(f"Restoring Docker container: {' '.join(cmd)}")
            proc_result = _run_cmd(cmd, check=False)

            if proc_result.returncode != 0:
                result.status = CheckpointStatus.FAILED
                result.error_message = "Docker checkpoint restore failed"
                result.error_output = proc_result.stderr
                logger.error(f"Docker restore failed: {proc_result.stderr}")
            else:
                result.status = CheckpointStatus.COMPLETED
                result.completed_at = datetime.now()
                result.new_container_id = container_id
                logger.info("Docker container restored successfully")

        except Exception as e:
            result.status = CheckpointStatus.FAILED
            result.error_message = str(e)
            logger.error(f"Docker restore error: {e}")

        return result


class NspawnCheckpointExecutor:
    """CRIU-based checkpoint executor for systemd-nspawn containers.

    Uses machinectl to get container information and CRIU directly
    for checkpoint/restore operations.

    Requirements:
        - systemd-nspawn container running
        - CRIU installed
        - Root privileges
    """

    def __init__(self) -> None:
        """Initialize the executor."""
        self._criu = CRIUExecutor()

    @staticmethod
    def is_available() -> bool:
        """Check if nspawn checkpoint is available."""
        return (
            CRIUExecutor.is_available()
            and shutil.which("machinectl") is not None
        )

    def get_leader_pid(self, machine_name: str) -> int | None:
        """Get the leader PID for a systemd-nspawn container.

        Args:
            machine_name: Name of the machine/container

        Returns:
            Leader PID or None if not found/not running
        """
        try:
            result = _run_cmd(
                ["machinectl", "show", machine_name, "--property=Leader"],
                check=False,
            )
            if result.returncode == 0 and "=" in result.stdout:
                pid_str = result.stdout.strip().split("=")[1]
                if pid_str and pid_str != "0":
                    return int(pid_str)
        except Exception as e:
            logger.error(f"Error getting leader PID: {e}")

        return None

    def get_machine_state(self, machine_name: str) -> str | None:
        """Get the state of a machine.

        Args:
            machine_name: Name of the machine/container

        Returns:
            State string or None
        """
        try:
            result = _run_cmd(
                ["machinectl", "show", machine_name, "--property=State"],
                check=False,
            )
            if result.returncode == 0 and "=" in result.stdout:
                return result.stdout.strip().split("=")[1]
        except Exception as e:
            logger.error(f"Error getting machine state: {e}")

        return None

    def freeze_container(self, machine_name: str) -> bool:
        """Freeze a systemd-nspawn container before checkpoint.

        Uses systemctl freeze to pause the container.

        Args:
            machine_name: Name of the container

        Returns:
            True if frozen successfully
        """
        try:
            result = _run_cmd(
                ["systemctl", "freeze", f"systemd-nspawn@{machine_name}.service"],
                check=False,
            )
            return result.returncode == 0
        except Exception as e:
            logger.error(f"Error freezing container: {e}")
            return False

    def thaw_container(self, machine_name: str) -> bool:
        """Thaw (unfreeze) a systemd-nspawn container.

        Args:
            machine_name: Name of the container

        Returns:
            True if thawed successfully
        """
        try:
            result = _run_cmd(
                ["systemctl", "thaw", f"systemd-nspawn@{machine_name}.service"],
                check=False,
            )
            return result.returncode == 0
        except Exception as e:
            logger.error(f"Error thawing container: {e}")
            return False

    def checkpoint_create(
        self,
        machine_name: str,
        image_dir: Path,
        checkpoint_name: str,
        *,
        leave_running: bool = False,
        tcp_established: bool = False,
    ) -> CheckpointResult:
        """Create a checkpoint for a systemd-nspawn container.

        Args:
            machine_name: Name of the machine/container
            image_dir: Directory to store CRIU images
            checkpoint_name: Name for this checkpoint
            leave_running: Keep container running after checkpoint
            tcp_established: Checkpoint TCP connections

        Returns:
            CheckpointResult with operation status
        """
        CRIUExecutor.require_available()

        checkpoint_id = str(uuid4())
        result = CheckpointResult(
            container_name=machine_name,
            checkpoint_name=checkpoint_name,
            checkpoint_id=checkpoint_id,
            started_at=datetime.now(),
            status=CheckpointStatus.PENDING,
            leave_running=leave_running,
        )

        # Get leader PID
        pid = self.get_leader_pid(machine_name)
        if pid is None:
            result.status = CheckpointStatus.FAILED
            result.error_message = f"Container {machine_name} is not running"
            return result

        result.pid_at_checkpoint = pid
        result.container_was_running = True

        # Freeze container
        result.status = CheckpointStatus.FREEZING
        logger.info(f"Freezing container {machine_name}")

        if not self.freeze_container(machine_name):
            result.status = CheckpointStatus.FAILED
            result.error_message = "Failed to freeze container"
            return result

        try:
            # Run CRIU dump
            criu_result = self._criu.dump(
                pid=pid,
                image_dir=image_dir,
                checkpoint_name=checkpoint_name,
                container_name=machine_name,
                leave_running=leave_running,
                tcp_established=tcp_established,
                shell_job=False,  # nspawn container, not shell job
            )

            # Copy result status
            result.status = criu_result.status
            result.completed_at = criu_result.completed_at
            result.size_bytes = criu_result.size_bytes
            result.image_dir = criu_result.image_dir
            result.checksum = criu_result.checksum
            result.error_message = criu_result.error_message
            result.error_output = criu_result.error_output

        finally:
            # Thaw container if leave_running or on failure
            if leave_running or result.status == CheckpointStatus.FAILED:
                logger.info(f"Thawing container {machine_name}")
                self.thaw_container(machine_name)

        return result

    def restore(
        self,
        machine_name: str,
        image_dir: Path,
        checkpoint_name: str,
        *,
        tcp_established: bool = False,
    ) -> RestoreResult:
        """Restore a systemd-nspawn container from checkpoint.

        Note: This is complex for nspawn as we need to restore
        into the container's namespace. For now, this provides
        basic restore functionality.

        Args:
            machine_name: Name of the machine/container
            image_dir: Directory containing CRIU images
            checkpoint_name: Name of checkpoint
            tcp_established: Restore TCP connections

        Returns:
            RestoreResult with operation status
        """
        CRIUExecutor.require_available()

        result = RestoreResult(
            container_name=machine_name,
            checkpoint_name=checkpoint_name,
            checkpoint_id="",
            started_at=datetime.now(),
            status=CheckpointStatus.RESTORING,
        )

        # For nspawn, restore is more complex - we need to:
        # 1. Ensure the container/service exists
        # 2. Restore into the container's namespace

        criu_result = self._criu.restore(
            image_dir=image_dir,
            checkpoint_name=checkpoint_name,
            container_name=machine_name,
            tcp_established=tcp_established,
            shell_job=False,
        )

        result.status = criu_result.status
        result.completed_at = criu_result.completed_at
        result.error_message = criu_result.error_message
        result.error_output = criu_result.error_output

        return result
