#!/usr/bin/env python3
"""Create and verify byte-complete, symlink-safe release backup snapshots."""

from __future__ import annotations

import argparse
import errno
import hashlib
import json
import os
import re
import stat
from dataclasses import dataclass
from pathlib import Path, PurePosixPath


MANIFEST_HEADER = "relative_path\tsize_bytes\tmtime_ns\tsha256"
SAFE_LABEL = re.compile(r"^[A-Za-z0-9][A-Za-z0-9._-]{0,63}$")
FICLONE = 0x40049409
FICLONE_FALLBACK_ERRORS = {
    errno.EXDEV,
    errno.EOPNOTSUPP,
    errno.ENOTTY,
    errno.EINVAL,
    errno.ENOSYS,
}


@dataclass(frozen=True)
class SourceEntry:
    path: Path
    relative_path: str
    stat_result: os.stat_result
    is_directory: bool


@dataclass(frozen=True)
class ManifestEntry:
    relative_path: str
    size_bytes: int
    mtime_ns: int
    sha256: str


def _safe_relative(value: str) -> str:
    if not value or "\t" in value or "\n" in value or "\r" in value:
        raise RuntimeError(f"Unsupported snapshot path: {value!r}")
    candidate = PurePosixPath(value)
    if candidate.is_absolute() or any(part in {"", ".", ".."} for part in candidate.parts):
        raise RuntimeError(f"Unsafe snapshot path: {value!r}")
    return candidate.as_posix()


def _collect(root: Path) -> list[SourceEntry]:
    entries: list[SourceEntry] = []
    for current, directory_names, file_names in os.walk(root, topdown=True, followlinks=False):
        directory_names.sort()
        file_names.sort()
        current_path = Path(current)
        for name, is_directory in [
            *((name, True) for name in directory_names),
            *((name, False) for name in file_names),
        ]:
            path = current_path / name
            details = path.lstat()
            relative = _safe_relative(path.relative_to(root).as_posix())
            if stat.S_ISLNK(details.st_mode):
                raise RuntimeError(f"Release snapshot refuses symlinked content: {relative}")
            if is_directory and not stat.S_ISDIR(details.st_mode):
                raise RuntimeError(f"Snapshot directory changed during inventory: {relative}")
            if not is_directory and not stat.S_ISREG(details.st_mode):
                raise RuntimeError(f"Release snapshot refuses non-regular content: {relative}")
            entries.append(SourceEntry(path, relative, details, is_directory))
    return entries


def _same_file_state(before: os.stat_result, after: os.stat_result) -> bool:
    return (
        before.st_dev,
        before.st_ino,
        before.st_size,
        before.st_mtime_ns,
    ) == (
        after.st_dev,
        after.st_ino,
        after.st_size,
        after.st_mtime_ns,
    )


def _sha256(path: Path) -> str:
    digest = hashlib.sha256()
    with path.open("rb") as handle:
        for chunk in iter(lambda: handle.read(1024 * 1024), b""):
            digest.update(chunk)
    return digest.hexdigest()


def _copy_all(source_descriptor: int, destination_descriptor: int) -> None:
    while True:
        value = os.read(source_descriptor, 1024 * 1024)
        if not value:
            return
        view = memoryview(value)
        while view:
            written = os.write(destination_descriptor, view)
            if written <= 0:
                raise RuntimeError("Snapshot copy stopped before writing all bytes")
            view = view[written:]


def _clone_or_copy(entry: SourceEntry, destination: Path) -> ManifestEntry:
    source_flags = os.O_RDONLY | getattr(os, "O_BINARY", 0) | getattr(os, "O_NOFOLLOW", 0)
    destination_flags = os.O_WRONLY | os.O_CREAT | os.O_EXCL | getattr(os, "O_BINARY", 0)
    source_descriptor = os.open(entry.path, source_flags)
    destination_descriptor = -1
    try:
        opened = os.fstat(source_descriptor)
        if not stat.S_ISREG(opened.st_mode) or not _same_file_state(entry.stat_result, opened):
            raise RuntimeError(f"Snapshot file changed before copying: {entry.relative_path}")
        destination.parent.mkdir(parents=True, exist_ok=True)
        destination_descriptor = os.open(destination, destination_flags, stat.S_IMODE(opened.st_mode))
        cloned = False
        if os.name == "posix":
            try:
                import fcntl

                fcntl.ioctl(destination_descriptor, FICLONE, source_descriptor)
                cloned = True
            except OSError as exc:
                if exc.errno not in FICLONE_FALLBACK_ERRORS:
                    raise
        if not cloned:
            os.lseek(source_descriptor, 0, os.SEEK_SET)
            os.ftruncate(destination_descriptor, 0)
            _copy_all(source_descriptor, destination_descriptor)
        os.fsync(destination_descriptor)
        after = os.fstat(source_descriptor)
        if not _same_file_state(opened, after):
            raise RuntimeError(f"Snapshot file changed while copying: {entry.relative_path}")
    except BaseException:
        if destination_descriptor >= 0:
            os.close(destination_descriptor)
            destination_descriptor = -1
        destination.unlink(missing_ok=True)
        raise
    finally:
        if destination_descriptor >= 0:
            os.close(destination_descriptor)
        os.close(source_descriptor)

    os.chmod(destination, stat.S_IMODE(entry.stat_result.st_mode) & ~0o222, follow_symlinks=False)
    retained_times = (entry.stat_result.st_atime_ns, entry.stat_result.st_mtime_ns)
    try:
        os.utime(destination, ns=retained_times, follow_symlinks=False)
    except NotImplementedError:
        # Windows does not expose no-follow utime. The destination was created
        # exclusively above; recheck it before using the portable call.
        if destination.is_symlink():
            destination.unlink(missing_ok=True)
            raise RuntimeError(f"Snapshot destination became a symlink: {entry.relative_path}")
        os.utime(destination, ns=retained_times)
    source_checksum = _sha256(entry.path)
    snapshot_checksum = _sha256(destination)
    final_source = entry.path.lstat()
    if not _same_file_state(entry.stat_result, final_source):
        raise RuntimeError(f"Snapshot file changed during checksum verification: {entry.relative_path}")
    if source_checksum != snapshot_checksum:
        raise RuntimeError(f"Snapshot checksum differs from source: {entry.relative_path}")
    return ManifestEntry(
        relative_path=entry.relative_path,
        size_bytes=entry.stat_result.st_size,
        mtime_ns=entry.stat_result.st_mtime_ns,
        sha256=snapshot_checksum,
    )


