"""Filesystem capture — baseline + post-detonation diff.

Per FN-A0.1 decision #4: ``capture filesystem-baseline`` refuses on
a running VM by default; ``--allow-running`` is the explicit opt-in.
``capture filesystem-diff`` always allows running but tags the JSON
with ``"vm_state": "running", "disk_view": "inconsistent"``.

The qemu-nbd / mount machinery is hidden behind an injectable
``mount_disk_fn`` context manager so unit tests can substitute a
plain ``tmp_path`` mountpoint instead of needing root + a kernel
NBD module.
"""

from __future__ import annotations

import contextlib
import hashlib
import logging
import os
from datetime import UTC, datetime
from fnmatch import fnmatch
from pathlib import Path
from typing import TYPE_CHECKING

from mgf.common.config import BaseModel, ConfigDict
from mgf.common.exceptions import AppConfigError, AppConfigErrorKind
from mgf.common.versioned import VersionedFileStore
from mgf.vm.core.exceptions import OperationError

from mgf.vm.forensics.case import Artifact, ArtifactKind

if TYPE_CHECKING:
    from collections.abc import Callable, Iterator, Sequence


__all__ = [
    "BASELINE_SCHEMA_VERSION",
    "DIFF_SCHEMA_VERSION",
    "FileEntry",
    "FilesystemBaseline",
    "FilesystemDiff",
    "capture_filesystem_baseline",
    "capture_filesystem_diff",
    "load_baseline",
    "walk_tree",
]


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


BASELINE_SCHEMA_VERSION = 1
DIFF_SCHEMA_VERSION = 1


# ---------------------------------------------------------------------------
# Data shapes
#
# Pydantic frozen ``BaseModel`` + ``extra="forbid"`` gives us
# unknown-key rejection, immutability, and a free JSON serializer
# that ``VersionedFileStore`` plugs into directly. The
# ``schema_version`` field is the document-root version stamped by
# VFS on save (we pass ``version_field="schema_version"`` in
# ``_baseline_store`` / ``_diff_store`` to override the default
# ``"version"`` field name).
# ---------------------------------------------------------------------------


class FileEntry(BaseModel):
    """One file in a baseline / diff."""

    model_config = ConfigDict(extra="forbid", frozen=True)

    path: str  # POSIX-style, relative to mount root
    size_bytes: int
    mtime: float
    sha256: str = ""  # "" for symlinks / special files


class FilesystemBaseline(BaseModel):
    """Read-only snapshot of a guest's filesystem at a moment in time."""

    model_config = ConfigDict(extra="forbid", frozen=True)

    captured_at: datetime
    vm_state: str  # "running" | "stopped"
    disk_view: str  # "consistent" | "inconsistent"
    includes: tuple[str, ...] = ()
    excludes: tuple[str, ...] = ()
    entries: tuple[FileEntry, ...] = ()
    schema_version: int = BASELINE_SCHEMA_VERSION

    def by_path(self) -> dict[str, FileEntry]:
        return {e.path: e for e in self.entries}


class FilesystemDiff(BaseModel):
    """The result of comparing a current state against a baseline."""

    model_config = ConfigDict(extra="forbid", frozen=True)

    captured_at: datetime
    baseline_captured_at: datetime
    vm_state: str
    disk_view: str
    added: tuple[FileEntry, ...] = ()
    modified: tuple[FileEntry, ...] = ()
    deleted: tuple[str, ...] = ()  # paths only — file is gone
    schema_version: int = DIFF_SCHEMA_VERSION


# ---------------------------------------------------------------------------
# Public capture API
# ---------------------------------------------------------------------------


def capture_filesystem_baseline(
    vm_name: str,
    *,
    is_running: bool,
    disk_paths: Sequence[Path],
    case_dir: Path,
    includes: Sequence[str] = (),
    excludes: Sequence[str] = (),
    allow_running: bool = False,
    mount_disk_fn: Callable[[Path], contextlib.AbstractContextManager[Path]] | None = None,
    captured_at: datetime | None = None,
) -> Artifact:
    """Walk the guest's disk(s) read-only; persist a baseline JSON.

    The caller supplies the VM's ``is_running`` flag and ``disk_paths`` (the
    app reads them off the SDK ``client.get_vm().status``; D-G22). Refuses a
    running VM unless ``allow_running=True`` (FN-A0.1 #4). The mount seam is
    injectable for testing — production wraps ``qemu-nbd``; tests pass a fake
    context manager that yields a pre-populated ``tmp_path``.
    """
    if is_running and not allow_running:
        msg = (
            f"refusing to baseline {vm_name!r}: VM is running. "
            f"Stop the VM first (`vmanager vm stop {vm_name}`), "
            f"or pass --allow-running to acknowledge that the disk view "
            f"may be inconsistent."
        )
        raise OperationError(msg)

    if not disk_paths:
        msg = f"VM {vm_name!r} has no disks attached"
        raise OperationError(msg)

    if mount_disk_fn is None:
        mount_disk_fn = _qemu_nbd_mount

    when = captured_at or datetime.now(UTC)
    fs_dir = case_dir / "filesystem"
    fs_dir.mkdir(parents=True, exist_ok=True)
    out_path = fs_dir / "baseline.json"

    all_entries: list[FileEntry] = []
    for disk_path in disk_paths:
        with mount_disk_fn(disk_path) as mount:
            entries = walk_tree(mount, includes=includes, excludes=excludes)
            all_entries.extend(entries)

    baseline = FilesystemBaseline(
        captured_at=when,
        vm_state="running" if is_running else "stopped",
        disk_view="inconsistent" if is_running else "consistent",
        includes=tuple(includes),
        excludes=tuple(excludes),
        entries=tuple(all_entries),
    )
    out_path.parent.mkdir(parents=True, exist_ok=True)
    _baseline_store(out_path).save(baseline)

    return Artifact(
        kind=ArtifactKind.FS_BASELINE,
        path=out_path,
        sha256=_hash_file(out_path),
        size_bytes=out_path.stat().st_size,
        captured_at=when,
        captured_by="mgf.vm.forensics.capture.filesystem",
        metadata={
            "vm_state": baseline.vm_state,
            "disk_view": baseline.disk_view,
            "entry_count": len(baseline.entries),
            "includes": list(baseline.includes),
            "excludes": list(baseline.excludes),
        },
    )


