"""Guest asset lookup and per-VM writable root filesystem handling."""

from __future__ import annotations

import re
import shutil
import tempfile
from dataclasses import dataclass
from pathlib import Path

from .config import Settings
from .errors import UvmError
from .integrity import load_manifest, verify_file
from .system import CommandRunner


_FIRECRACKER_DEMO_PUBLIC_KEY = (
    "ssh-rsa "
    "AAAAB3NzaC1yc2EAAAADAQABAAABAQCirWKrc1zDyvZufHGinIRNoeIot+C3idANxtqZyDHL9mYm"
    "NeQzx9CjbjMgSDJ3xhIPP9mu3MP4Py/u3X5Wey98zN3EPboKdGRf6T2fFviK4i0q85LueDtsK0"
    "IoWR459w87tC9NMwPb27C8jPqFod6nWfccdhEdM+veKkFh4Dk5TrPfYHDayXDPFEdz7jl0GedEH"
    "fP9w11LPfa66D7731CdD3tMHAWLYxYmeXo58RXUaP6AgUK8uF/hL+E21q+wgTNPOQuRn2ekOjdu"
    "J34oHkJ2i48tLqKdGKU6RwfFc3rW3TkeYTUqi3UY9EnwNWFJicez+nZ5bhr5KvRsfSZ2QvBj"
)
_MODE_PATTERN = re.compile(r"Mode:\s+0*([0-7]+)")
_OWNER_PATTERN = re.compile(r"User:\s+(\d+)\s+Group:\s+(\d+)")
_HOST_KEY_TYPES = (
    ("rsa", "3072"),
    ("ecdsa", "256"),
    ("ed25519", None),
)


@dataclass(frozen=True, slots=True)
class GuestAssets:
    kernel: Path
    rootfs: Path


@dataclass(frozen=True, slots=True)
class _GuestFile:
    path: Path
    mode: int
    uid: int
    gid: int


