"""Tests for :mod:`mgf.vm.forensics.capture._nbd` (Phase 17, FN-Chunk-5a).

Every subprocess call is intercepted via the injectable ``run_fn``;
filesystem probes go through the injectable ``probe_fn``. The real
``qemu-nbd`` / ``mount`` / ``modprobe`` binaries are never spawned.
"""

from __future__ import annotations

from pathlib import Path
from typing import TYPE_CHECKING

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

from mgf.vm.forensics.capture._nbd import (
    MAX_NBD_DEVICES,
    connect_disk,
    disconnect_disk,
    discover_partitions,
    ensure_nbd_module,
    find_free_nbd_device,
    mount_disk,
    mount_readonly,
    umount,
)

if TYPE_CHECKING:
    from collections.abc import Callable


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


class _RunRecorder:
    """Records every argv passed to run_fn; returns a configurable result."""

    def __init__(
        self,
        *,
        stdout: str = "",
        stderr: str = "",
        returncode: int = 0,
        raise_on: Callable[[list[str]], CommandError | None] | None = None,
    ) -> None:
        self.calls: list[list[str]] = []
        self._stdout = stdout
        self._stderr = stderr
        self._returncode = returncode
        self._raise_on = raise_on

    def __call__(self, argv: list[str], **_kw: object) -> CommandResult:
        self.calls.append(list(argv))
        if self._raise_on is not None:
            err = self._raise_on(argv)
            if err is not None:
                raise err
        return CommandResult(
            argv=list(argv),
            returncode=self._returncode,
            stdout=self._stdout,
            stderr=self._stderr,
        )


def _cmd_error(stderr: str = "") -> CommandError:
    return CommandError(argv=["fake"], returncode=1, stderr=stderr)


def _ok_run(stdout: str = "") -> _RunRecorder:
    return _RunRecorder(stdout=stdout)


# ---------------------------------------------------------------------------
# ensure_nbd_module
# ---------------------------------------------------------------------------


class TestEnsureNbdModule:
    def test_no_op_when_already_loaded(self) -> None:
        run = _ok_run()

        ensure_nbd_module(run_fn=run, probe_fn=lambda p: p == Path("/dev/nbd0"))
        assert run.calls == []

    def test_modprobe_invoked_when_module_missing(self) -> None:
        run = _ok_run()
        # Probe says no /dev/nbd0 the first time, yes after modprobe.
        seen: list[Path] = []

        def probe(p: Path) -> bool:
            seen.append(p)
            return seen.count(p) > 1  # second probe sees the device

        ensure_nbd_module(run_fn=run, probe_fn=probe)
        assert run.calls == [["sudo", "-n", "modprobe", "nbd", "max_part=8"]]

    def test_skip_sudo_prefix_when_disabled(self) -> None:
        run = _ok_run()
        seen: list[Path] = []

        def probe(p: Path) -> bool:
            seen.append(p)
            return seen.count(p) > 1

        ensure_nbd_module(sudo=False, run_fn=run, probe_fn=probe)
        assert run.calls == [["modprobe", "nbd", "max_part=8"]]

    def test_modprobe_failure_raises_actionable_error(self) -> None:
        def fail(_: list[str]) -> CommandError:
            return _cmd_error("modprobe failed: not found")

        run = _RunRecorder(raise_on=fail)

        with pytest.raises(OperationError) as exc:
            ensure_nbd_module(run_fn=run, probe_fn=lambda _p: False)
        assert "passwordless sudo" in str(exc.value)
        assert "modprobe nbd" in str(exc.value)

    def test_post_modprobe_still_missing_raises(self) -> None:
        run = _ok_run()
        with pytest.raises(OperationError) as exc:
            ensure_nbd_module(run_fn=run, probe_fn=lambda _p: False)
        assert "/dev/nbd0 still does not exist" in str(exc.value)


# ---------------------------------------------------------------------------
# find_free_nbd_device
# ---------------------------------------------------------------------------