def capture_filesystem_diff(
    vm_name: str,
    *,
    is_running: bool,
    disk_paths: Sequence[Path],
    case_dir: Path,
    includes: Sequence[str] = (),
    excludes: Sequence[str] = (),
    mount_disk_fn: Callable[[Path], contextlib.AbstractContextManager[Path]] | None = None,
    captured_at: datetime | None = None,
) -> Artifact:
    """Compute a diff against the previously-captured baseline.

    The caller supplies ``is_running`` + ``disk_paths`` (off the SDK
    ``client.get_vm().status``; D-G22). Always allows running VMs (FN-A0.1 #4)
    but tags the diff JSON with ``"disk_view": "inconsistent"`` so the analyst
    sees the caveat in the report.

    Reads the baseline from ``<case>/filesystem/baseline.json``; raises
    :class:`OperationError` if it doesn't exist.
    """
    fs_dir = case_dir / "filesystem"
    baseline_path = fs_dir / "baseline.json"
    if not baseline_path.exists():
        msg = (
            f"no baseline at {baseline_path}. "
            f"Run `vmanager forensics capture filesystem-baseline` first."
        )
        raise OperationError(msg)

    baseline = load_baseline(baseline_path)

    if not disk_paths:
        msg = f"VM {vm_name!r} has no disks attached"
        raise OperationError(msg)

    if mount_disk_fn is None:
        mount_disk_fn = _qemu_nbd_mount

    when = captured_at or datetime.now(UTC)

    current_entries: list[FileEntry] = []
    for disk_path in disk_paths:
        with mount_disk_fn(disk_path) as mount:
            entries = walk_tree(mount, includes=includes, excludes=excludes)
            current_entries.extend(entries)

    diff = _compute_diff(
        baseline=baseline,
        current=current_entries,
        captured_at=when,
        vm_state="running" if is_running else "stopped",
    )

    out_path = fs_dir / "diff.json"
    out_path.parent.mkdir(parents=True, exist_ok=True)
    _diff_store(out_path).save(diff)

    return Artifact(
        kind=ArtifactKind.FS_DIFF,
        path=out_path,
        sha256=_hash_file(out_path),
        size_bytes=out_path.stat().st_size,
        captured_at=when,
        captured_by="mgf.vm.forensics.capture.filesystem",
        metadata={
            "vm_state": diff.vm_state,
            "disk_view": diff.disk_view,
            "added_count": len(diff.added),
            "modified_count": len(diff.modified),
            "deleted_count": len(diff.deleted),
        },
    )


# ---------------------------------------------------------------------------
# Comparison
# ---------------------------------------------------------------------------


def _compute_diff(
    *,
    baseline: FilesystemBaseline,
    current: list[FileEntry],
    captured_at: datetime,
    vm_state: str,
) -> FilesystemDiff:
    baseline_by_path = baseline.by_path()
    current_by_path = {e.path: e for e in current}

    added: list[FileEntry] = []
    modified: list[FileEntry] = []
    deleted: list[str] = []

    for path, cur in current_by_path.items():
        prev = baseline_by_path.get(path)
        if prev is None:
            added.append(cur)
            continue
        if prev.sha256 != cur.sha256 or prev.size_bytes != cur.size_bytes:
            modified.append(cur)

    for path in baseline_by_path:
        if path not in current_by_path:
            deleted.append(path)

    return FilesystemDiff(
        captured_at=captured_at,
        baseline_captured_at=baseline.captured_at,
        vm_state=vm_state,
        # If the diff was taken on a running VM, the disk view is
        # always inconsistent regardless of the baseline's state.
        disk_view="inconsistent" if vm_state == "running" else "consistent",
        added=tuple(sorted(added, key=lambda e: e.path)),
        modified=tuple(sorted(modified, key=lambda e: e.path)),
        deleted=tuple(sorted(deleted)),
    )