class ImageStore:
    """Treat downloaded guest assets as templates, never as writable VM disks."""

    def __init__(self, settings: Settings, runner: CommandRunner | None = None) -> None:
        self._settings = settings
        self._runner = runner or CommandRunner()

    def installed_assets(self) -> GuestAssets:
        kernel = self._settings.kernel_image
        rootfs = self._settings.rootfs_image
        if not kernel.exists() or not rootfs.exists():
            raise UvmError("guest assets missing. Run: sudo uvm install")
        manifest = load_manifest(self._settings.integrity_manifest_path)
        kernel_checksum = self._settings.kernel_sha256
        rootfs_checksum = self._settings.rootfs_sha256
        if manifest.verified or self._settings.allow_unverified_downloads:
            kernel_checksum = kernel_checksum or manifest.checksums.get("kernel")
            rootfs_checksum = rootfs_checksum or manifest.checksums.get("rootfs")
        verify_file(
            kernel,
            kernel_checksum,
            "guest kernel",
            "UVM_KERNEL_SHA256",
            allow_unverified=self._settings.allow_unverified_downloads,
        )
        verify_file(
            rootfs,
            rootfs_checksum,
            "guest root filesystem",
            "UVM_ROOTFS_SHA256",
            allow_unverified=self._settings.allow_unverified_downloads,
        )
        return GuestAssets(kernel=kernel, rootfs=rootfs)

    def create_vm_disk(self, source: Path, destination: Path) -> Path:
        """Copy the template before Firecracker opens it read-write for a VM."""

        if destination.exists():
            raise UvmError(f"VM disk already exists: {destination}")
        destination.parent.mkdir(parents=True, exist_ok=True)
        try:
            shutil.copy2(source, destination)
            destination.chmod(0o600)
        except OSError as error:
            try:
                destination.unlink()
            except FileNotFoundError:
                pass
            raise UvmError(f"could not create VM disk {destination}: {error}") from error
        return destination

    def provision_credentials(self, disk: Path, username: str, password: str) -> None:
        """Set an existing guest account password in an offline ext4 disk."""

        password_hash = self._password_hash(password)
        try:
            with tempfile.TemporaryDirectory(prefix="uvm-credentials-") as temporary_name:
                temporary = Path(temporary_name)
                passwd = self._read_guest_file(disk, "/etc/passwd", temporary / "passwd")
                shadow = self._read_guest_file(disk, "/etc/shadow", temporary / "shadow")
                sshd_config = self._read_guest_file(
                    disk,
                    "/etc/ssh/sshd_config",
                    temporary / "sshd_config",
                )
                authorized_keys = self._read_guest_file(
                    disk,
                    "/root/.ssh/authorized_keys",
                    temporary / "authorized_keys",
                    required=False,
                )
                assert passwd is not None
                assert shadow is not None
                assert sshd_config is not None

                account_names = {
                    line.split(":", 1)[0]
                    for line in passwd.path.read_text(encoding="utf-8").splitlines()
                    if ":" in line
                }
                if username not in account_names:
                    raise UvmError(f"guest user does not exist in rootfs: {username}")

                shadow.path.write_text(
                    _set_shadow_password(
                        shadow.path.read_text(encoding="utf-8"),
                        username,
                        password_hash,
                    ),
                    encoding="utf-8",
                )
                shadow.path.chmod(shadow.mode)

                sshd_config.path.write_text(
                    _enable_ssh_password_authentication(
                        sshd_config.path.read_text(encoding="utf-8"),
                        username,
                    ),
                    encoding="utf-8",
                )
                sshd_config.path.chmod(sshd_config.mode)

                self._write_guest_file(disk, "/etc/shadow", shadow)
                self._write_guest_file(disk, "/etc/ssh/sshd_config", sshd_config)
                self._replace_guest_host_keys(disk, temporary)
                if authorized_keys is not None:
                    self._remove_firecracker_demo_key(disk, authorized_keys)
        except (OSError, UnicodeError) as error:
            raise UvmError(f"could not provision guest credentials: {error}") from error

    def _password_hash(self, password: str) -> str:
        result = self._runner.run(
            ("openssl", "passwd", "-6", "-stdin"),
            capture=True,
            input_text=f"{password}\n",
            sensitive=True,
        )
        password_hash = result.stdout.strip()
        if not password_hash.startswith("$6$") or any(
            character in password_hash for character in ("\n", "\r", ":")
        ):
            raise UvmError("openssl returned an invalid guest password hash")
        return password_hash

    def _read_guest_file(
        self,
        disk: Path,
        guest_path: str,
        destination: Path,
        *,
        required: bool = True,
    ) -> _GuestFile | None:
        stat_result = self._runner.run(
            ("debugfs", "-R", f"stat {guest_path}", disk),
            check=False,
            capture=True,
        )
        mode_match = _MODE_PATTERN.search(stat_result.stdout)
        owner_match = _OWNER_PATTERN.search(stat_result.stdout)
        if mode_match is None or owner_match is None:
            output = f"{stat_result.stdout}\n{stat_result.stderr}".lower()
            if not required and "file not found" in output:
                return None
            if required and "file not found" in output:
                raise UvmError(f"guest rootfs is missing required file: {guest_path}")
            raise UvmError(f"could not inspect guest rootfs file: {guest_path}")

        self._runner.run(
            ("debugfs", "-R", f"dump {guest_path} {destination}", disk),
            capture=True,
        )
        if not destination.is_file():
            raise UvmError(f"could not read guest rootfs file: {guest_path}")
        return _GuestFile(
            path=destination,
            mode=int(mode_match.group(1), 8),
            uid=int(owner_match.group(1)),
            gid=int(owner_match.group(2)),
        )

    def _write_guest_file(self, disk: Path, guest_path: str, source: _GuestFile) -> None:
        self._runner.run(("debugfs", "-w", "-R", f"rm {guest_path}", disk), capture=True)
        self._runner.run(
            ("debugfs", "-w", "-R", f"write {source.path} {guest_path}", disk),
            capture=True,
        )
        for field, value in (
            ("mode", f"0{0o100000 | source.mode:o}"),
            ("uid", str(source.uid)),
            ("gid", str(source.gid)),
        ):
            self._runner.run(
                (
                    "debugfs",
                    "-w",
                    "-R",
                    f"set_inode_field {guest_path} {field} {value}",
                    disk,
                ),
                capture=True,
            )

        verification = source.path.with_name(f"{source.path.name}.verify")
        written = self._read_guest_file(disk, guest_path, verification)
        assert written is not None
        if (
            verification.read_bytes() != source.path.read_bytes()
            or written.mode != source.mode
            or written.uid != source.uid
            or written.gid != source.gid
        ):
            raise UvmError(f"could not verify updated guest rootfs file: {guest_path}")

    def _replace_guest_host_keys(self, disk: Path, temporary: Path) -> None:
        for key_type, bits in _HOST_KEY_TYPES:
            name = f"ssh_host_{key_type}_key"
            private_path = temporary / name
            command: list[str | Path] = [
                "ssh-keygen",
                "-q",
                "-t",
                key_type,
                "-N",
                "",
                "-C",
                "",
                "-f",
                private_path,
            ]
            if bits is not None:
                command[4:4] = ["-b", bits]
            self._runner.run(command, capture=True)

            for suffix, default_mode in (("", 0o600), (".pub", 0o644)):
                guest_path = f"/etc/ssh/{name}{suffix}"
                existing = self._read_guest_file(
                    disk,
                    guest_path,
                    temporary / f"existing-{name}{suffix}",
                    required=False,
                )
                generated = private_path.with_name(f"{name}{suffix}")
                source = _GuestFile(
                    path=generated,
                    mode=existing.mode if existing is not None else default_mode,
                    uid=existing.uid if existing is not None else 0,
                    gid=existing.gid if existing is not None else 0,
                )
                self._write_guest_file(disk, guest_path, source)

    def _remove_firecracker_demo_key(self, disk: Path, authorized_keys: _GuestFile) -> None:
        contents = authorized_keys.path.read_text(encoding="utf-8")
        updated = _without_firecracker_demo_key(contents)
        if updated == contents:
            return
        if updated is not None:
            authorized_keys.path.write_text(updated, encoding="utf-8")
            authorized_keys.path.chmod(authorized_keys.mode)
            self._write_guest_file(
                disk,
                "/root/.ssh/authorized_keys",
                authorized_keys,
            )
            return

        self._runner.run(
            ("debugfs", "-w", "-R", "rm /root/.ssh/authorized_keys", disk),
            capture=True,
        )
        if self._read_guest_file(
            disk,
            "/root/.ssh/authorized_keys",
            authorized_keys.path.with_name("authorized_keys.verify"),
            required=False,
        ) is not None:
            raise UvmError("could not remove the insecure Firecracker demo SSH key")


