"""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 Settings
from .errors import UvmError
from .images import GuestAssets
from .integrity import load_manifest, require_checksum, sha256_file, verify_file, write_manifest
from .state import StateStore
from .system import CommandRunner, check_kvm, ensure_data_directories, require_root


_UBUNTU_RELEASE = "jammy"
_ROOTFS_RECIPE = "uvm-ubuntu-jammy-v1\n"
_GUEST_PACKAGES = (
    "ca-certificates", "openssh-server", "sudo", "systemd-sysv", "udev",
    "dbus", "libpam-systemd", "iproute2", "iputils-ping", "curl",
)
# Keep the image mountable by the older Firecracker quickstart guest kernels.
_EXT4_FEATURES = (
    "none,has_journal,ext_attr,resize_inode,dir_index,filetype,extent,64bit,"
    "flex_bg,sparse_super,large_file,huge_file,dir_nlink,extra_isize,metadata_csum"
)


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_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_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,
        )
        if self._settings.rootfs_url:
            require_checksum(
                self._settings.rootfs_sha256,
                "guest root filesystem",
                "UVM_ROOTFS_SHA256",
                allow_unverified=self._settings.allow_unverified_downloads,
            )
        elif self._settings.rootfs_sha256:
            raise UvmError(
                "UVM_ROOTFS_SHA256 requires UVM_ROOTFS_URL for a prebuilt image; "
                "unset UVM_ROOTFS_SHA256 to build the default Ubuntu rootfs"
            )

    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 if self._settings.rootfs_url else True,
                )
            ),
        )

    def install_apt_packages(self) -> None:
        packages = (
            "curl",
            "jq",
            "iproute2",
            "iptables",
            "nftables",
            "conntrack",
            "coreutils",
            "e2fsprogs",
            "util-linux",
            "ca-certificates",
            "openssl",
            "openssh-client",
        )
        if not self._settings.rootfs_url:
            packages += ("debootstrap", "ubuntu-keyring")
        self._runner.run(("apt-get", "update"))
        self._runner.run(("apt-get", "install", "-y", *packages))

    def install_firecracker_binary(self, *, force: bool = False) -> Path:
        architecture = self._settings.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 not self._settings.rootfs_url:
            self._install_ubuntu_rootfs(rootfs, force=force)
        elif 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,
        )
        if self._settings.rootfs_url:
            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 _install_ubuntu_rootfs(self, destination: Path, *, force: bool) -> None:
        if destination.is_file() and not force:
            manifest = load_manifest(self._settings.integrity_manifest_path)
            checksum = manifest.checksums.get("rootfs")
            if manifest.verified and checksum:
                verify_file(destination, checksum, "Ubuntu rootfs", "integrity manifest", allow_unverified=False)
                marker = self._runner.run(
                    ("debugfs", "-R", "cat /etc/uvm-image", destination), capture=True,
                )
                if marker.stdout == _ROOTFS_RECIPE:
                    return
        # Rebuild old demo images and unverified caches rather than promoting their trust.
        self._build_ubuntu_rootfs(destination)

    def _build_ubuntu_rootfs(self, destination: Path) -> None:
        architecture = "arm64" if self._settings.architecture == "aarch64" else "amd64"
        mirror = (
            "https://ports.ubuntu.com/ubuntu-ports" if architecture == "arm64"
            else "https://archive.ubuntu.com/ubuntu"
        )
        if self._emit is not None:
            self._emit("Building Ubuntu 22.04 rootfs with sudo, OpenSSH, and apt/dpkg metadata")
        destination.parent.mkdir(parents=True, exist_ok=True)
        try:
            with tempfile.TemporaryDirectory(prefix=".ubuntu-build-", dir=destination.parent) as name:
                temporary = Path(name)
                root = temporary / "root"
                policy = root / "usr/sbin/policy-rc.d"
                policy.parent.mkdir(parents=True)
                policy.write_text("#!/bin/sh\nexit 101\n", encoding="utf-8")
                policy.chmod(0o755)
                # Package scripts run natively, with private mounts/PIDs and service startup disabled.
                self._runner.run((
                    "unshare", "--mount", "--propagation", "private", "--pid", "--fork",
                    "--kill-child", "--mount-proc", "env", "DEBIAN_FRONTEND=noninteractive",
                    "debootstrap", f"--arch={architecture}", "--variant=minbase", "--merged-usr",
                    "--components=main,universe", f"--include={','.join(_GUEST_PACKAGES)}",
                    f"--extra-suites={_UBUNTU_RELEASE}-updates,{_UBUNTU_RELEASE}-security",
                    "--force-check-gpg", "--keyring=/usr/share/keyrings/ubuntu-archive-keyring.gpg",
                    _UBUNTU_RELEASE, root, mirror,
                ))
                policy.unlink()
                self._configure_ubuntu_rootfs(root, mirror)
                self._validate_ubuntu_rootfs(root)

                disk = temporary / "rootfs.ext4"
                with disk.open("wb") as image:
                    image.truncate(2 * 1024 * 1024 * 1024)
                disk.chmod(0o600)
                self._runner.run((
                    "mkfs.ext4", "-q", "-F", "-b", "4096", "-O", _EXT4_FEATURES,
                    "-E", "lazy_itable_init=0,lazy_journal_init=0", "-d", root, disk,
                ), capture=True)
                self._runner.run(("e2fsck", "-f", "-n", disk), capture=True)
                with disk.open("rb") as image:
                    os.fsync(image.fileno())
                os.replace(disk, destination)
                directory = os.open(destination.parent, os.O_RDONLY)
                try:
                    os.fsync(directory)
                finally:
                    os.close(directory)
        except (OSError, UnicodeError) as error:
            raise UvmError(f"could not build Ubuntu guest rootfs: {error}") from error

    def _configure_ubuntu_rootfs(self, root: Path, mirror: str) -> None:
        files = {
            "etc/apt/sources.list": "".join(
                f"deb {mirror} {suite} main universe\n"
                for suite in (_UBUNTU_RELEASE, f"{_UBUNTU_RELEASE}-updates", f"{_UBUNTU_RELEASE}-security")
            ),
            "etc/fstab": "/dev/vda / ext4 defaults 0 1\n",
            "etc/hostname": "uvm\n",
            "etc/hosts": "127.0.0.1 localhost\n127.0.1.1 uvm\n",
            "etc/resolv.conf": "nameserver 1.1.1.1\nnameserver 8.8.8.8\n",
            "etc/machine-id": "",
            "etc/uvm-image": _ROOTFS_RECIPE,
            "etc/systemd/network/10-uvm.network": (
                "[Match]\nName=eth0\n\n[Network]\n"
                "KeepConfiguration=static\nDHCP=no\nLinkLocalAddressing=no\nIPv6AcceptRA=no\n"
            ),
        }
        for relative, contents in files.items():
            path = root / relative
            path.parent.mkdir(parents=True, exist_ok=True)
            path.unlink(missing_ok=True)  # In particular, replace a resolver symlink with guest DNS.
            path.write_text(contents, encoding="utf-8")
            path.chmod(0o644)
        machine_id = root / "var/lib/dbus/machine-id"
        machine_id.parent.mkdir(parents=True, exist_ok=True)
        machine_id.unlink(missing_ok=True)
        machine_id.symlink_to("../../../etc/machine-id")
        for key in (root / "etc/ssh").glob("ssh_host_*_key*"):
            key.unlink()
        self._runner.run(("systemctl", "--root", root, "enable", "ssh.service", "systemd-networkd.service"), capture=True)
        self._runner.run(("systemctl", "--root", root, "set-default", "multi-user.target"), capture=True)

    @staticmethod
    def _validate_ubuntu_rootfs(root: Path) -> None:
        for relative in (
            "etc/passwd", "etc/shadow", "etc/group", "etc/sudoers", "etc/ssh/sshd_config",
            "usr/bin/sudo", "usr/bin/apt-get", "usr/sbin/sshd", "usr/lib/systemd/systemd",
            "var/lib/dpkg/status", "var/lib/dpkg/info/apt.list", "var/lib/dpkg/info/sudo.list",
        ):
            if not (root / relative).is_file():
                raise UvmError(f"built Ubuntu rootfs is missing /{relative}")
        installed = set()
        for paragraph in (root / "var/lib/dpkg/status").read_text(encoding="utf-8").split("\n\n"):
            fields = dict(line.split(": ", 1) for line in paragraph.splitlines() if ": " in line)
            if fields.get("Status") == "install ok installed":
                installed.add(fields.get("Package"))
        if not {"apt", "sudo", "openssh-server", "systemd-sysv"}.issubset(installed):
            raise UvmError("built Ubuntu rootfs has incomplete apt/dpkg package metadata")
        if not (root / "usr/bin/sudo").stat().st_mode & 0o4000:
            raise UvmError("built Ubuntu rootfs sudo is missing its setuid permission")

    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