#!/usr/bin/env python3
"""
Systemd generator for coherent Synex Btrfs snapshot boots.

When the system is booted from the root member of a Synex snapshot set,
this generator exposes the other members of the same set.

The root snapshot remains mounted directly and read-only.

Writable runtime areas required for a usable recovery desktop are
provided through ephemeral OverlayFS mounts:

    /home
    /var/log

The lower layers always point to the read-only historical Btrfs
snapshots. Upper and work directories live under /run and therefore
disappear after reboot.

Normal boots are left completely untouched.
"""

from __future__ import annotations

from dataclasses import dataclass
import os
from pathlib import Path
import re
import shutil
import subprocess
import sys


SNAPSHOT_STORAGE_SUBVOLUME = "@snapshots"
SNAPSHOT_MEMBERS_DIRECTORY = "subvolumes"

FSTAB_PATH = Path("/etc/fstab")
CMDLINE_PATH = Path("/proc/cmdline")

OVERLAY_BASE = Path(
    "/run/synex-snapshots/overlay"
)

OVERLAY_TARGETS = {
    "/home": "home",
    "/var": "var",
    "/var/log": "log",
}

OVERLAY_PREPARE_UNIT = (
    "synex-snapshots-overlay-prepare.service"
)

DEBUG = (
    os.environ.get(
        "SYNEX_SNAPSHOTS_GENERATOR_DEBUG",
        "",
    )
    == "1"
)


@dataclass(frozen=True)
class FstabEntry:
    """Relevant information from one /etc/fstab entry."""

    source: str
    target: str
    filesystem: str
    options: str


@dataclass(frozen=True)
class MountPlan:
    """One historical mount to expose during snapshot boot."""

    source: str
    target: str
    options: str
    snapshot_path: str
    unit_name: str
    overlay_key: str | None


def debug(message: str) -> None:
    """Print diagnostic output only when explicitly requested."""

    if DEBUG:
        print(
            f"synex-snapshots-generator: {message}",
            file=sys.stderr,
        )


def decode_fstab_field(value: str) -> str:
    """Decode standard escaped characters used in fstab fields."""

    replacements = (
        ("\\040", " "),
        ("\\011", "\t"),
        ("\\012", "\n"),
        ("\\134", "\\"),
    )

    result = value

    for encoded, decoded in replacements:
        result = result.replace(
            encoded,
            decoded,
        )

    return result


def read_fstab() -> list[FstabEntry]:
    """Read usable entries from /etc/fstab."""

    try:
        contents = FSTAB_PATH.read_text(
            encoding="utf-8",
        )
    except (OSError, UnicodeError) as exc:
        debug(
            f"cannot read {FSTAB_PATH}: {exc}"
        )
        return []

    entries: list[FstabEntry] = []

    for raw_line in contents.splitlines():
        line = raw_line.strip()

        if not line or line.startswith("#"):
            continue

        fields = line.split()

        if len(fields) < 4:
            continue

        source, target, filesystem, options = (
            fields[:4]
        )

        entries.append(
            FstabEntry(
                source=decode_fstab_field(source),
                target=decode_fstab_field(target),
                filesystem=filesystem,
                options=decode_fstab_field(options),
            )
        )

    return entries


def normalize_subvolume_path(
    path: str,
) -> str:
    """Return a filesystem-relative Btrfs subvolume path."""

    return path.strip().lstrip("/").rstrip("/")


def get_option_value(
    options: str,
    key: str,
) -> str | None:
    """Return one comma-separated mount option value."""

    prefix = f"{key}="

    for option in options.split(","):
        option = option.strip()

        if option.startswith(prefix):
            return option[len(prefix):]

    return None


def source_uuid(
    source: str,
) -> str | None:
    """Extract a UUID from supported fstab source forms."""

    if source.startswith("UUID="):
        uuid = source[5:].strip()

        return uuid or None

    prefix = "/dev/disk/by-uuid/"

    if source.startswith(prefix):
        uuid = source[len(prefix):].strip()

        return uuid or None

    return None


