#!/usr/bin/env python3

"""Crash-safe snapshots of Panama's desktop and Home preference stores.

A restore is a two-file transaction. Its fixed journal and artifacts live at
`$XDG_STATE_HOME/panama/transactions/settings-restore`; they contain no
caller-provided paths. The journal is fsynced before either destination changes
and is removed only after both replacements are durable. Every invocation
recovers an incomplete transaction before doing any other work.
"""

from __future__ import annotations

import json
import os
import re
import stat
import sys
import tempfile
import time
import fcntl
from contextlib import contextmanager
from datetime import datetime
from pathlib import Path
from typing import Any, Iterator, NoReturn


HOME = Path(os.environ.get("HOME", str(Path.home())))
CONFIG_ROOT = Path(os.environ.get("XDG_CONFIG_HOME", str(HOME / ".config")))
STATE_ROOT = Path(os.environ.get("XDG_STATE_HOME", str(HOME / ".local/state")))
SETTINGS = CONFIG_ROOT / "panama/settings.json"
HOME_STATE = STATE_ROOT / "panama/panama-home.json"
BACKUP_DIR = STATE_ROOT / "panama/backups"
TRANSACTION_PARENT = STATE_ROOT / "panama/transactions"
TRANSACTION_DIR = TRANSACTION_PARENT / "settings-restore"
JOURNAL = TRANSACTION_DIR / "journal.json"
LOCK_FILE = TRANSACTION_PARENT / "settings-backup.lock"
KEEP = 15
SNAPSHOT_RE = re.compile(r"^settings-[0-9]{8}-[0-9]{9}\.json$")
ENTITY_RE = re.compile(r"^light\.[a-z0-9_]+$")


class BackupError(RuntimeError):
    pass


def fail(message: str) -> NoReturn:
    raise BackupError(message)


def fsync_directory(path: Path) -> None:
    descriptor = os.open(path, os.O_RDONLY | os.O_DIRECTORY)
    try:
        os.fsync(descriptor)
    finally:
        os.close(descriptor)


def ensure_directory(path: Path) -> None:
    path.mkdir(parents=True, exist_ok=True)
    if path.is_symlink() or not path.is_dir():
        fail(f"{path} is not a safe directory.")


@contextmanager
def process_lock() -> Iterator[None]:
    ensure_directory(TRANSACTION_PARENT)
    if LOCK_FILE.is_symlink():
        fail("The settings transaction lock is a symbolic link.")
    descriptor = os.open(
        LOCK_FILE,
        os.O_RDWR | os.O_CREAT | getattr(os, "O_NOFOLLOW", 0),
        0o600,
    )
    try:
        os.fchmod(descriptor, 0o600)
        fcntl.flock(descriptor, fcntl.LOCK_EX)
        yield
    finally:
        fcntl.flock(descriptor, fcntl.LOCK_UN)
        os.close(descriptor)


def is_present(path: Path) -> bool:
    return path.exists() or path.is_symlink()


def require_regular(path: Path, label: str) -> None:
    if path.is_symlink():
        fail(f"{label} is a symbolic link and cannot be used safely.")
    try:
        mode = path.stat().st_mode
    except FileNotFoundError:
        fail(f"{label} is missing.")
    if not stat.S_ISREG(mode):
        fail(f"{label} is not a regular file.")


def read_json(path: Path, label: str) -> dict[str, Any]:
    require_regular(path, label)
    try:
        value = json.loads(path.read_text(encoding="utf-8"))
    except (OSError, UnicodeError, json.JSONDecodeError) as error:
        raise BackupError(f"{label} is not valid JSON.") from error
    if not isinstance(value, dict):
        fail(f"{label} is not a JSON object.")
    return value


