#!/usr/bin/env python3
"""Guard Raspberry Pi OS /boot/firmware/config.txt health.

Priority: desktop / DRM must stay healthy. HS-idle gpio stanzas are secondary.
If applying our block would (or did) destroy Bookworm KMS settings, refuse or
restore from a golden backup on rootfs.

Golden backup lives on ext4 (/var/lib/ethergold/), not on the boot FAT, so a
vfat glitch cannot erase the recovery copy.

CLI:
  python3 boot_config_guard.py ensure-golden [--config PATH]
  python3 boot_config_guard.py apply-hs-idle [--config PATH]
  python3 boot_config_guard.py remove-hs-idle [--config PATH]
  python3 boot_config_guard.py repair [--config PATH]
  python3 boot_config_guard.py check [--config PATH]
"""

from __future__ import annotations

import argparse
import os
import sys
import tempfile
from pathlib import Path
from typing import Iterable, List, Optional, Sequence, Tuple

MARK_B = "# BEGIN DYZINETEK_HS_OUT_IDLE"
MARK_E = "# END DYZINETEK_HS_OUT_IDLE"

# Desktop/KMS must survive any of our edits. Missing any → unhealthy.
REQUIRED_SUBSTRINGS: Tuple[str, ...] = (
    "dtoverlay=vc4-kms-v3d",
    "arm_64bit=1",
)

MIN_HEALTHY_BYTES = 500

def _golden_dir() -> Path:
    # Tests/override: ETHERGOLD_BOOT_GOLDEN_DIR=/tmp/... (ext4 preferred; not boot FAT)
    override = os.environ.get("ETHERGOLD_BOOT_GOLDEN_DIR", "").strip()
    if override:
        return Path(override)
    return Path("/var/lib/ethergold")


def _golden_path() -> Path:
    return _golden_dir() / "boot-config-golden.txt"


def _golden_meta() -> Path:
    return _golden_dir() / "boot-config-golden.meta"


# Module-level aliases refreshed on each use via helpers (keep names for callers).
GOLDEN_DIR = Path("/var/lib/ethergold")
GOLDEN_PATH = GOLDEN_DIR / "boot-config-golden.txt"
GOLDEN_META = GOLDEN_DIR / "boot-config-golden.meta"

HS_IDLE_BLOCK = (
    f"{MARK_B}\n"
    "# Idle-drive OIS_OUT (BCM10) and M-S_OUT (BCM11) from firmware (before Linux).\n"
    "# Safe on older POG-COM-PI-HAT (pads unconnected). Required on Rev E+ (never high-Z).\n"
    "gpio=10=op,dl\n"
    "gpio=11=op,dl\n"
    f"{MARK_E}\n"
)


def resolve_config_path(explicit: Optional[str] = None) -> Path:
    if explicit:
        return Path(explicit)
    for candidate in (Path("/boot/firmware/config.txt"), Path("/boot/config.txt")):
        if candidate.is_file():
            return candidate
    raise FileNotFoundError("no /boot/firmware/config.txt or /boot/config.txt")


def read_text(path: Path) -> str:
    return path.read_text(encoding="utf-8", errors="replace")


def health_problems(text: str) -> List[str]:
    problems: List[str] = []
    raw = text.encode("utf-8", errors="replace")
    if len(raw) < MIN_HEALTHY_BYTES:
        problems.append(f"too_small:{len(raw)}<{MIN_HEALTHY_BYTES}")
    for needle in REQUIRED_SUBSTRINGS:
        if needle not in text:
            problems.append(f"missing:{needle}")
    # Marker pair integrity (corrupt splice)
    if (MARK_B in text) ^ (MARK_E in text):
        problems.append("hs_idle_marker_mismatch")
    return problems


def is_healthy(text: str) -> bool:
    return not health_problems(text)


def atomic_write(path: Path, text: str) -> None:
    path.parent.mkdir(parents=True, exist_ok=True)
    data = text if text.endswith("\n") else text + "\n"
    # Prefer same-dir temp so replace is atomic on the target filesystem.
    fd, tmp_name = tempfile.mkstemp(
        prefix=".config.txt.",
        suffix=".tmp",
        dir=str(path.parent),
    )
    tmp_path = Path(tmp_name)
    try:
        with os.fdopen(fd, "w", encoding="utf-8", newline="\n") as handle:
            handle.write(data)
            handle.flush()
            os.fsync(handle.fileno())
        os.replace(tmp_path, path)
        # Best-effort fsync directory (may fail on vfat; ignore)
        try:
            dir_fd = os.open(str(path.parent), os.O_RDONLY)
            try:
                os.fsync(dir_fd)
            finally:
                os.close(dir_fd)
        except OSError:
            pass
    except Exception:
        try:
            tmp_path.unlink(missing_ok=True)
        except OSError:
            pass
        raise


def save_golden(text: str, source: Path) -> Path:
    if not is_healthy(text):
        raise ValueError("refusing to save unhealthy config as golden")
    gdir = _golden_dir()
    gpath = _golden_path()
    gdir.mkdir(parents=True, exist_ok=True)
    atomic_write(gpath, text)
    meta = (
        f"source={source}\n"
        f"bytes={len(text.encode('utf-8'))}\n"
        f"required={','.join(REQUIRED_SUBSTRINGS)}\n"
    )
    atomic_write(_golden_meta(), meta)
    return gpath


def load_golden() -> Optional[str]:
    gpath = _golden_path()
    if not gpath.is_file():
        return None
    text = read_text(gpath)
    if not is_healthy(text):
        return None
    return text