def detect_snapshot_boot() -> tuple[str, str, str] | None:
    """
    Detect a Synex snapshot root from the kernel command line.

    Returns:
        (root_uuid, set_name, root_member)

    Expected root subvolume:

        @snapshots/<set>/subvolumes/<root-member>
    """

    try:
        cmdline = CMDLINE_PATH.read_text(
            encoding="utf-8",
        ).strip()
    except (OSError, UnicodeError) as exc:
        debug(
            f"cannot read {CMDLINE_PATH}: {exc}"
        )
        return None

    root_uuid: str | None = None
    root_subvolume: str | None = None

    for token in cmdline.split():
        if token.startswith("root=UUID="):
            root_uuid = token[len("root=UUID="):]

        elif token.startswith("rootflags="):
            rootflags = token[len("rootflags="):]

            for option in rootflags.split(","):
                if option.startswith("subvol="):
                    root_subvolume = (
                        option[len("subvol="):]
                    )
                    break

    if not root_uuid or not root_subvolume:
        debug(
            "normal boot: no snapshot root detected"
        )
        return None

    normalized_root = normalize_subvolume_path(
        root_subvolume
    )

    pattern = re.compile(
        rf"^{re.escape(SNAPSHOT_STORAGE_SUBVOLUME)}"
        rf"/(?P<set_name>[^/]+)"
        rf"/{re.escape(SNAPSHOT_MEMBERS_DIRECTORY)}"
        rf"/(?P<root_member>[^/]+)$"
    )

    match = pattern.fullmatch(
        normalized_root
    )

    if match is None:
        debug(
            "normal boot: root subvolume is not "
            "inside Synex snapshot storage"
        )
        return None

    set_name = match.group("set_name")

    if not re.fullmatch(
        r"[A-Za-z0-9._-]+",
        set_name,
    ):
        debug(
            "snapshot boot rejected: invalid snapshot set name"
        )
        return None

    return (
        root_uuid,
        set_name,
        match.group("root_member"),
    )


def list_btrfs_subvolumes() -> set[str] | None:
    """Return all Btrfs subvolume paths visible from root."""

    btrfs_command = shutil.which("btrfs")

    if btrfs_command is None:
        debug("btrfs command not found")
        return None

    try:
        result = subprocess.run(
            [
                btrfs_command,
                "subvolume",
                "list",
                "/",
            ],
            stdout=subprocess.PIPE,
            stderr=subprocess.PIPE,
            text=True,
            check=False,
            timeout=10,
            env={
                **os.environ,
                "LC_ALL": "C",
                "LANG": "C",
            },
        )
    except (
        OSError,
        subprocess.TimeoutExpired,
    ) as exc:
        debug(
            f"cannot list Btrfs subvolumes: {exc}"
        )
        return None

    if result.returncode != 0:
        error = result.stderr.strip()

        debug(
            "btrfs subvolume list failed"
            + (
                f": {error}"
                if error
                else ""
            )
        )

        return None

    paths: set[str] = set()

    for line in result.stdout.splitlines():
        marker = " path "

        if marker not in line:
            continue

        _, path = line.split(
            marker,
            1,
        )

        normalized = normalize_subvolume_path(
            path
        )

        if normalized:
            paths.add(normalized)

    return paths


def systemd_mount_unit_name(
    mountpoint: str,
) -> str | None:
    """Convert an absolute path into its systemd .mount unit name."""

    systemd_escape = shutil.which(
        "systemd-escape"
    )

    if systemd_escape is None:
        debug("systemd-escape command not found")
        return None

    try:
        result = subprocess.run(
            [
                systemd_escape,
                "--path",
                "--suffix=mount",
                mountpoint,
            ],
            stdout=subprocess.PIPE,
            stderr=subprocess.PIPE,
            text=True,
            check=False,
            timeout=5,
            env={
                **os.environ,
                "LC_ALL": "C",
                "LANG": "C",
            },
        )
    except (
        OSError,
        subprocess.TimeoutExpired,
    ) as exc:
        debug(
            f"cannot escape mountpoint {mountpoint}: {exc}"
        )
        return None

    if result.returncode != 0:
        debug(
            f"cannot derive systemd unit for {mountpoint}"
        )
        return None

    unit_name = result.stdout.strip()

    if not unit_name.endswith(".mount"):
        return None

    return unit_name