def valid_home(value: Any) -> bool:
    if not isinstance(value, dict):
        return False
    initialized = value.get("initialized")
    favorites = value.get("favorites")
    if not isinstance(initialized, bool) or not isinstance(favorites, list):
        return False
    if not initialized and favorites:
        return False
    seen: set[str] = set()
    for favorite in favorites:
        if not isinstance(favorite, dict):
            return False
        entity_id = favorite.get("id")
        alias = favorite.get("alias")
        if (
            not isinstance(entity_id, str)
            or ENTITY_RE.fullmatch(entity_id) is None
            or not isinstance(alias, str)
            or entity_id in seen
        ):
            return False
        seen.add(entity_id)
    return True


def validate_home(value: Any, label: str) -> dict[str, Any]:
    if not valid_home(value):
        fail(f"{label} does not contain valid Home favourites.")
    return value


def json_bytes(value: Any) -> bytes:
    return (json.dumps(value, indent=2, ensure_ascii=False) + "\n").encode("utf-8")


def atomic_write_bytes(path: Path, content: bytes) -> None:
    ensure_directory(path.parent)
    if path.is_symlink():
        fail(f"{path} is a symbolic link and cannot be replaced safely.")
    descriptor, temporary_name = tempfile.mkstemp(prefix=f".{path.name}.", dir=path.parent)
    temporary = Path(temporary_name)
    try:
        os.fchmod(descriptor, 0o600)
        with os.fdopen(descriptor, "wb") as stream:
            stream.write(content)
            stream.flush()
            os.fsync(stream.fileno())
        os.replace(temporary, path)
        fsync_directory(path.parent)
    finally:
        if temporary.exists() or temporary.is_symlink():
            temporary.unlink()


def atomic_write_json(path: Path, value: Any) -> None:
    atomic_write_bytes(path, json_bytes(value))


def durable_remove(path: Path) -> None:
    if path.exists() or path.is_symlink():
        path.unlink()
        fsync_directory(path.parent)


def transaction_path(name: str) -> Path:
    if name not in {
        "journal.json",
        "desktop.old",
        "desktop.new",
        "home.old",
        "home.new",
    }:
        fail("The restore transaction contains an unknown artifact name.")
    path = TRANSACTION_DIR / name
    resolved_parent = path.parent.resolve(strict=False)
    if resolved_parent != TRANSACTION_DIR.resolve(strict=False):
        fail("The restore transaction escaped its contained state directory.")
    return path


def clean_transaction_artifacts() -> None:
    if not TRANSACTION_DIR.exists() and not TRANSACTION_DIR.is_symlink():
        return
    if TRANSACTION_DIR.is_symlink() or not TRANSACTION_DIR.is_dir():
        fail("The restore transaction path is not a safe directory.")
    for child in list(TRANSACTION_DIR.iterdir()):
        if child.name not in {
            "journal.json",
            "desktop.old",
            "desktop.new",
            "home.old",
            "home.new",
        } and not child.name.startswith(".journal.json."):
            fail("The restore transaction directory contains an unknown artifact.")
        if child.is_dir() and not child.is_symlink():
            fail("The restore transaction contains an unexpected directory.")
        child.unlink()
    fsync_directory(TRANSACTION_DIR)
    TRANSACTION_DIR.rmdir()
    fsync_directory(TRANSACTION_PARENT)


def clean_stale_atomic_files() -> None:
    locations = (
        (SETTINGS.parent, (".settings.json.",)),
        (HOME_STATE.parent, (".panama-home.json.",)),
        (BACKUP_DIR, (".settings-",)),
    )
    for directory, prefixes in locations:
        if not directory.exists():
            continue
        if directory.is_symlink() or not directory.is_dir():
            fail(f"{directory} is not a safe directory.")
        changed = False
        for child in directory.iterdir():
            if not any(child.name.startswith(prefix) for prefix in prefixes):
                continue
            # Only Panama's hidden atomic-write names are eligible. A matching
            # directory is unexpected and is never recursively removed.
            if child.is_dir() and not child.is_symlink():
                fail("A stale settings temporary path is an unexpected directory.")
            child.unlink()
            changed = True
        if changed:
            fsync_directory(directory)