class TestFindFreeNbdDevice:
    def test_returns_first_free_device(self) -> None:
        # /dev/nbd0 + /dev/nbd1 exist, but only nbd0 has a `pid` (in use).
        # nbd1 has no pid file → free.
        existing = {
            Path("/dev/nbd0"),
            Path("/dev/nbd1"),
            Path("/sys/block/nbd0/pid"),  # nbd0 in use
        }

        dev = find_free_nbd_device(probe_fn=lambda p: p in existing)
        assert dev == Path("/dev/nbd1")

    def test_skips_nonexistent_devices(self) -> None:
        # Only /dev/nbd5 exists (kernel happens to not publish 0..4)
        existing = {Path("/dev/nbd5")}
        dev = find_free_nbd_device(probe_fn=lambda p: p in existing)
        assert dev == Path("/dev/nbd5")

    def test_all_busy_raises(self) -> None:
        # Every device has a pid file → all in use.
        def probe(p: Path) -> bool:
            return p.parts[1] == "dev" or p.name == "pid"

        with pytest.raises(OperationError) as exc:
            find_free_nbd_device(probe_fn=probe, max_devices=2)
        assert "in use" in str(exc.value)
        assert "nbds_max" in str(exc.value)


# ---------------------------------------------------------------------------
# connect_disk / disconnect_disk
# ---------------------------------------------------------------------------


class TestConnectDisk:
    def test_argv_shape_with_sudo(self, tmp_path: Path) -> None:
        run = _ok_run()
        disk = tmp_path / "guest.qcow2"
        disk.touch()

        connect_disk(disk, Path("/dev/nbd3"), run_fn=run)
        assert run.calls == [
            [
                "sudo",
                "-n",
                "qemu-nbd",
                "--read-only",
                "--connect=/dev/nbd3",
                str(disk),
            ],
        ]

    def test_read_only_can_be_disabled(self, tmp_path: Path) -> None:
        run = _ok_run()
        disk = tmp_path / "guest.qcow2"
        disk.touch()
        connect_disk(disk, Path("/dev/nbd1"), read_only=False, run_fn=run)
        # No --read-only flag.
        assert "--read-only" not in run.calls[0]

    def test_failure_wrapped_with_actionable_hint(self, tmp_path: Path) -> None:
        def fail(_: list[str]) -> CommandError:
            return _cmd_error("command not found: qemu-nbd")

        run = _RunRecorder(raise_on=fail)
        with pytest.raises(OperationError) as exc:
            connect_disk(tmp_path / "g.qcow2", Path("/dev/nbd0"), run_fn=run)
        assert "qemu-utils" in str(exc.value)


class TestDisconnectDisk:
    def test_argv_shape(self) -> None:
        run = _ok_run()
        disconnect_disk(Path("/dev/nbd2"), run_fn=run)
        assert run.calls == [
            ["sudo", "-n", "qemu-nbd", "--disconnect", "/dev/nbd2"],
        ]

    def test_failure_is_swallowed_and_logged(self, caplog: pytest.LogCaptureFixture) -> None:
        def fail(_: list[str]) -> CommandError:
            return _cmd_error("device busy")

        run = _RunRecorder(raise_on=fail)
        # Must NOT raise — cleanup is best-effort.
        disconnect_disk(Path("/dev/nbd0"), run_fn=run)
        assert any("disconnect failed" in r.message for r in caplog.records)


# ---------------------------------------------------------------------------
# discover_partitions
# ---------------------------------------------------------------------------


