#!/usr/bin/env python3

"""Snapshots, through snapper.

A point in time you can go back to, per btrfs subvolume. This machine already
had snapper running hourly, but only for / -- /home is a separate subvolume and
had no configuration at all, so six hundred snapshots existed and not one of
them contained a document.

Deliberately absent: rollback. snapper's rollback works by changing the btrfs
default subvolume, and this system's fstab pins subvol=root and subvol=home
explicitly, which overrides it -- so a rollback would appear to succeed and
change nothing after a reboot. Restoring files and folders out of a snapshot
needs no reboot, cannot affect booting, and covers the cases people actually
hit.

    panama-snapshots snapshot
    panama-snapshots create CONFIG DESCRIPTION
    panama-snapshots delete CONFIG NUMBER
    panama-snapshots set-retention CONFIG HOURLY DAILY WEEKLY
    panama-snapshots set-timeline CONFIG true|false
    panama-snapshots browse CONFIG NUMBER [RELATIVE_PATH]
    panama-snapshots restore CONFIG NUMBER RELATIVE_PATH
"""

from __future__ import annotations

import json
import os
import re
import shutil
import subprocess
import sys
from pathlib import Path

CONFIG_NAME = re.compile(r"^[A-Za-z0-9_.-]{1,64}$")

# Directory listings are for choosing something to restore, not for browsing a
# terabyte. A folder with more entries than this is reported as truncated.
BROWSE_LIMIT = 400


class BoundaryError(RuntimeError):
    """A user-visible validation or snapper failure."""


def run(command: list[str], timeout: float = 30.0) -> subprocess.CompletedProcess:
    try:
        return subprocess.run(command, capture_output=True, text=True,
                              timeout=timeout, check=False)
    except subprocess.TimeoutExpired as error:
        raise BoundaryError("snapper did not answer in time.") from error
    except OSError as error:
        raise BoundaryError("snapper is not available.") from error


def require_config(name: str) -> str:
    if not CONFIG_NAME.fullmatch(name or ""):
        raise BoundaryError("That is not a snapshot configuration.")
    return name


def require_number(value: str) -> int:
    if not str(value).isdigit():
        raise BoundaryError("That is not a snapshot.")
    number = int(value)
    # 0 is the live filesystem, not a snapshot, and must never be a target.
    if number < 1:
        raise BoundaryError("That is the current state, not a snapshot.")
    return number


def config_names() -> list[str]:
    result = run(["snapper", "list-configs"])
    if result.returncode != 0:
        return []
    names = []
    for line in result.stdout.splitlines()[2:]:
        parts = [part.strip() for part in line.split("│")]
        if len(parts) >= 2 and CONFIG_NAME.fullmatch(parts[0]):
            names.append(parts[0])
    return names


def config_settings(name: str) -> dict:
    result = run(["snapper", "-c", name, "get-config"])
    if result.returncode != 0:
        return {}
    values = {}
    for line in result.stdout.splitlines()[2:]:
        parts = [part.strip() for part in line.split("│")]
        if len(parts) >= 2:
            values[parts[0]] = parts[1]
    return values


def snapshots_for(name: str) -> list[dict]:
    result = run(["snapper", "-c", name, "list"], timeout=45)
    if result.returncode != 0:
        return []
    entries = []
    for line in result.stdout.splitlines()[2:]:
        parts = [part.strip() for part in line.split("│")]
        if len(parts) < 7 or not parts[0].isdigit():
            continue
        number = int(parts[0])
        if number == 0:
            continue
        entries.append({
            "number": number,
            "kind": parts[1],
            "date": parts[3],
            "user": parts[4],
            "cleanup": parts[5],
            "description": parts[6],
            # A snapshot with no cleanup algorithm is not on the timeline's
            # list to remove, which is what "kept" means to someone reading it.
            "kept": parts[5] == "",
        })
    entries.sort(key=lambda entry: entry["number"], reverse=True)
    return entries


def btrfs_subvolumes() -> list[dict]:
    """Mounted btrfs subvolumes, so the page can name what is NOT protected."""
    result = run(["findmnt", "-t", "btrfs", "-J", "-o", "TARGET,OPTIONS"])
    if result.returncode != 0:
        return []
    try:
        payload = json.loads(result.stdout)
    except json.JSONDecodeError:
        return []

    found = []

    def walk(nodes):
        for node in nodes:
            options = str(node.get("options", ""))
            target = str(node.get("target", ""))
            match = re.search(r"subvol=(/[^,]*)", options)
            # .snapshots holds the snapshots themselves and is not a thing to
            # protect; listing it would offer to snapshot the snapshots.
            if match and "/.snapshots" not in target:
                found.append({"path": target, "subvolume": match.group(1)})
            walk(node.get("children", []))

    walk(payload.get("filesystems", []))
    return found