# ---------------------------------------------------------------------------
# Walking
# ---------------------------------------------------------------------------


def walk_tree(
    root: Path,
    *,
    includes: Sequence[str] = (),
    excludes: Sequence[str] = (),
) -> list[FileEntry]:
    """Walk ``root`` recursively; return one :class:`FileEntry` per regular file.

    ``includes`` is a sequence of subpaths to walk (relative to
    ``root``). Empty includes means "walk the whole tree."
    ``excludes`` is a sequence of fnmatch patterns matched against
    the relative path; matching files are skipped.

    Symlinks are NOT followed (per FN-A0.1's §6 invariants); special
    files (devices / FIFOs / sockets) are NOT walked.
    """
    if not root.exists():
        return []

    roots: list[Path] = []
    if includes:
        for inc in includes:
            # Strip leading `/` so "/etc" is treated as relative to the
            # mount root — `root / "/etc"` would otherwise discard
            # `root` because `/etc` is absolute.
            sub = root / inc.lstrip("/")
            if sub.exists():
                roots.append(sub)
    else:
        roots.append(root)

    out: list[FileEntry] = []
    seen_inodes: set[tuple[int, int]] = set()  # cycle detection

    for sub_root in roots:
        for current_dir, _dirs, files in os.walk(sub_root, followlinks=False):
            for name in files:
                full = Path(current_dir) / name
                # rel-to-root path with leading slash so the JSON is portable.
                try:
                    rel = "/" + str(full.relative_to(root)).replace(os.sep, "/")
                except ValueError:
                    continue

                if any(fnmatch(rel, pat) or fnmatch(name, pat) for pat in excludes):
                    continue

                try:
                    st = full.lstat()
                except OSError:
                    continue
                # Skip non-regular files entirely.
                if not _is_regular_file(st.st_mode):
                    continue
                # Cycle detection (in case the host accidentally points
                # us at a recursive structure via bind mounts).
                key = (st.st_dev, st.st_ino)
                if key in seen_inodes:
                    continue
                seen_inodes.add(key)

                try:
                    digest = _hash_file(full)
                except OSError as e:
                    LOG.warning(
                        "skipping unreadable file during walk",
                        extra={"path": str(full), "error": str(e)},
                    )
                    continue

                out.append(
                    FileEntry(
                        path=rel,
                        size_bytes=st.st_size,
                        mtime=st.st_mtime,
                        sha256=digest,
                    ),
                )

    out.sort(key=lambda e: e.path)
    return out


def _is_regular_file(mode: int) -> bool:
    import stat

    return stat.S_ISREG(mode)


# ---------------------------------------------------------------------------
# qemu-nbd integration (production default; tests inject a fake)
# ---------------------------------------------------------------------------


@contextlib.contextmanager
def _qemu_nbd_mount(disk_path: Path) -> Iterator[Path]:
    """Default mount seam — wraps qemu-nbd + read-only loop-mount.

    See :mod:`mgf.vm.forensics.capture._nbd` for the implementation.
    Tests bypass this seam entirely by passing ``mount_disk_fn=`` to the
    public capture functions.
    """
    from mgf.vm.forensics.capture._nbd import mount_disk

    with mount_disk(disk_path) as root:
        yield root


# ---------------------------------------------------------------------------
# JSON serialization (delegated to VersionedFileStore)
# ---------------------------------------------------------------------------


def _baseline_store(path: Path) -> VersionedFileStore[FilesystemBaseline]:
    return VersionedFileStore(
        path,
        schema=FilesystemBaseline,
        current_version=BASELINE_SCHEMA_VERSION,
        version_field="schema_version",
        file_format="json",
    )


def _diff_store(path: Path) -> VersionedFileStore[FilesystemDiff]:
    return VersionedFileStore(
        path,
        schema=FilesystemDiff,
        current_version=DIFF_SCHEMA_VERSION,
        version_field="schema_version",
        file_format="json",
    )


def load_baseline(path: Path) -> FilesystemBaseline:
    """Load a baseline JSON from disk; raise :class:`OperationError` on schema mismatch."""
    try:
        return _baseline_store(path).load()
    except AppConfigError as e:
        # mgf-common 0.15+ tags every VFS error with ``e.kind`` so
        # consumers translate via a typed enum instead of grepping
        # the message format (closes upstream PAPER-11). Test pins
        # like ``match="unreadable"`` and ``match="schema_version"``
        # are preserved by mapping each kind back to VManager's
        # legacy wording.
        if e.kind in {AppConfigErrorKind.JSON_PARSE_FAILURE, AppConfigErrorKind.YAML_PARSE_FAILURE}:
            msg = f"baseline at {path} is unreadable: {e}"
        else:
            msg = (
                f"baseline at {path} has unsupported schema_version; "
                f"this build expects {BASELINE_SCHEMA_VERSION}."
            )
        raise OperationError(msg) from e


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


def _hash_file(path: Path, *, chunk_size: int = 1 << 16) -> str:
    h = hashlib.sha256()
    with path.open("rb") as f:
        while True:
            chunk = f.read(chunk_size)
            if not chunk:
                break
            h.update(chunk)
    return h.hexdigest()