class TestDiscoverPartitions:
    def test_returns_partition_devices(self) -> None:
        existing = {
            Path("/dev/nbd2p1"),
            Path("/dev/nbd2p2"),
            Path("/dev/nbd2p3"),
        }
        sleep_calls: list[float] = []

        parts = discover_partitions(
            Path("/dev/nbd2"),
            probe_fn=lambda p: p in existing,
            sleep_fn=sleep_calls.append,
        )
        assert parts == (
            Path("/dev/nbd2p1"),
            Path("/dev/nbd2p2"),
            Path("/dev/nbd2p3"),
        )
        # No sleep — partitions visible on first scan.
        assert sleep_calls == []

    def test_polls_until_partitions_appear(self) -> None:
        # On the first 3 calls return no partitions; on the 4th return p1.
        scan_count = [0]

        def probe(p: Path) -> bool:
            if p == Path("/dev/nbd0p1"):
                scan_count[0] += 1
                return scan_count[0] >= 4
            return False

        sleep_calls: list[float] = []
        parts = discover_partitions(
            Path("/dev/nbd0"),
            probe_fn=probe,
            sleep_fn=sleep_calls.append,
            poll_interval=0.1,
            settle_seconds=2.0,
        )
        assert parts == (Path("/dev/nbd0p1"),)
        # Three sleeps before the fourth scan succeeded.
        assert len(sleep_calls) >= 3

    def test_no_partitions_returns_disk_itself(self) -> None:
        # Raw filesystem on the disk — no partition table, no /dev/nbdNp* nodes.
        sleep_calls: list[float] = []
        parts = discover_partitions(
            Path("/dev/nbd5"),
            probe_fn=lambda _p: False,
            sleep_fn=sleep_calls.append,
            settle_seconds=0.3,
            poll_interval=0.1,
        )
        assert parts == (Path("/dev/nbd5"),)


# ---------------------------------------------------------------------------
# mount_readonly
# ---------------------------------------------------------------------------


class TestMountReadonly:
    def test_ext4_uses_noload(self, tmp_path: Path) -> None:
        run = _RunRecorder(stdout='/dev/nbd0p1: TYPE="ext4"')
        mount_dir = tmp_path / "p1"

        ok = mount_readonly(Path("/dev/nbd0p1"), mount_dir, run_fn=run)
        assert ok is True
        # Two argv: blkid probe, then mount.
        assert any("blkid" in c for c in run.calls)
        mount_call = next(c for c in run.calls if "mount" in c)
        assert "ro,noload" in mount_call

    def test_xfs_falls_back_to_plain_ro(self, tmp_path: Path) -> None:
        # blkid says xfs; first mount with noload fails; second mount with ro succeeds.
        def behaviour(argv: list[str]) -> CommandError | None:
            if "mount" in argv and "ro,noload" in argv:
                return _cmd_error("mount: bad option `noload'")
            return None

        run = _RunRecorder(stdout='TYPE="xfs"', raise_on=behaviour)
        ok = mount_readonly(Path("/dev/nbd1p1"), tmp_path / "p1", run_fn=run)
        assert ok is True
        mount_calls = [c for c in run.calls if "mount" in c and "umount" not in c[0]]
        assert len(mount_calls) == 2
        assert "ro,noload" in mount_calls[0]
        assert "ro,noload" not in mount_calls[1]
        assert "ro" in mount_calls[1]

    def test_swap_partition_returns_false_without_mounting(self, tmp_path: Path) -> None:
        run = _RunRecorder(stdout='TYPE="swap"')
        ok = mount_readonly(Path("/dev/nbd0p2"), tmp_path / "p2", run_fn=run)
        assert ok is False
        # Only the blkid probe ran — no mount attempt.
        assert all("mount" not in c for c in run.calls)

    def test_unrecoverable_mount_failure_returns_false(self, tmp_path: Path) -> None:
        def always_fail(_: list[str]) -> CommandError:
            return _cmd_error("mount: garbage filesystem")

        run = _RunRecorder(stdout='TYPE="ext4"', raise_on=always_fail)
        ok = mount_readonly(Path("/dev/nbd9p1"), tmp_path / "p1", run_fn=run)
        assert ok is False


# ---------------------------------------------------------------------------
# umount
# ---------------------------------------------------------------------------


