"""Installation of host prerequisites, Firecracker, and guest assets."""

from __future__ import annotations

import os
import shutil
import tarfile
import tempfile
import urllib.error
import urllib.request
from collections.abc import Callable
from pathlib import Path

from .config import DEFAULT_KERNEL_URL, DEFAULT_ROOTFS_URL, Settings
from .errors import UvmError
from .images import GuestAssets
from .integrity import require_checksum, sha256_file, verify_file, write_manifest
from .state import StateStore
from .system import CommandRunner, check_kvm, ensure_data_directories, require_root


class Installer:
    """Install exactly the local host dependencies used by the uvm CLI."""

    def __init__(
        self,
        settings: Settings,
        runner: CommandRunner,
        state_store: StateStore,
        *,
        emit: Callable[[str], None] | None = print,
    ) -> None:
        self._settings = settings
        self._runner = runner
        self._state_store = state_store
        self._emit = emit

    def install(self, *, force_assets: bool = False) -> tuple[Path, GuestAssets]:
        require_root()
        ensure_data_directories(self._settings)
        with self._state_store.operation_lock():
            return self._install_locked(force_assets=force_assets)

    def _install_locked(self, *, force_assets: bool) -> tuple[Path, GuestAssets]:
        self._validate_guest_asset_architecture()
        self._validate_integrity_policy()
        if self._settings.allow_unverified_downloads and self._emit is not None:
            self._emit(
                "WARNING: downloads are not checksum verified. Set UVM_*_SHA256 values "
                "and UVM_ALLOW_UNVERIFIED_DOWNLOADS=0 to enforce verification."
            )
        self.install_apt_packages()
        check_kvm()
        firecracker = self.install_firecracker_binary(force=force_assets)
        assets = self.install_guest_assets(force=force_assets)
        self._write_integrity_manifest(firecracker, assets)
        return firecracker, assets

    def _validate_guest_asset_architecture(self) -> None:
        architecture = os.uname().machine
        default_kernel = self._settings.kernel_url == DEFAULT_KERNEL_URL
        default_rootfs = self._settings.rootfs_url == DEFAULT_ROOTFS_URL
        if architecture == "aarch64" and (default_kernel or default_rootfs):
            raise UvmError(
                "default guest assets support x86_64 only; set both UVM_KERNEL_URL "
                "and UVM_ROOTFS_URL to compatible aarch64 assets"
            )

    def _validate_integrity_policy(self) -> None:
        require_checksum(
            self._settings.firecracker_sha256,
            "Firecracker archive",
            "UVM_FIRECRACKER_SHA256",
            allow_unverified=self._settings.allow_unverified_downloads,
        )
        require_checksum(
            self._settings.kernel_sha256,
            "guest kernel",
            "UVM_KERNEL_SHA256",
            allow_unverified=self._settings.allow_unverified_downloads,
        )
        require_checksum(
            self._settings.rootfs_sha256,
            "guest root filesystem",
            "UVM_ROOTFS_SHA256",
            allow_unverified=self._settings.allow_unverified_downloads,
        )

    def _write_integrity_manifest(self, firecracker: Path, assets: GuestAssets) -> None:
        write_manifest(
            self._settings.integrity_manifest_path,
            {
                "firecracker": sha256_file(firecracker, "Firecracker binary"),
                "jailer": sha256_file(self._settings.jailer_binary, "Jailer binary"),
                "kernel": sha256_file(assets.kernel, "guest kernel"),
                "rootfs": sha256_file(assets.rootfs, "guest root filesystem"),
            },
            verified=all(
                (
                    self._settings.firecracker_sha256,
                    self._settings.kernel_sha256,
                    self._settings.rootfs_sha256,
                )
            ),
        )

    def install_apt_packages(self) -> None:
        packages = (
            "curl",
            "jq",
            "iproute2",
            "iptables",
            "e2fsprogs",
            "util-linux",
            "ca-certificates",
            "openssl",
            "openssh-client",
        )
        self._runner.run(("apt-get", "update"))
        self._runner.run(("apt-get", "install", "-y", *packages))

    def install_firecracker_binary(self, *, force: bool = False) -> Path:
        architecture = os.uname().machine
        if architecture not in ("x86_64", "aarch64"):
            raise UvmError(f"unsupported host architecture: {architecture}")

        firecracker = self._settings.firecracker_binary
        jailer = self._settings.jailer_binary
        if (
            firecracker.exists()
            and jailer.exists()
            and self._settings.allow_unverified_downloads
            and not force
        ):
            return firecracker

        archive = self._settings.base / (
            f"firecracker-{self._settings.firecracker_version}-{architecture}.tgz"
        )
        url = (
            "https://github.com/firecracker-microvm/firecracker/releases/download/"
            f"{self._settings.firecracker_version}/"
            f"firecracker-{self._settings.firecracker_version}-{architecture}.tgz"
        )
        if force or not archive.exists():
            self.download(
                url,
                archive,
                expected_sha256=self._settings.firecracker_sha256,
                label="Firecracker archive",
                environment_name="UVM_FIRECRACKER_SHA256",
            )
        verify_file(
            archive,
            self._settings.firecracker_sha256,
            "Firecracker archive",
            "UVM_FIRECRACKER_SHA256",
            allow_unverified=self._settings.allow_unverified_downloads,
        )

        try:
            with tarfile.open(archive, "r:gz") as tar:
                self._extract_archive_safely(tar, self._settings.base)
        except (OSError, tarfile.TarError) as error:
            raise UvmError(f"could not extract Firecracker archive {archive}: {error}") from error

        release_dir = self._settings.base / (
            f"release-{self._settings.firecracker_version}-{architecture}"
        )
        source_firecracker = self._find_release_binary(release_dir, "firecracker-")
        source_jailer = self._find_release_binary(release_dir, "jailer-")
        if source_firecracker is None or source_jailer is None:
            raise UvmError("Firecracker archive layout was not recognized")

        try:
            self._settings.bin_dir.mkdir(parents=True, exist_ok=True)
            shutil.copy2(source_firecracker, firecracker)
            firecracker.chmod(0o755)
            shutil.copy2(source_jailer, jailer)
            jailer.chmod(0o755)
        except OSError as error:
            raise UvmError(f"could not install Firecracker binaries: {error}") from error
        return firecracker

    def install_guest_assets(self, *, force: bool = False) -> GuestAssets:
        kernel = self._settings.kernel_image
        rootfs = self._settings.rootfs_image
        if force or not kernel.exists():
            self.download(
                self._settings.kernel_url,
                kernel,
                expected_sha256=self._settings.kernel_sha256,
                label="guest kernel",
                environment_name="UVM_KERNEL_SHA256",
            )
        if force or not rootfs.exists():
            self.download(
                self._settings.rootfs_url,
                rootfs,
                expected_sha256=self._settings.rootfs_sha256,
                label="guest root filesystem",
                environment_name="UVM_ROOTFS_SHA256",
            )
        verify_file(
            kernel,
            self._settings.kernel_sha256,
            "guest kernel",
            "UVM_KERNEL_SHA256",
            allow_unverified=self._settings.allow_unverified_downloads,
        )
        verify_file(
            rootfs,
            self._settings.rootfs_sha256,
            "guest root filesystem",
            "UVM_ROOTFS_SHA256",
            allow_unverified=self._settings.allow_unverified_downloads,
        )
        return GuestAssets(kernel=kernel, rootfs=rootfs)

    def download(
        self,
        url: str,
        destination: Path,
        *,
        expected_sha256: str | None,
        label: str,
        environment_name: str,
    ) -> None:
        """Download to a sibling temporary file before atomically publishing it."""

        if self._emit is not None:
            self._emit(f"Downloading {url}")
        require_checksum(
            expected_sha256,
            label,
            environment_name,
            allow_unverified=self._settings.allow_unverified_downloads,
        )
        destination.parent.mkdir(parents=True, exist_ok=True)
        descriptor, temporary_name = tempfile.mkstemp(
            prefix=f".{destination.name}.",
            suffix=".tmp",
            dir=destination.parent,
        )
        try:
            with os.fdopen(descriptor, "wb") as temporary_file:
                try:
                    with urllib.request.urlopen(url, timeout=60) as response:
                        shutil.copyfileobj(response, temporary_file)
                except (OSError, urllib.error.URLError) as error:
                    raise UvmError(f"could not download {url}: {error}") from error
                temporary_file.flush()
                os.fsync(temporary_file.fileno())
            verify_file(
                Path(temporary_name),
                expected_sha256,
                label,
                environment_name,
                allow_unverified=self._settings.allow_unverified_downloads,
            )
            os.replace(temporary_name, destination)
        finally:
            try:
                os.unlink(temporary_name)
            except FileNotFoundError:
                pass

    @staticmethod
    def _extract_archive_safely(tar: tarfile.TarFile, destination: Path) -> None:
        root = destination.resolve()
        for member in tar.getmembers():
            target = (destination / member.name).resolve()
            if target != root and root not in target.parents:
                raise UvmError("Firecracker archive contains an unsafe path")
        tar.extractall(destination, filter="data")

    @staticmethod
    def _find_release_binary(release_dir: Path, prefix: str) -> Path | None:
        if not release_dir.exists():
            return None
        candidates = sorted(
            path
            for path in release_dir.rglob(f"{prefix}*")
            if path.is_file() and path.name.startswith(prefix)
        )
        return candidates[0] if candidates else None