def historical_btrfs_options(
    options: str,
    snapshot_path: str,
) -> str:
    """
    Replace the canonical subvolume with a historical snapshot.

    Historical layers are explicitly mounted read-only even though
    Synex snapshots already carry the Btrfs ro=true property.
    """

    result: list[str] = []
    replaced = False

    for option in options.split(","):
        option = option.strip()

        if not option:
            continue

        if option.startswith("subvolid="):
            continue

        if option == "rw":
            continue

        if option == "ro":
            continue

        if option.startswith("subvol="):
            result.append(
                f"subvol=/{snapshot_path}"
            )
            replaced = True
            continue

        result.append(option)

    if not replaced:
        result.insert(
            0,
            f"subvol=/{snapshot_path}",
        )

    result.append("ro")

    return ",".join(result)


def add_local_fs_requirement(
    early_directory: Path,
    unit_name: str,
) -> bool:
    """Make local-fs.target require one generated mount unit."""

    try:
        requires_directory = (
            early_directory
            / "local-fs.target.requires"
        )

        requires_directory.mkdir(
            parents=True,
            exist_ok=True,
        )

        dependency = (
            requires_directory
            / unit_name
        )

        if dependency.is_symlink() or dependency.exists():
            dependency.unlink()

        dependency.symlink_to(
            f"../{unit_name}"
        )

    except OSError as exc:
        debug(
            f"cannot link {unit_name} into local-fs.target: {exc}"
        )
        return False

    return True


def write_prepare_service(
    early_directory: Path,
    plans: list[MountPlan],
) -> bool:
    """Generate the oneshot service that prepares OverlayFS runtime dirs."""

    overlay_plans = [
        plan
        for plan in plans
        if plan.overlay_key is not None
    ]

    if not overlay_plans:
        return True

    directories: list[str] = []
    before_units: list[str] = []

    for plan in overlay_plans:
        assert plan.overlay_key is not None

        base = (
            OVERLAY_BASE
            / plan.overlay_key
        )

        lower = base / "lower"
        upper = base / "upper"
        work = base / "work"

        directories.extend(
            [
                str(lower),
                str(upper),
                str(work),
            ]
        )

        lower_unit = systemd_mount_unit_name(
            str(lower)
        )

        if lower_unit is None:
            return False

        before_units.extend(
            [
                lower_unit,
                plan.unit_name,
            ]
        )

    mkdir_command = (
        shutil.which("mkdir")
        or "/usr/bin/mkdir"
    )

    modprobe_command = shutil.which(
        "modprobe"
    )

    unit_lines = [
        "[Unit]",
        "Description=Prepare Synex snapshot OverlayFS runtime directories",
        "DefaultDependencies=no",
        "After=local-fs-pre.target",
        f"Before={' '.join(sorted(set(before_units)))}",
        "",
        "[Service]",
        "Type=oneshot",
    ]

    if modprobe_command is not None:
        unit_lines.append(
            f"ExecStart=-{modprobe_command} overlay"
        )

    unit_lines.append(
        "ExecStart="
        + mkdir_command
        + " -p "
        + " ".join(directories)
    )

    unit_lines.extend(
        [
            "RemainAfterExit=yes",
            "",
        ]
    )

    try:
        (
            early_directory
            / OVERLAY_PREPARE_UNIT
        ).write_text(
            "\n".join(unit_lines),
            encoding="utf-8",
        )
    except OSError as exc:
        debug(
            f"cannot generate {OVERLAY_PREPARE_UNIT}: {exc}"
        )
        return False

    return True