def free_space() -> dict:
    try:
        usage = shutil.disk_usage("/home")
    except OSError:
        return {"freeBytes": 0, "totalBytes": 0}
    return {"freeBytes": usage.free, "totalBytes": usage.total}


def timeline_running() -> bool:
    return run(["systemctl", "is-active", "snapper-timeline.timer"]).stdout.strip() == "active"


def snapshot() -> dict:
    configs = []
    protected_paths = set()
    for name in config_names():
        settings = config_settings(name)
        subvolume = settings.get("SUBVOLUME", "")
        protected_paths.add(subvolume)
        configs.append({
            "name": name,
            "subvolume": subvolume,
            "timelineEnabled": settings.get("TIMELINE_CREATE", "no") == "yes",
            # Empty when this user cannot read the config at all, which is a
            # different state from "no snapshots".
            "readable": bool(settings),
            "limits": {
                "hourly": int(settings.get("TIMELINE_LIMIT_HOURLY") or 0),
                "daily": int(settings.get("TIMELINE_LIMIT_DAILY") or 0),
                "weekly": int(settings.get("TIMELINE_LIMIT_WEEKLY") or 0),
                "monthly": int(settings.get("TIMELINE_LIMIT_MONTHLY") or 0),
                "yearly": int(settings.get("TIMELINE_LIMIT_YEARLY") or 0),
            },
            "snapshots": snapshots_for(name) if settings else [],
        })
    configs.sort(key=lambda entry: entry["subvolume"])

    unprotected = [entry for entry in btrfs_subvolumes()
                   if entry["path"] not in protected_paths]

    return {
        "configs": configs,
        "unprotected": unprotected,
        "timelineRunning": timeline_running(),
        "space": free_space(),
        "error": "",
    }


def snapshot_root(config: str, number: int) -> Path:
    settings = config_settings(config)
    subvolume = settings.get("SUBVOLUME", "")
    if not subvolume:
        raise BoundaryError("That snapshot configuration cannot be read.")
    path = Path(subvolume) / ".snapshots" / str(number) / "snapshot"
    if not path.is_dir():
        raise BoundaryError("That snapshot is not available.")
    return path


def safe_relative(root: Path, relative: str) -> Path:
    """Resolve a path inside a snapshot, refusing anything that escapes it.

    The caller is a settings page passing a path a person clicked, and ".." in
    the wrong place would read or restore from outside the snapshot entirely.
    """
    candidate = (root / relative.lstrip("/")).resolve()
    if candidate != root.resolve() and root.resolve() not in candidate.parents:
        raise BoundaryError("That path is not inside the snapshot.")
    return candidate


def browse(config: str, number: str, relative: str) -> dict:
    root = snapshot_root(require_config(config), require_number(number))
    target = safe_relative(root, relative)
    if not target.is_dir():
        raise BoundaryError("That is not a folder in this snapshot.")

    entries = []
    truncated = False
    try:
        with os.scandir(target) as scan:
            for item in scan:
                if len(entries) >= BROWSE_LIMIT:
                    truncated = True
                    break
                try:
                    is_dir = item.is_dir(follow_symlinks=False)
                    size = 0 if is_dir else item.stat(follow_symlinks=False).st_size
                except OSError:
                    continue
                entries.append({"name": item.name, "directory": is_dir, "bytes": size})
    except PermissionError as error:
        raise BoundaryError("That folder cannot be read from this snapshot.") from error
    except OSError as error:
        raise BoundaryError("That folder could not be listed.") from error

    entries.sort(key=lambda entry: (not entry["directory"], entry["name"].lower()))
    return {"path": relative, "entries": entries, "truncated": truncated, "error": ""}


