"""Internal qemu-nbd + read-only loop-mount machinery for FN-A3.

Production default for :func:`mgf.vm.forensics.capture.filesystem._qemu_nbd_mount`.
Tests don't go through this module — they inject a fake ``mount_disk_fn``
that yields a pre-populated ``tmp_path``. This module is exercised by
its own unit tests via injectable ``run_fn`` / ``probe_fn`` / ``sleep_fn``
seams; the real qemu-nbd binary is never spawned in CI.

Sudo policy
-----------

``qemu-nbd --connect``, ``modprobe nbd``, ``mount``, and ``umount`` all
need root. The smoke flow assumes the analyst either runs as root or
has passwordless sudo for these binaries. We always prefix the argv
with ``sudo -n`` so a missing sudoers entry surfaces as a clean
``CommandError`` from :func:`run_command` rather than an interactive
prompt blocking the CLI. Pass ``sudo=False`` from tests / from a root
process to skip the prefix.

Cleanup
-------

The :func:`mount_disk` context manager always unmounts every mounted
partition + disconnects the NBD device on exit, even when the body
raises. Cleanup errors are logged at WARNING and swallowed — the
in-flight exception (if any) wins.
"""

from __future__ import annotations

import contextlib
import logging
import re
import tempfile
import time
from pathlib import Path
from typing import TYPE_CHECKING

from mgf.vm.core.common.shell import CommandResult, run_command
from mgf.vm.core.exceptions import CommandError, OperationError

if TYPE_CHECKING:
    from collections.abc import Callable, Iterator


__all__ = [
    "MAX_NBD_DEVICES",
    "PARTITION_SETTLE_SECONDS",
    "connect_disk",
    "disconnect_disk",
    "discover_partitions",
    "ensure_nbd_module",
    "find_free_nbd_device",
    "mount_disk",
    "mount_readonly",
    "umount",
]


LOG = logging.getLogger("mgf.vm.forensics.capture._nbd")

# Kernel default — the `nbd` module exposes /dev/nbd0 through /dev/nbd15
# unless loaded with a higher `nbds_max=`.
MAX_NBD_DEVICES = 16

# How long to wait after qemu-nbd --connect for the kernel to publish
# /dev/nbdNp* partition device nodes. partprobe is async via udev.
PARTITION_SETTLE_SECONDS = 5.0
PARTITION_POLL_INTERVAL = 0.1

# Partition types we never try to mount — listed by blkid TYPE= value.
SKIP_FS_TYPES = frozenset({"swap", "linux_raid_member", "LVM2_member"})

RunFn = "Callable[..., CommandResult]"


# ---------------------------------------------------------------------------
# Low-level operations (each independently injectable for tests)
# ---------------------------------------------------------------------------


def ensure_nbd_module(
    *,
    sudo: bool = True,
    run_fn: Callable[..., CommandResult] = run_command,
    probe_fn: Callable[[Path], bool] = Path.exists,
) -> None:
    """Load the ``nbd`` kernel module if /dev/nbd0 is not present.

    Idempotent. ``modprobe nbd max_part=8`` so the kernel auto-probes
    partition tables on connect — without this the per-partition
    devices ``/dev/nbdNp1`` etc. never appear.
    """
    if probe_fn(Path("/dev/nbd0")):
        return
    argv = _maybe_sudo(["modprobe", "nbd", "max_part=8"], sudo=sudo)
    try:
        run_fn(argv)
    except CommandError as e:
        msg = (
            "could not load the `nbd` kernel module "
            f"(needed for qemu-nbd --connect): {e}. "
            "Either run as root, or configure passwordless sudo for `modprobe nbd`."
        )
        raise OperationError(msg) from e
    if not probe_fn(Path("/dev/nbd0")):
        msg = (
            "modprobe succeeded but /dev/nbd0 still does not exist. "
            "Your kernel may not ship the `nbd` module."
        )
        raise OperationError(msg)


def find_free_nbd_device(
    *,
    probe_fn: Callable[[Path], bool] = Path.exists,
    pid_dir: Path = Path("/sys/block"),
    max_devices: int = MAX_NBD_DEVICES,
) -> Path:
    """Return the first ``/dev/nbdN`` whose ``/sys/block/nbdN/pid`` is empty.

    A non-empty ``pid`` file means qemu-nbd has already connected to
    that device. We scan from 0 upward and return the first free one.
    """
    for i in range(max_devices):
        dev = Path(f"/dev/nbd{i}")
        if not probe_fn(dev):
            continue
        pid_file = pid_dir / f"nbd{i}" / "pid"
        # If the pid sysfs file does not exist, the device is unused.
        if not probe_fn(pid_file):
            return dev
    msg = (
        f"all {max_devices} /dev/nbd* devices are in use. "
        "Disconnect a stale capture, or boot with `nbd nbds_max=N` for more."
    )
    raise OperationError(msg)