def _link_verified_prior(
    entry: SourceEntry,
    destination: Path,
    prior_root: Path,
    prior_manifest: dict[str, ManifestEntry],
) -> ManifestEntry | None:
    retained = prior_manifest.get(entry.relative_path)
    if retained is None or retained.size_bytes != entry.stat_result.st_size:
        return None
    source_checksum = _sha256(entry.path)
    final_source = entry.path.lstat()
    if not _same_file_state(entry.stat_result, final_source):
        raise RuntimeError(f"Snapshot file changed during prior comparison: {entry.relative_path}")
    if source_checksum != retained.sha256:
        return None
    prior_path = prior_root / entry.relative_path
    try:
        prior_details = prior_path.lstat()
    except FileNotFoundError:
        return None
    if not stat.S_ISREG(prior_details.st_mode) or prior_details.st_size != retained.size_bytes:
        raise RuntimeError(f"Prior snapshot file is not reusable: {entry.relative_path}")
    if _sha256(prior_path) != retained.sha256:
        raise RuntimeError(f"Prior snapshot checksum changed: {entry.relative_path}")
    destination.parent.mkdir(parents=True, exist_ok=True)
    os.link(prior_path, destination, follow_symlinks=False)
    if destination.stat().st_size != retained.size_bytes or _sha256(destination) != retained.sha256:
        destination.unlink(missing_ok=True)
        raise RuntimeError(f"Hard-linked snapshot verification failed: {entry.relative_path}")
    return ManifestEntry(
        relative_path=entry.relative_path,
        size_bytes=entry.stat_result.st_size,
        mtime_ns=entry.stat_result.st_mtime_ns,
        sha256=source_checksum,
    )


def read_manifest(path: Path) -> dict[str, ManifestEntry]:
    lines = path.read_text(encoding="utf-8").splitlines()
    if not lines or lines[0] != MANIFEST_HEADER:
        raise RuntimeError(f"Snapshot inventory has an invalid header: {path}")
    entries: dict[str, ManifestEntry] = {}
    for line in lines[1:]:
        fields = line.split("\t")
        if len(fields) != 4:
            raise RuntimeError(f"Snapshot inventory has an invalid row: {line!r}")
        relative_path, size_text, mtime_text, checksum = fields
        relative_path = _safe_relative(relative_path)
        if relative_path in entries:
            raise RuntimeError(f"Snapshot inventory repeats a path: {relative_path}")
        try:
            size_bytes = int(size_text)
            mtime_ns = int(mtime_text)
        except ValueError as exc:
            raise RuntimeError(f"Snapshot inventory has invalid metadata: {relative_path}") from exc
        if size_bytes < 0 or mtime_ns < 0 or not re.fullmatch(r"[0-9a-f]{64}", checksum):
            raise RuntimeError(f"Snapshot inventory has invalid retained state: {relative_path}")
        entries[relative_path] = ManifestEntry(relative_path, size_bytes, mtime_ns, checksum)
    return entries


def verify_snapshot(snapshot_path: Path, manifest_path: Path) -> None:
    root = snapshot_path.expanduser().resolve(strict=True)
    if not root.is_dir():
        raise RuntimeError(f"Snapshot path is not a directory: {root}")
    expected = read_manifest(manifest_path)
    observed_entries = _collect(root)
    observed_files = {item.relative_path: item for item in observed_entries if not item.is_directory}
    extra = sorted(set(observed_files) - set(expected))
    missing = sorted(set(expected) - set(observed_files))
    if extra:
        raise RuntimeError(f"Snapshot contains unmanifested files: {', '.join(extra[:10])}")
    if missing:
        raise RuntimeError(f"Snapshot omits manifested files: {', '.join(missing[:10])}")
    for relative, retained in expected.items():
        current = observed_files[relative]
        if current.stat_result.st_size != retained.size_bytes:
            raise RuntimeError(f"Snapshot size differs for: {relative}")
        if _sha256(current.path) != retained.sha256:
            raise RuntimeError(f"Snapshot checksum differs for: {relative}")