def write_direct_mount_unit(
    early_directory: Path,
    *,
    plan: MountPlan,
    root_uuid: str,
) -> bool:
    """Generate a direct read-only historical Btrfs mount."""

    unit_contents = (
        "[Unit]\n"
        f"Description=Synex historical snapshot member for {plan.target}\n"
        "Before=local-fs.target\n"
        "\n"
        "[Mount]\n"
        f"What=/dev/disk/by-uuid/{root_uuid}\n"
        f"Where={plan.target}\n"
        "Type=btrfs\n"
        f"Options={plan.options}\n"
    )

    try:
        (
            early_directory
            / plan.unit_name
        ).write_text(
            unit_contents,
            encoding="utf-8",
        )
    except OSError as exc:
        debug(
            f"cannot generate {plan.unit_name}: {exc}"
        )
        return False

    return add_local_fs_requirement(
        early_directory,
        plan.unit_name,
    )


def write_overlay_mount_units(
    early_directory: Path,
    *,
    plan: MountPlan,
    root_uuid: str,
) -> bool:
    """
    Generate the historical lower Btrfs mount and final OverlayFS mount.
    """

    if plan.overlay_key is None:
        return False

    base = (
        OVERLAY_BASE
        / plan.overlay_key
    )

    lower = base / "lower"
    upper = base / "upper"
    work = base / "work"

    lower_unit_name = (
        systemd_mount_unit_name(
            str(lower)
        )
    )

    if lower_unit_name is None:
        return False

    lower_contents = (
        "[Unit]\n"
        f"Description=Synex historical lower layer for {plan.target}\n"
        f"Requires={OVERLAY_PREPARE_UNIT}\n"
        f"After={OVERLAY_PREPARE_UNIT}\n"
        f"Before={plan.unit_name} local-fs.target\n"
        "\n"
        "[Mount]\n"
        f"What=/dev/disk/by-uuid/{root_uuid}\n"
        f"Where={lower}\n"
        "Type=btrfs\n"
        f"Options={plan.options}\n"
    )

    overlay_contents = (
        "[Unit]\n"
        f"Description=Synex recovery overlay for {plan.target}\n"
        f"Requires={lower_unit_name}\n"
        f"After={lower_unit_name}\n"
        "Before=local-fs.target\n"
        "\n"
        "[Mount]\n"
        "What=overlay\n"
        f"Where={plan.target}\n"
        "Type=overlay\n"
        f"Options=lowerdir={lower},upperdir={upper},workdir={work}\n"
    )

    try:
        (
            early_directory
            / lower_unit_name
        ).write_text(
            lower_contents,
            encoding="utf-8",
        )

        (
            early_directory
            / plan.unit_name
        ).write_text(
            overlay_contents,
            encoding="utf-8",
        )

    except OSError as exc:
        debug(
            f"cannot generate overlay for {plan.target}: {exc}"
        )
        return False

    return add_local_fs_requirement(
        early_directory,
        plan.unit_name,
    )


def build_mount_plans(
    entries: list[FstabEntry],
    *,
    root_uuid: str,
    set_name: str,
    available_subvolumes: set[str],
) -> list[MountPlan]:
    """Build all valid historical mount plans for this snapshot set."""

    plans: list[MountPlan] = []

    for entry in entries:
        if entry.filesystem != "btrfs":
            continue

        if entry.target in {
            "/",
            "/.snapshots",
        }:
            continue

        entry_uuid = source_uuid(
            entry.source
        )

        if (
            entry_uuid is None
            or entry_uuid.lower()
            != root_uuid.lower()
        ):
            continue

        canonical_subvolume = get_option_value(
            entry.options,
            "subvol",
        )

        if canonical_subvolume is None:
            continue

        canonical_subvolume = (
            normalize_subvolume_path(
                canonical_subvolume
            )
        )

        if not canonical_subvolume:
            continue

        snapshot_path = (
            f"{SNAPSHOT_STORAGE_SUBVOLUME}/"
            f"{set_name}/"
            f"{SNAPSHOT_MEMBERS_DIRECTORY}/"
            f"{canonical_subvolume}"
        )

        if snapshot_path not in available_subvolumes:
            debug(
                f"no snapshot member for {entry.target}: "
                f"{snapshot_path}"
            )
            continue

        unit_name = systemd_mount_unit_name(
            entry.target
        )

        if unit_name is None:
            continue

        plans.append(
            MountPlan(
                source=entry.source,
                target=entry.target,
                options=historical_btrfs_options(
                    entry.options,
                    snapshot_path,
                ),
                snapshot_path=snapshot_path,
                unit_name=unit_name,
                overlay_key=OVERLAY_TARGETS.get(
                    entry.target
                ),
            )
        )

    return plans