def restore(config: str, number: str, relative: str) -> dict:
    """Copy something out of a snapshot, keeping whatever is there now.

    The current version is moved aside rather than overwritten. A restore that
    destroys the thing you were about to compare it against is how people lose
    the work they were trying to save.
    """
    name = require_config(config)
    index = require_number(number)
    root = snapshot_root(name, index)
    source = safe_relative(root, relative)
    if not source.exists():
        raise BoundaryError("That is not in this snapshot.")

    settings = config_settings(name)
    live_root = Path(settings.get("SUBVOLUME", ""))
    destination = safe_relative(live_root, relative)

    kept = ""
    if destination.exists():
        kept = str(destination) + f".before-restore-{index}"
        suffix = 1
        while Path(kept).exists():
            suffix += 1
            kept = str(destination) + f".before-restore-{index}-{suffix}"
        try:
            os.rename(destination, kept)
        except OSError as error:
            raise BoundaryError("The current version could not be set aside.") from error

    try:
        if source.is_dir():
            shutil.copytree(source, destination, symlinks=True)
        else:
            destination.parent.mkdir(parents=True, exist_ok=True)
            shutil.copy2(source, destination, follow_symlinks=False)
    except OSError as error:
        raise BoundaryError("That could not be restored.") from error

    return {"restored": str(destination), "keptAs": kept}


def create(config: str, description: str) -> None:
    if len(description) > 200 or "\n" in description:
        raise BoundaryError("That description cannot be used.")
    result = run(["snapper", "-c", require_config(config), "create",
                  "--description", description or "manual snapshot"], timeout=120)
    if result.returncode != 0:
        raise BoundaryError(_refusal(result, "The snapshot could not be taken."))


def delete(config: str, number: str) -> None:
    result = run(["snapper", "-c", require_config(config), "delete",
                  str(require_number(number))], timeout=120)
    if result.returncode != 0:
        raise BoundaryError(_refusal(result, "That snapshot could not be removed."))


def set_retention(config: str, hourly: str, daily: str, weekly: str) -> None:
    values = []
    for label, value in (("HOURLY", hourly), ("DAILY", daily), ("WEEKLY", weekly)):
        if not str(value).isdigit() or int(value) > 999:
            raise BoundaryError("Keep counts must be whole numbers.")
        values.append(f"TIMELINE_LIMIT_{label}={int(value)}")
    result = run(["snapper", "-c", require_config(config), "set-config", *values])
    if result.returncode != 0:
        raise BoundaryError(_refusal(result, "The keep counts could not be changed."))


def set_timeline(config: str, enabled: str) -> None:
    result = run(["snapper", "-c", require_config(config), "set-config",
                  f"TIMELINE_CREATE={'yes' if enabled == 'true' else 'no'}"])
    if result.returncode != 0:
        raise BoundaryError(_refusal(result, "Automatic snapshots could not be changed."))


def _refusal(result: subprocess.CompletedProcess, fallback: str) -> str:
    text = (result.stderr or result.stdout or "").strip().splitlines()
    if not text:
        return fallback
    last = text[-1]
    if "permission" in last.lower():
        return "This account is not allowed to change that configuration."
    return last[:200]


def main(arguments: list[str]) -> int:
    try:
        if arguments == ["snapshot"]:
            print(json.dumps(snapshot(), separators=(",", ":")))
            return 0
        if len(arguments) in (3, 4) and arguments[0] == "browse":
            print(json.dumps(browse(arguments[1], arguments[2],
                                    arguments[3] if len(arguments) == 4 else ""),
                             separators=(",", ":")))
            return 0
        if len(arguments) == 4 and arguments[0] == "restore":
            outcome = restore(arguments[1], arguments[2], arguments[3])
            state = snapshot()
            state["restored"] = outcome
            print(json.dumps(state, separators=(",", ":")))
            return 0

        if len(arguments) == 3 and arguments[0] == "create":
            create(arguments[1], arguments[2])
        elif len(arguments) == 3 and arguments[0] == "delete":
            delete(arguments[1], arguments[2])
        elif len(arguments) == 5 and arguments[0] == "set-retention":
            set_retention(arguments[1], arguments[2], arguments[3], arguments[4])
        elif len(arguments) == 3 and arguments[0] == "set-timeline":
            set_timeline(arguments[1], arguments[2])
        else:
            raise BoundaryError(
                "Usage: panama-snapshots snapshot | create CONFIG DESCRIPTION | "
                "delete CONFIG NUMBER | set-retention CONFIG HOURLY DAILY WEEKLY | "
                "set-timeline CONFIG true|false | browse CONFIG NUMBER [PATH] | "
                "restore CONFIG NUMBER PATH")
    except BoundaryError as error:
        try:
            state = snapshot()
        except BoundaryError:
            state = {"configs": [], "unprotected": [], "timelineRunning": False,
                     "space": {"freeBytes": 0, "totalBytes": 0}}
        state["error"] = str(error)
        print(json.dumps(state, separators=(",", ":")))
        return 0

    print(json.dumps(snapshot(), separators=(",", ":")))
    return 0


if __name__ == "__main__":
    raise SystemExit(main(sys.argv[1:]))