def create_snapshot(
    source: Path,
    snapshot_path: Path,
    manifest_path: Path,
    *,
    label: str,
    link_dest_snapshot: Path | None = None,
    link_dest_manifest: Path | None = None,
) -> None:
    if not SAFE_LABEL.fullmatch(label):
        raise RuntimeError(f"Unsafe snapshot label: {label!r}")
    root = source.expanduser()
    if root.is_symlink():
        raise RuntimeError(f"Release snapshot refuses a symlinked root: {root}")
    root = root.resolve(strict=True)
    if not root.is_dir():
        raise RuntimeError(f"Snapshot source is not a directory: {root}")
    snapshot_path = snapshot_path.expanduser().resolve()
    manifest_path = manifest_path.expanduser().resolve()
    for output in (snapshot_path, manifest_path):
        try:
            output.relative_to(root)
        except ValueError:
            pass
        else:
            raise RuntimeError("Release snapshot output must not be inside its source tree")
    if snapshot_path.exists():
        raise RuntimeError(f"Snapshot destination already exists: {snapshot_path}")
    snapshot_path.mkdir(parents=True, exist_ok=False)
    manifest_path.parent.mkdir(parents=True, exist_ok=True)

    prior_root: Path | None = None
    prior_manifest: dict[str, ManifestEntry] = {}
    if (link_dest_snapshot is None) != (link_dest_manifest is None):
        raise RuntimeError("Prior snapshot and manifest must be supplied together")
    if link_dest_snapshot is not None and link_dest_manifest is not None:
        prior_root = link_dest_snapshot.expanduser().resolve(strict=True)
        prior_manifest = read_manifest(link_dest_manifest.expanduser().resolve(strict=True))

    initial = _collect(root)
    retained: list[ManifestEntry] = []
    for entry in initial:
        destination = snapshot_path / entry.relative_path
        if entry.is_directory:
            destination.mkdir(parents=True, exist_ok=False)
            os.chmod(destination, stat.S_IMODE(entry.stat_result.st_mode), follow_symlinks=False)
            continue
        linked = (
            _link_verified_prior(entry, destination, prior_root, prior_manifest)
            if prior_root is not None
            else None
        )
        retained.append(linked or _clone_or_copy(entry, destination))
    final = _collect(root)
    if [(item.relative_path, item.is_directory) for item in initial] != [
        (item.relative_path, item.is_directory) for item in final
    ]:
        raise RuntimeError(f"Snapshot source contents changed while backup was running: {root}")
    with manifest_path.open("w", encoding="utf-8", newline="\n") as handle:
        handle.write(f"{MANIFEST_HEADER}\n")
        for entry in retained:
            handle.write(
                f"{entry.relative_path}\t{entry.size_bytes}\t{entry.mtime_ns}\t{entry.sha256}\n"
            )
    verify_snapshot(snapshot_path, manifest_path)


def verify_backup(backup_dir: Path) -> None:
    root = backup_dir.expanduser().resolve(strict=True)
    payload = json.loads((root / "manifest.json").read_text(encoding="utf-8"))
    for label in ("storage", "models"):
        requested = payload.get(f"{label}_inventory_requested") is True
        snapshotted = payload.get(f"{label}_snapshot_requested") is True
        if requested != snapshotted:
            raise RuntimeError(f"Backup manifest does not bind the {label} inventory to a snapshot")
        if requested:
            verify_snapshot(root / f"{label}-snapshot", root / f"{label}-manifest.tsv")


def parse_args() -> argparse.Namespace:
    parser = argparse.ArgumentParser(description=__doc__)
    subparsers = parser.add_subparsers(dest="command", required=True)
    create = subparsers.add_parser("create")
    create.add_argument("--source", type=Path, required=True)
    create.add_argument("--snapshot", type=Path, required=True)
    create.add_argument("--manifest", type=Path, required=True)
    create.add_argument("--label", required=True)
    create.add_argument("--link-dest-snapshot", type=Path)
    create.add_argument("--link-dest-manifest", type=Path)
    verify = subparsers.add_parser("verify")
    verify.add_argument("--snapshot", type=Path, required=True)
    verify.add_argument("--manifest", type=Path, required=True)
    verify_backup_parser = subparsers.add_parser("verify-backup")
    verify_backup_parser.add_argument("--backup-dir", type=Path, required=True)
    return parser.parse_args()


def main() -> int:
    args = parse_args()
    if args.command == "create":
        create_snapshot(
            args.source,
            args.snapshot,
            args.manifest,
            label=args.label,
            link_dest_snapshot=args.link_dest_snapshot,
            link_dest_manifest=args.link_dest_manifest,
        )
    elif args.command == "verify":
        verify_snapshot(args.snapshot, args.manifest)
    else:
        verify_backup(args.backup_dir)
    return 0


if __name__ == "__main__":
    raise SystemExit(main())