def ensure_golden(config: Path) -> Tuple[str, str]:
    """Return (status, detail). status: ok|created|repaired|fatal."""
    text = read_text(config) if config.is_file() else ""
    gpath = _golden_path()
    if is_healthy(text):
        if not gpath.is_file() or not is_healthy(read_text(gpath)):
            save_golden(text, config)
            return "created", f"golden saved from healthy {config}"
        # Refresh golden when live is healthy and larger/newer content
        golden = read_text(gpath)
        if text != golden and is_healthy(text):
            save_golden(text, config)
            return "updated", f"golden refreshed from healthy {config}"
        return "ok", f"healthy {config}; golden present"
    # Unhealthy live config — restore if possible
    golden = load_golden()
    if golden is None:
        return (
            "fatal",
            f"UNHEALTHY {config} problems={health_problems(text)}; "
            f"no usable golden at {gpath}",
        )
    atomic_write(config, golden)
    return "repaired", f"restored {config} from {gpath}"


def splice_hs_idle(text: str) -> str:
    block = HS_IDLE_BLOCK
    if MARK_B in text:
        pre, rest = text.split(MARK_B, 1)
        if MARK_E in rest:
            rest = rest.split(MARK_E, 1)[1]
            if rest.startswith("\n"):
                rest = rest[1:]
        else:
            rest = ""
        text = pre.rstrip() + "\n\n" + block + rest.lstrip("\n")
    else:
        if text and not text.endswith("\n"):
            text += "\n"
        text += "\n" + block
    return text


def remove_hs_idle(text: str) -> str:
    if MARK_B not in text:
        return text
    pre, rest = text.split(MARK_B, 1)
    if MARK_E in rest:
        rest = rest.split(MARK_E, 1)[1]
        if rest.startswith("\n"):
            rest = rest[1:]
    else:
        rest = ""
    out = pre.rstrip() + ("\n" if pre.strip() else "") + rest.lstrip("\n")
    return out


def apply_hs_idle(config: Path) -> Tuple[str, str]:
    status, detail = ensure_golden(config)
    if status == "fatal":
        return status, detail
    text = read_text(config)
    if not is_healthy(text):
        return "fatal", f"post-ensure still unhealthy: {health_problems(text)}"
    new_text = splice_hs_idle(text)
    if not is_healthy(new_text):
        return (
            "fatal",
            "refusing write: spliced text would be unhealthy "
            f"{health_problems(new_text)}",
        )
    # Extra safety: never shrink below golden size by >50% (catch empty-read splice)
    golden = load_golden()
    if golden is not None:
        g_len = len(golden.encode("utf-8"))
        n_len = len(new_text.encode("utf-8"))
        if g_len >= MIN_HEALTHY_BYTES and n_len < max(MIN_HEALTHY_BYTES, g_len // 2):
            return (
                "fatal",
                f"refusing write: new size {n_len} << golden {g_len}",
            )
    atomic_write(config, new_text)
    # Verify on-disk readback
    verify = read_text(config)
    if not is_healthy(verify):
        # Emergency restore
        if golden is not None:
            atomic_write(config, golden)
            return (
                "fatal",
                "write verify failed; restored golden. "
                f"problems={health_problems(verify)}",
            )
        return "fatal", f"write verify failed; no golden. problems={health_problems(verify)}"
    save_golden(verify, config)
    return "ok", f"HS idle applied to {config}; golden updated"


def remove_hs_idle_cmd(config: Path) -> Tuple[str, str]:
    status, detail = ensure_golden(config)
    if status == "fatal":
        return status, detail
    text = read_text(config)
    new_text = remove_hs_idle(text)
    if not is_healthy(new_text):
        return (
            "fatal",
            "refusing uninstall write: result unhealthy "
            f"{health_problems(new_text)}",
        )
    atomic_write(config, new_text)
    verify = read_text(config)
    if not is_healthy(verify):
        golden = load_golden()
        if golden is not None:
            atomic_write(config, golden)
        return "fatal", "uninstall verify failed; restored golden if available"
    if is_healthy(verify):
        save_golden(verify, config)
    return "ok", f"HS idle removed from {config}"


def repair(config: Path) -> Tuple[str, str]:
    return ensure_golden(config)


def check(config: Path) -> Tuple[str, str]:
    text = read_text(config) if config.is_file() else ""
    problems = health_problems(text)
    golden_ok = load_golden() is not None
    if problems:
        return (
            "unhealthy",
            f"problems={problems} golden={'yes' if golden_ok else 'NO'} "
            f"path={_golden_path()}",
        )
    return (
        "ok",
        f"healthy bytes={len(text.encode('utf-8'))} golden={'yes' if golden_ok else 'NO'} "
        f"path={_golden_path()}",
    )


def _print_result(status: str, detail: str) -> int:
    print(f"boot_config_guard: {status}: {detail}")
    if status in ("fatal", "unhealthy"):
        return 1
    return 0


def main(argv: Optional[Sequence[str]] = None) -> int:
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument(
        "command",
        choices=("ensure-golden", "apply-hs-idle", "remove-hs-idle", "repair", "check"),
    )
    parser.add_argument("--config", default=None, help="Override config.txt path")
    args = parser.parse_args(argv)
    try:
        config = resolve_config_path(args.config)
    except FileNotFoundError as exc:
        print(f"boot_config_guard: fatal: {exc}", file=sys.stderr)
        return 1

    if args.command == "ensure-golden":
        return _print_result(*ensure_golden(config))
    if args.command == "apply-hs-idle":
        return _print_result(*apply_hs_idle(config))
    if args.command == "remove-hs-idle":
        return _print_result(*remove_hs_idle_cmd(config))
    if args.command == "repair":
        return _print_result(*repair(config))
    if args.command == "check":
        return _print_result(*check(config))
    return 2


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