def connect_disk(
    disk_path: Path,
    nbd_dev: Path,
    *,
    sudo: bool = True,
    read_only: bool = True,
    run_fn: Callable[..., CommandResult] = run_command,
) -> None:
    """``qemu-nbd --connect=<nbd_dev> <disk>`` (read-only by default)."""
    argv = ["qemu-nbd"]
    if read_only:
        argv.append("--read-only")
    argv.extend([f"--connect={nbd_dev}", str(disk_path)])
    try:
        run_fn(_maybe_sudo(argv, sudo=sudo))
    except CommandError as e:
        msg = (
            f"qemu-nbd --connect={nbd_dev} {disk_path} failed: {e}. "
            "Is qemu-nbd installed (`sudo apt install qemu-utils`)?"
        )
        raise OperationError(msg) from e


def disconnect_disk(
    nbd_dev: Path,
    *,
    sudo: bool = True,
    run_fn: Callable[..., CommandResult] = run_command,
) -> None:
    """``qemu-nbd --disconnect <nbd_dev>``. Best-effort — logs on failure."""
    argv = _maybe_sudo(["qemu-nbd", "--disconnect", str(nbd_dev)], sudo=sudo)
    try:
        run_fn(argv)
    except CommandError as e:
        LOG.warning(
            "qemu-nbd --disconnect failed",
            extra={"nbd_dev": str(nbd_dev), "error": str(e)},
        )


def discover_partitions(
    nbd_dev: Path,
    *,
    probe_fn: Callable[[Path], bool] = Path.exists,
    sleep_fn: Callable[[float], None] = time.sleep,
    settle_seconds: float = PARTITION_SETTLE_SECONDS,
    poll_interval: float = PARTITION_POLL_INTERVAL,
) -> tuple[Path, ...]:
    """Return mountable block-device paths under ``nbd_dev``.

    For a partitioned disk: ``(/dev/nbdNp1, /dev/nbdNp2, ...)``. Polls
    for up to ``settle_seconds`` because the kernel publishes partition
    devices via udev *after* the qemu-nbd connect returns.

    For a disk with no partition table (raw filesystem): ``(/dev/nbdN,)``.
    """
    deadline = settle_seconds
    elapsed = 0.0
    parts: list[Path] = []
    while elapsed <= deadline:
        parts = _scan_partitions(nbd_dev, probe_fn=probe_fn)
        if parts:
            return tuple(parts)
        sleep_fn(poll_interval)
        elapsed += poll_interval
    # No partitions — assume the disk holds a raw filesystem (no MBR/GPT).
    return (nbd_dev,)


def _scan_partitions(
    nbd_dev: Path,
    *,
    probe_fn: Callable[[Path], bool],
) -> list[Path]:
    parts: list[Path] = []
    base = nbd_dev.name  # "nbd3"
    parent = nbd_dev.parent  # "/dev"
    # Try p1..p15 — same upper bound as the kernel default `max_part=8`
    # but harmless to overscan because we just probe existence.
    for i in range(1, 16):
        candidate = parent / f"{base}p{i}"
        if probe_fn(candidate):
            parts.append(candidate)
    return parts


def mount_readonly(
    part_dev: Path,
    mount_dir: Path,
    *,
    sudo: bool = True,
    run_fn: Callable[..., CommandResult] = run_command,
) -> bool:
    """``mount -o ro,noload <part> <mount_dir>``. Returns False on a "skip me" filesystem.

    Tries ``ro,noload`` first (correct for ext2/3/4 — skips journal
    replay which would write to the disk). Falls back to plain ``ro``
    on filesystems that reject ``noload`` (xfs, btrfs).

    Returns False if the partition has a filesystem we never mount
    (swap, raid metadata, lvm pv); True on a successful mount.
    """
    fstype = _probe_fstype(part_dev, sudo=sudo, run_fn=run_fn)
    if fstype in SKIP_FS_TYPES:
        LOG.info(
            "skipping unmountable partition",
            extra={"part_dev": str(part_dev), "fstype": fstype},
        )
        return False

    mount_dir.mkdir(parents=True, exist_ok=True)
    for opts in ("ro,noload", "ro"):
        argv = _maybe_sudo(
            ["mount", "-o", opts, str(part_dev), str(mount_dir)],
            sudo=sudo,
        )
        try:
            run_fn(argv)
            return True
        except CommandError as e:
            # If `noload` was rejected, fall through and retry plain `ro`.
            if opts == "ro,noload" and "noload" in str(e):
                continue
            # Any other failure on the second attempt — give up on this partition.
            LOG.warning(
                "mount failed; skipping partition",
                extra={"part_dev": str(part_dev), "opts": opts, "error": str(e)},
            )
            return False
    return False