def main() -> int:
    """Generate coherent mounts for a Synex snapshot boot."""

    if len(sys.argv) != 4:
        debug(
            "expected systemd generator arguments: "
            "normal-dir early-dir late-dir"
        )
        return 1

    early_directory = Path(
        sys.argv[2]
    )

    snapshot_boot = detect_snapshot_boot()

    if snapshot_boot is None:
        return 0

    (
        root_uuid,
        set_name,
        root_member,
    ) = snapshot_boot

    debug(
        "snapshot boot detected: "
        f"set={set_name}, "
        f"root={root_member}, "
        f"uuid={root_uuid}"
    )

    entries = read_fstab()

    if not entries:
        debug("no usable fstab entries found")
        return 0

    root_entry = next(
        (
            entry
            for entry in entries
            if (
                entry.target == "/"
                and entry.filesystem == "btrfs"
            )
        ),
        None,
    )

    if root_entry is None:
        debug(
            "snapshot boot rejected: "
            "fstab root is not Btrfs"
        )
        return 0

    fstab_root_uuid = source_uuid(
        root_entry.source
    )

    if (
        fstab_root_uuid is None
        or fstab_root_uuid.lower()
        != root_uuid.lower()
    ):
        debug(
            "snapshot boot rejected: "
            "root UUID does not match fstab"
        )
        return 0

    canonical_root = get_option_value(
        root_entry.options,
        "subvol",
    )

    if canonical_root is None:
        debug(
            "snapshot boot rejected: "
            "root fstab entry has no subvol option"
        )
        return 0

    canonical_root = normalize_subvolume_path(
        canonical_root
    )

    if canonical_root != root_member:
        debug(
            "snapshot boot rejected: "
            f"snapshot root member '{root_member}' "
            f"does not match canonical root '{canonical_root}'"
        )
        return 0

    available_subvolumes = (
        list_btrfs_subvolumes()
    )

    if available_subvolumes is None:
        return 0

    plans = build_mount_plans(
        entries,
        root_uuid=root_uuid,
        set_name=set_name,
        available_subvolumes=available_subvolumes,
    )

    if not plans:
        debug(
            "snapshot root has no matching historical "
            "secondary members"
        )
        return 0

    try:
        early_directory.mkdir(
            parents=True,
            exist_ok=True,
        )
    except OSError as exc:
        debug(
            f"cannot create generator output directory: {exc}"
        )
        return 0

    if not write_prepare_service(
        early_directory,
        plans,
    ):
        debug(
            "cannot prepare OverlayFS runtime service"
        )
        return 0

    generated = 0

    for plan in plans:
        if plan.overlay_key is not None:
            success = write_overlay_mount_units(
                early_directory,
                plan=plan,
                root_uuid=root_uuid,
            )

            if success:
                generated += 1

                debug(
                    f"generated overlay {plan.unit_name}: "
                    f"{plan.target} -> /{plan.snapshot_path}"
                )

            continue

        success = write_direct_mount_unit(
            early_directory,
            plan=plan,
            root_uuid=root_uuid,
        )

        if success:
            generated += 1

            debug(
                f"generated historical {plan.unit_name}: "
                f"{plan.target} -> /{plan.snapshot_path}"
            )

    debug(
        f"generation complete: {generated} mount unit(s)"
    )

    return 0


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