def validate_journal_side(value: Any) -> dict[str, bool]:
    if not isinstance(value, dict):
        fail("The restore journal is malformed.")
    if set(value) != {"touch", "oldPresent", "newPresent"}:
        fail("The restore journal is malformed.")
    if not all(isinstance(value[key], bool) for key in value):
        fail("The restore journal is malformed.")
    return value


def read_journal() -> dict[str, Any]:
    value = read_json(JOURNAL, "The restore journal")
    if set(value) != {"version", "desktop", "home"} or value.get("version") != 1:
        fail("The restore journal uses an unsupported format.")
    return {
        "version": 1,
        "desktop": validate_journal_side(value.get("desktop")),
        "home": validate_journal_side(value.get("home")),
    }


def target_for(store: str) -> Path:
    if store == "desktop":
        return SETTINGS
    if store == "home":
        return HOME_STATE
    fail("The restore journal names an unknown store.")


def apply_artifact(store: str, generation: str, present: bool) -> None:
    target = target_for(store)
    if present:
        artifact = transaction_path(f"{store}.{generation}")
        require_regular(artifact, "A restore transaction artifact")
        atomic_write_bytes(target, artifact.read_bytes())
    else:
        ensure_directory(target.parent)
        if target.is_symlink():
            fail(f"{target} is a symbolic link and cannot be replaced safely.")
        durable_remove(target)


def recover_transaction() -> None:
    ensure_directory(TRANSACTION_PARENT)
    if not TRANSACTION_DIR.exists() and not TRANSACTION_DIR.is_symlink():
        return
    if TRANSACTION_DIR.is_symlink() or not TRANSACTION_DIR.is_dir():
        fail("The restore transaction path is not a safe directory.")
    if not JOURNAL.exists() and not JOURNAL.is_symlink():
        clean_transaction_artifacts()
        return

    journal = read_journal()
    for store in ("desktop", "home"):
        side = journal[store]
        if side["touch"]:
            apply_artifact(store, "old", side["oldPresent"])

    # Journal absence is the durable commit marker for recovery too. If a
    # second power loss occurs above, the journal remains and recovery retries.
    durable_remove(JOURNAL)
    clean_transaction_artifacts()


def is_v2_side(value: Any) -> bool:
    return (
        isinstance(value, dict)
        and isinstance(value.get("present"), bool)
        and (not value["present"] or isinstance(value.get("data"), dict))
    )


def is_v2_envelope(value: Any) -> bool:
    return (
        isinstance(value, dict)
        and value.get("version") == 2
        and is_v2_side(value.get("desktop"))
        and is_v2_side(value.get("home"))
    )


def validate_snapshot(value: dict[str, Any]) -> tuple[str, dict[str, Any]]:
    if not is_v2_envelope(value):
        return "legacy", value
    if value["home"]["present"]:
        validate_home(value["home"]["data"], "That snapshot")
    return "versioned", value


def current_store(path: Path, label: str, *, home_store: bool = False) -> tuple[bool, Any]:
    if not is_present(path):
        return False, None
    value = read_json(path, label)
    if home_store:
        validate_home(value, label)
    return True, value


def next_snapshot_path() -> Path:
    ensure_directory(BACKUP_DIR)
    while True:
        stamp = datetime.now().strftime("%Y%m%d-%H%M%S%f")[:18]
        candidate = BACKUP_DIR / f"settings-{stamp}.json"
        if not is_present(candidate):
            return candidate
        time.sleep(0.002)


def prune_snapshots() -> None:
    snapshots = sorted(
        (
            path
            for path in BACKUP_DIR.iterdir()
            if SNAPSHOT_RE.fullmatch(path.name)
            and path.is_file()
            and not path.is_symlink()
        ),
        key=lambda path: path.stat().st_mtime_ns,
        reverse=True,
    )
    for old in snapshots[KEEP:]:
        durable_remove(old)