class TestUmount:
    def test_argv_shape(self, tmp_path: Path) -> None:
        run = _ok_run()
        umount(tmp_path / "m", run_fn=run)
        assert run.calls == [["sudo", "-n", "umount", str(tmp_path / "m")]]

    def test_failure_swallowed(self, tmp_path: Path, caplog: pytest.LogCaptureFixture) -> None:
        def fail(_: list[str]) -> CommandError:
            return _cmd_error("device busy")

        run = _RunRecorder(raise_on=fail)
        umount(tmp_path / "m", run_fn=run)
        assert any("umount failed" in r.message for r in caplog.records)


# ---------------------------------------------------------------------------
# mount_disk — full context manager
# ---------------------------------------------------------------------------


class TestMountDisk:
    def test_full_lifecycle_partitioned_disk(self, tmp_path: Path) -> None:
        disk = tmp_path / "g.qcow2"
        disk.touch()

        # /dev/nbd0 already loaded; partitions p1 + p2 ext4.
        existing = {
            Path("/dev/nbd0"),
            Path("/dev/nbd0p1"),
            Path("/dev/nbd0p2"),
        }
        run = _RunRecorder(stdout='TYPE="ext4"')

        with mount_disk(
            disk,
            run_fn=run,
            probe_fn=lambda p: p in existing,
            sleep_fn=lambda _s: None,
            find_free_nbd_fn=lambda: Path("/dev/nbd0"),
        ) as root:
            assert root.exists()
            # Two partitions mounted as p1, p2.
            children = sorted(c.name for c in root.iterdir())
            assert children == ["p1", "p2"]

        # Verify cleanup ran: two umounts, then disconnect.
        op_sequence = [c[2] if c[:2] == ["sudo", "-n"] else c[0] for c in run.calls]
        umount_indices = [i for i, op in enumerate(op_sequence) if op == "umount"]
        disconnect_index = next(
            i for i, c in enumerate(run.calls) if "qemu-nbd" in c and "--disconnect" in c
        )
        assert len(umount_indices) == 2
        assert all(i < disconnect_index for i in umount_indices)

    def test_cleanup_runs_when_body_raises(self, tmp_path: Path) -> None:
        disk = tmp_path / "g.qcow2"
        disk.touch()

        existing = {Path("/dev/nbd0"), Path("/dev/nbd0p1")}
        run = _RunRecorder(stdout='TYPE="ext4"')

        with (
            pytest.raises(RuntimeError, match="boom"),
            mount_disk(
                disk,
                run_fn=run,
                probe_fn=lambda p: p in existing,
                sleep_fn=lambda _s: None,
                find_free_nbd_fn=lambda: Path("/dev/nbd0"),
            ),
        ):
            msg = "boom"
            raise RuntimeError(msg)

        # Disconnect MUST have run despite the body raising.
        assert any("qemu-nbd" in c and "--disconnect" in c for c in run.calls)

    def test_raw_filesystem_yields_one_partition_subdir(self, tmp_path: Path) -> None:
        disk = tmp_path / "raw.img"
        disk.touch()

        # The disk is a raw filesystem — no partition table, no
        # /dev/nbd0pN entries. discover_partitions must fall back to
        # treating /dev/nbd0 itself as the mountable target.
        existing = {Path("/dev/nbd0")}
        run = _RunRecorder(stdout='TYPE="ext4"')

        with mount_disk(
            disk,
            run_fn=run,
            probe_fn=lambda p: p in existing,
            sleep_fn=lambda _s: None,
            find_free_nbd_fn=lambda: Path("/dev/nbd0"),
        ) as root:
            children = sorted(c.name for c in root.iterdir())
            assert children == ["p1"]


# ---------------------------------------------------------------------------
# Constants
# ---------------------------------------------------------------------------


class TestConstants:
    def test_max_nbd_devices_matches_kernel_default(self) -> None:
        assert MAX_NBD_DEVICES == 16