def umount(
    mount_dir: Path,
    *,
    sudo: bool = True,
    run_fn: Callable[..., CommandResult] = run_command,
) -> None:
    """``umount <mount_dir>``. Best-effort — logs on failure."""
    argv = _maybe_sudo(["umount", str(mount_dir)], sudo=sudo)
    try:
        run_fn(argv)
    except CommandError as e:
        LOG.warning(
            "umount failed",
            extra={"mount_dir": str(mount_dir), "error": str(e)},
        )


# ---------------------------------------------------------------------------
# High-level context manager — what filesystem.py wires in by default
# ---------------------------------------------------------------------------


@contextlib.contextmanager
def mount_disk(
    disk_path: Path,
    *,
    sudo: bool = True,
    run_fn: Callable[..., CommandResult] = run_command,
    probe_fn: Callable[[Path], bool] = Path.exists,
    sleep_fn: Callable[[float], None] = time.sleep,
    find_free_nbd_fn: Callable[[], Path] | None = None,
) -> Iterator[Path]:
    """Read-only mount of a qcow2 / raw disk image via qemu-nbd.

    Yields the path to a temp directory holding every successfully
    mounted partition. For partitioned disks: ``<tmpdir>/p1``,
    ``<tmpdir>/p2``, … For raw-filesystem disks: ``<tmpdir>/p0`` (one
    entry, mounted at the disk itself).

    On exit, every partition is unmounted and the NBD device is
    disconnected. Errors during cleanup are logged at WARNING and
    swallowed.
    """
    ensure_nbd_module(sudo=sudo, run_fn=run_fn, probe_fn=probe_fn)
    if find_free_nbd_fn is None:
        nbd_dev = find_free_nbd_device(probe_fn=probe_fn)
    else:
        nbd_dev = find_free_nbd_fn()

    connect_disk(disk_path, nbd_dev, sudo=sudo, run_fn=run_fn)

    mounted: list[Path] = []
    tmp_root = Path(tempfile.mkdtemp(prefix="vmanager-fs-"))
    try:
        partitions = discover_partitions(
            nbd_dev,
            probe_fn=probe_fn,
            sleep_fn=sleep_fn,
        )
        for idx, part in enumerate(partitions, start=1):
            sub = tmp_root / f"p{idx}"
            if mount_readonly(part, sub, sudo=sudo, run_fn=run_fn):
                mounted.append(sub)
        yield tmp_root
    finally:
        for sub in mounted:
            umount(sub, sudo=sudo, run_fn=run_fn)
        disconnect_disk(nbd_dev, sudo=sudo, run_fn=run_fn)
        # Clean tmp dir even if some umount failed — leftover mountpoints
        # would just be empty subdirs since umount lazily detaches on
        # busy. Best-effort: if anything is genuinely still mounted we
        # leave it for manual cleanup rather than rmtree-ing through it.
        with contextlib.suppress(OSError):
            for sub in tmp_root.iterdir():
                with contextlib.suppress(OSError):
                    sub.rmdir()
            tmp_root.rmdir()


# ---------------------------------------------------------------------------
# Helpers
# ---------------------------------------------------------------------------


_BLKID_TYPE_RE = re.compile(r'\bTYPE="([^"]+)"')


def _probe_fstype(
    part_dev: Path,
    *,
    sudo: bool,
    run_fn: Callable[..., CommandResult],
) -> str:
    """Return the filesystem type via ``blkid``, or "" on any failure."""
    argv = _maybe_sudo(["blkid", str(part_dev)], sudo=sudo)
    try:
        result = run_fn(argv)
    except CommandError:
        return ""
    match = _BLKID_TYPE_RE.search(result.stdout)
    return match.group(1) if match else ""


def _maybe_sudo(argv: list[str], *, sudo: bool) -> list[str]:
    """Prefix ``argv`` with ``sudo -n`` if ``sudo=True``.

    ``-n`` makes sudo non-interactive — a missing sudoers entry surfaces
    as a clean ``CommandError`` rather than blocking the CLI on a TTY
    password prompt.
    """
    if not sudo:
        return list(argv)
    return ["sudo", "-n", *argv]