def save_snapshot(*, require_any: bool, validate: bool) -> Path | None:
    try:
        desktop_present, desktop = current_store(
            SETTINGS, "The current settings file"
        )
        home_present, home = current_store(
            HOME_STATE, "The current Home state file", home_store=True
        )
    except BackupError:
        if validate:
            raise
        return None

    if not desktop_present and not home_present:
        if require_any:
            fail("No Panama settings exist to back up.")
        return None

    envelope: dict[str, Any] = {
        "version": 2,
        "desktop": {"present": desktop_present},
        "home": {"present": home_present},
    }
    if desktop_present:
        envelope["desktop"]["data"] = desktop
    if home_present:
        envelope["home"]["data"] = home

    destination = next_snapshot_path()
    atomic_write_json(destination, envelope)
    prune_snapshots()
    return destination


def snapshot_source(name: str) -> Path:
    if SNAPSHOT_RE.fullmatch(name) is None:
        fail("Not a snapshot name.")
    ensure_directory(BACKUP_DIR)
    candidate = BACKUP_DIR / name
    require_regular(candidate, "That snapshot")
    if candidate.resolve(strict=True).parent != BACKUP_DIR.resolve(strict=True):
        fail("That snapshot is outside the backup directory.")
    return candidate


def stage_artifact(name: str, content: bytes) -> None:
    path = transaction_path(name)
    if path.exists() or path.is_symlink():
        fail("A stale restore transaction artifact was not recovered.")
    descriptor = os.open(path, os.O_WRONLY | os.O_CREAT | os.O_EXCL, 0o600)
    try:
        with os.fdopen(descriptor, "wb") as stream:
            stream.write(content)
            stream.flush()
            os.fsync(stream.fileno())
    finally:
        # fdopen owns the descriptor after construction.
        pass
    fsync_directory(TRANSACTION_DIR)


def capture_old(store: str, target: Path) -> bool:
    if not is_present(target):
        return False
    require_regular(target, f"The current {store} settings file")
    stage_artifact(f"{store}.old", target.read_bytes())
    return True


def prepare_transaction(
    desktop_present: bool,
    desktop_data: dict[str, Any] | None,
    home_touch: bool,
    home_present: bool,
    home_data: dict[str, Any] | None,
) -> dict[str, Any]:
    # Cleanup is active before the first artifact is created. A pre-journal
    # error removes every staged/rollback file; a process death is recovered as
    # stale preparation by the next invocation.
    ensure_directory(TRANSACTION_PARENT)
    clean_transaction_artifacts()
    ensure_directory(TRANSACTION_DIR)
    fsync_directory(TRANSACTION_PARENT)
    try:
        desktop_old = capture_old("desktop", SETTINGS)
        if desktop_present:
            stage_artifact("desktop.new", json_bytes(desktop_data))
        if os.environ.get("PANAMA_SETTINGS_BACKUP_TEST_FAIL") == "after-desktop-stage":
            fail("Injected failure after desktop staging.")

        home_old = capture_old("home", HOME_STATE) if home_touch else False
        if home_touch and home_present:
            stage_artifact("home.new", json_bytes(home_data))

        journal = {
            "version": 1,
            "desktop": {
                "touch": True,
                "oldPresent": desktop_old,
                "newPresent": desktop_present,
            },
            "home": {
                "touch": home_touch,
                "oldPresent": home_old,
                "newPresent": home_present,
            },
        }
        atomic_write_json(JOURNAL, journal)
        return journal
    except BaseException:
        # SIGKILL/os._exit bypass this block by design; the next invocation
        # cleans a pre-journal directory or recovers a journalled transaction.
        if not JOURNAL.exists() and not JOURNAL.is_symlink():
            clean_transaction_artifacts()
        raise


def commit_restore(journal: dict[str, Any]) -> None:
    try:
        desktop = journal["desktop"]
        apply_artifact("desktop", "new", desktop["newPresent"])
        if os.environ.get("PANAMA_SETTINGS_BACKUP_TEST_CRASH") == "after-desktop":
            os._exit(86)

        home = journal["home"]
        if home["touch"]:
            apply_artifact("home", "new", home["newPresent"])

        # Both targets and their parent directories are durable. Removing and
        # fsyncing the journal is the transaction's commit record.
        durable_remove(JOURNAL)
        clean_transaction_artifacts()
    except BaseException:
        # Ordinary failures roll back immediately. Process death leaves the
        # journal in place and takes this same path on the next invocation.
        recover_transaction()
        raise