def _set_shadow_password(contents: str, username: str, password_hash: str) -> str:
    lines = contents.splitlines()
    for index, line in enumerate(lines):
        fields = line.split(":")
        if fields[0] == username and len(fields) >= 2:
            fields[1] = password_hash
            lines[index] = ":".join(fields)
            return "\n".join(lines) + "\n"
    raise UvmError(f"guest rootfs has no shadow entry for user: {username}")


def _enable_ssh_password_authentication(contents: str, username: str) -> str:
    managed_directives = {"hostkey", "passwordauthentication", "permitrootlogin"}
    unmanaged_lines = []
    for line in contents.splitlines():
        stripped = line.lstrip()
        directive = stripped.split(None, 1)[0].lower() if stripped else ""
        if stripped.startswith("#") or directive not in managed_directives:
            unmanaged_lines.append(line)

    directives = [
        "HostKey /etc/ssh/ssh_host_rsa_key",
        "HostKey /etc/ssh/ssh_host_ecdsa_key",
        "HostKey /etc/ssh/ssh_host_ed25519_key",
        "PasswordAuthentication no",
    ]
    if username == "root":
        directives.append("PermitRootLogin yes")
    return (
        "# Managed by uvm\n"
        + "\n".join(directives)
        + f"\nMatch User {username}\n    PasswordAuthentication yes\nMatch all\n"
        + "\n".join(unmanaged_lines).rstrip()
        + "\n"
    )


def _without_firecracker_demo_key(contents: str) -> str | None:
    lines = contents.splitlines()
    retained = [
        line
        for line in lines
        if " ".join(line.split()[:2]) != _FIRECRACKER_DEMO_PUBLIC_KEY
    ]
    if len(retained) == len(lines):
        return contents
    return "\n".join(retained) + "\n" if retained else None