def write_live_home(text: str) -> None:
    try:
        value = json.loads(text)
    except json.JSONDecodeError as error:
        raise BackupError("The live Home state is not valid JSON.") from error
    validate_home(value, "The live Home state")
    atomic_write_json(HOME_STATE, value)


def command_save(arguments: list[str]) -> None:
    if arguments:
        write_live_home(arguments[0])
    destination = save_snapshot(require_any=True, validate=True)
    assert destination is not None
    print(json.dumps({"saved": destination.name}, separators=(",", ":")))


def snapshot_files() -> list[Path]:
    ensure_directory(BACKUP_DIR)
    return sorted(
        (
            path
            for path in BACKUP_DIR.iterdir()
            if SNAPSHOT_RE.fullmatch(path.name)
            and path.is_file()
            and not path.is_symlink()
        ),
        key=lambda path: path.stat().st_mtime_ns,
        reverse=True,
    )


def command_list() -> None:
    output: list[dict[str, Any]] = []
    for path in snapshot_files():
        try:
            value = read_json(path, "A snapshot")
            if is_v2_envelope(value):
                desktop = value["desktop"]
                keys = len(desktop["data"]) if desktop["present"] else 0
            else:
                keys = len(value)
        except BackupError:
            keys = 0
        raw = path.name.removeprefix("settings-").removesuffix(".json")
        pretty = (
            f"{raw[0:4]}-{raw[4:6]}-{raw[6:8]} "
            f"{raw[9:11]}:{raw[11:13]}:{raw[13:15]}"
        )
        output.append({"name": path.name, "when": pretty, "keys": keys})
    print(json.dumps(output, separators=(",", ":")))


def command_restore(arguments: list[str]) -> None:
    if not arguments:
        fail("Which snapshot?")
    name = arguments[0]
    source = snapshot_source(name)
    snapshot = read_json(source, "That snapshot")
    snapshot_format, value = validate_snapshot(snapshot)

    if snapshot_format == "versioned":
        desktop_present = value["desktop"]["present"]
        desktop_data = value["desktop"].get("data")
        home_touch = True
        home_present = value["home"]["present"]
        home_data = value["home"].get("data")
    else:
        desktop_present = True
        desktop_data = value
        home_touch = False
        home_present = False
        home_data = None

    # Restoring remains undoable, but a corrupt current file must not prevent a
    # known-good snapshot from recovering the desktop.
    save_snapshot(require_any=False, validate=False)
    journal = prepare_transaction(
        desktop_present,
        desktop_data,
        home_touch,
        home_present,
        home_data,
    )
    commit_restore(journal)

    if not home_touch:
        home_result: dict[str, Any] = {"preserve": True}
    elif home_present:
        home_result = {"present": True, "data": home_data}
    else:
        home_result = {"present": False}
    print(
        json.dumps(
            {"restored": name, "home": home_result},
            separators=(",", ":"),
        )
    )


def main() -> None:
    with process_lock():
        clean_stale_atomic_files()
        recover_transaction()
        command = sys.argv[1] if len(sys.argv) > 1 else "list"
        arguments = sys.argv[2:]
        if command == "save":
            command_save(arguments)
        elif command == "list":
            command_list()
        elif command == "restore":
            command_restore(arguments)
        else:
            fail("usage: panama-settings-backup [save|list|restore <name>]")


if __name__ == "__main__":
    try:
        main()
    except BackupError as error:
        print(str(error), file=sys.stderr)
        raise SystemExit(1) from error
    except OSError as error:
        print("The settings backup could not access its state files.", file=sys.stderr)
        raise SystemExit(1) from error