"""Root-only server deployment transactions. Commands are supplied by the Nix wrapper.

The deployment lock excludes other deployments and the WAN watchdog. A separate
systemd timer can stop a stuck worker before taking that lock and rolling back.
State and GC roots survive the invoking terminal, service and machine reboot.
"""

import argparse
import contextlib
import fcntl
import ipaddress
import json
import os
from pathlib import Path
import re
import signal
import socket
import subprocess
import sys
import time
import uuid

TERMINAL = {"committed", "rolled-back", "build-failed"}


def run(args, timeout=120):
    # A timeout must kill descendants too, not leave an activation writing in
    # parallel with recovery. systemd also kills the entire worker cgroup.
    proc = subprocess.Popen(
        args,
        stdout=subprocess.PIPE,
        stderr=subprocess.STDOUT,
        text=True,
        start_new_session=True,
    )
    try:
        output, _ = proc.communicate(timeout=timeout)
    except subprocess.TimeoutExpired:
        os.killpg(proc.pid, signal.SIGKILL)
        proc.communicate()
        raise RuntimeError(f"Command timed out: {args[0]}") from None
    if proc.returncode:
        raise RuntimeError(
            f"Command failed ({proc.returncode}): {' '.join(map(str, args))}\n{output}"
        )
    return output.strip()


def sync_directory(path):
    fd = os.open(path, os.O_RDONLY | os.O_DIRECTORY)
    try:
        os.fsync(fd)
    finally:
        os.close(fd)


def atomic_json(path, value):
    temp = path.with_suffix(".tmp")
    with temp.open("w") as out:
        json.dump(value, out)
        out.flush()
        os.fsync(out.fileno())
    os.replace(temp, path)
    sync_directory(path.parent)


class Deployment:
    def __init__(
        self,
        runner,
        state="/var/lib/fred-deploy",
        roots="/nix/var/nix/gcroots/fred-deploy",
        lock="/run/lock/fred-deploy.lock",
        command=run,
    ):
        self.runner = str(Path(runner).resolve())
        self.state, self.roots, self.lock = Path(state), Path(roots), Path(lock)
        self.command = command
        self.profile = Path("/nix/var/nix/profiles/system")
        self.current = Path("/run/current-system")
        self.state.mkdir(mode=0o700, parents=True, exist_ok=True)
        self.roots.mkdir(mode=0o700, parents=True, exist_ok=True)
        sync_directory(self.state.parent)
        sync_directory(self.roots.parent)

    @contextlib.contextmanager
    def exclusive(self):
        with self.lock.open("a") as lock:
            # Queue deployments; watchdog uses a nonblocking acquisition.
            fcntl.flock(lock, fcntl.LOCK_EX)
            yield

    def read(self):
        file = self.state / "last-result.json"
        return json.loads(file.read_text()) if file.exists() else None

    def save(self, record, phase, **fields):
        record.update(phase=phase, **fields)
        atomic_json(self.state / "last-result.json", record)
        print(f"Deployment {record['id']}: {phase}", flush=True)

    def pin(self, name, target):
        link = self.roots / name
        temporary = self.roots / (name + ".new")
        temporary.unlink(missing_ok=True)
        temporary.symlink_to(target)
        os.replace(temporary, link)
        sync_directory(self.roots)

    def cleanup(self):
        for name in ("candidate", "old-running", "old-boot", "runner"):
            (self.roots / name).unlink(missing_ok=True)

    def healthy(self):
        c = self.command
        c(
            [
                "systemctl",
                "is-active",
                "--quiet",
                "systemd-networkd",
                "dnsmasq",
                "adguardhome",
                "nginx",
                "sshd",
                "nftables",
                "router-ui",
            ],
            timeout=10,
        )
        addresses = json.loads(
            c(["ip", "-j", "-4", "address", "show", "dev", "eth0"], timeout=5)
        )
        if not any(
            a["local"] == "10.0.0.1" and a["prefixlen"] == 24
            for interface in addresses
            for a in interface["addr_info"]
        ):
            raise RuntimeError("LAN address 10.0.0.1/24 missing")
        routes = json.loads(
            c(["ip", "-j", "-4", "route", "show", "default", "dev", "eno1"], timeout=5)
        )
        if (
            not routes
            or Path("/proc/sys/net/ipv4/ip_forward").read_text().strip() != "1"
        ):
            raise RuntimeError("WAN route or IPv4 forwarding missing")
        rules = json.loads(
            c(["nft", "-j", "list", "table", "inet", "filter"], timeout=5)
        )
        chains = [r["chain"] for r in rules["nftables"] if "chain" in r]
        for hook in ("input", "forward"):
            if not any(
                ch.get("hook") == hook and ch.get("policy") == "drop" for ch in chains
            ):
                raise RuntimeError(f"Default-drop {hook} firewall missing")
        nat = json.loads(
            c(["nft", "-j", "list", "table", "ip", "router-nat"], timeout=5)
        )
        hooks = {r["chain"].get("hook") for r in nat["nftables"] if "chain" in r}
        if not {"prerouting", "postrouting"} <= hooks:
            raise RuntimeError("Router NAT hooks missing")
        expressions = [
            expr
            for item in nat["nftables"]
            if "rule" in item
            for expr in item["rule"].get("expr", [])
        ]
        if not all(
            any(action in expr for expr in expressions)
            for action in ("dnat", "masquerade")
        ):
            raise RuntimeError("Router DNAT or masquerade rules missing")
        try:
            c(["ping", "-c1", "-W3", "-I", "eno1", "1.1.1.1"], timeout=5)
        except RuntimeError:
            c(["ping", "-c1", "-W3", "-I", "eno1", "8.8.8.8"], timeout=5)
        answer = c(
            ["dig", "@10.0.0.1", "example.com", "A", "+short", "+time=2", "+tries=1"],
            timeout=5,
        )
        if not any(self.is_ip(line) for line in answer.splitlines()):
            raise RuntimeError("LAN DNS query failed")
        with socket.create_connection(("10.0.0.1", 22), timeout=3) as ssh:
            if not ssh.recv(128).startswith(b"SSH-2.0-"):
                raise RuntimeError("SSH did not return a protocol banner")
        c(
            [
                "curl",
                "--fail",
                "--silent",
                "--show-error",
                "--max-time",
                "5",
                "http://127.0.0.1:8086/",
            ],
            timeout=7,
        )
        code = c(
            [
                "curl",
                "--silent",
                "--show-error",
                "--max-time",
                "5",
                "--output",
                "/dev/null",
                "--write-out",
                "%{http_code}",
                "--resolve",
                "router.nordhammer.it:443:10.0.0.1",
                "https://router.nordhammer.it/",
            ],
            timeout=7,
        )
        if code not in {"200", "301", "302", "303", "307", "308", "401", "403"}:
            raise RuntimeError(f"Management proxy returned HTTP {code}")

    @staticmethod
    def is_ip(value):
        try:
            ipaddress.IPv4Address(value)
            return True
        except ValueError:
            return False

    def wait_healthy(self, duration=90):
        deadline = time.monotonic() + duration
        while True:
            try:
                self.healthy()
                return
            except Exception as error:
                if time.monotonic() >= deadline:
                    raise RuntimeError(
                        f"Router health checks failed: {error}"
                    ) from error
                print(f"Waiting for router health: {error}", flush=True)
                time.sleep(5)

    def disarm(self, record):
        # Do not stop the recovery service: it may already be checking state.
        self.command(
            ["systemctl", "stop", f"fred-rollback-{record['id']}.timer"], timeout=15
        )

    def rollback_locked(self, record):
        try:
            self.save(record, "rolling-back")
        except OSError as error:
            # A full/read-only state filesystem must not prevent restoration.
            print(f"Cannot record recovery start: {error}", flush=True)
        errors = []
        # Always attempt runtime recovery, even if bootloader restoration fails.
        for args in (
            ["nix-env", "--profile", str(self.profile), "--set", record["old_boot"]],
            [record["old_boot"] + "/bin/switch-to-configuration", "boot"],
            [record["old_running"] + "/bin/switch-to-configuration", "test"],
        ):
            try:
                self.command(args, timeout=120)
            except Exception as error:
                errors.append(str(error))
        try:
            self.wait_healthy()
        except Exception as error:
            errors.append(str(error))
        if errors:
            self.save(record, "rollback-failed", errors=errors)
            raise RuntimeError(
                "Recovery needs attention; GC roots retained: " + "; ".join(errors)
            )
        self.save(record, "rolled-back")
        # Cancellation failure cannot turn a completed rollback into success
        # for the deployment. A late timer sees terminal state and does nothing.
        try:
            self.disarm(record)
        except Exception as error:
            print(f"Could not cancel completed recovery timer: {error}", flush=True)
        self.cleanup()

    def recover(self, expected=None, boot=False):
        record = self.read()
        if (
            not record
            or record["phase"] in TERMINAL
            or (expected and expected != record["id"])
        ):
            return
        if boot and record.get("boot_id") == self.boot_id():
            return
        expected = record["id"]
        # Never wait behind a hung activation holding the transaction lock.
        # Timer and boot recovery execute in their own cgroups.
        try:
            self.command(["systemctl", "stop", record["unit"]], timeout=30)
        except RuntimeError:
            if (
                self.command(
                    [
                        "systemctl",
                        "show",
                        "--property=LoadState",
                        "--value",
                        record["unit"],
                    ]
                )
                != "not-found"
            ):
                raise
        with self.exclusive():
            record = self.read()
            if (
                record
                and record["phase"] not in TERMINAL
                and (not expected or expected == record["id"])
            ):
                if record["phase"] == "building":
                    self.save(
                        record,
                        "build-failed",
                        error="Builder interrupted; no activation attempted",
                    )
                    self.cleanup()
                else:
                    self.rollback_locked(record)

    @staticmethod
    def boot_id():
        return Path("/proc/sys/kernel/random/boot_id").read_text().strip()

    def deploy(self, ident, flake, reboot=False):
        with self.exclusive():
            previous = self.read()
            if (
                previous
                and previous["phase"] not in TERMINAL
                and previous["phase"] != "building"
                and previous.get("boot_id") != self.boot_id()
            ):
                # A boot-time update may win the lock before boot recovery.
                # The former worker cannot exist in this boot: restore first.
                self.rollback_locked(previous)
            if previous and previous["phase"] == "building":
                # Owning the lock proves the previous builder is gone. No
                # activation was attempted, so a fresh build is safe.
                self.save(
                    previous, "build-failed", error="Previous builder interrupted"
                )
                self.cleanup()
            if previous and previous["phase"] not in TERMINAL:
                raise RuntimeError(
                    "An interrupted deployment needs recovery first: sudo fred-deploy recover"
                )
            record = {
                "id": ident,
                "unit": f"fred-deploy-{ident}.service",
                "flake": flake,
                "boot_id": self.boot_id(),
                "old_running": str(self.current.resolve(strict=True)),
                "old_boot": str(self.profile.resolve(strict=True)),
            }
            self.pin("old-running", record["old_running"])
            self.pin("old-boot", record["old_boot"])
            self.pin("runner", str(Path(self.runner).parents[1]))
            self.save(record, "building")
            try:
                self.command(
                    [
                        "nix",
                        "build",
                        "--refresh",
                        "--out-link",
                        str(self.roots / "candidate"),
                        flake
                        + "#nixosConfigurations.FredOS-Mediaserver.config.system.build.toplevel",
                    ],
                    timeout=5400,
                )
                candidate = str((self.roots / "candidate").resolve(strict=True))
                if not Path(candidate, "bin/switch-to-configuration").is_file():
                    raise RuntimeError("Build did not produce a NixOS system")
                if not Path(
                    candidate, "etc/systemd/system/fred-deploy-recover.service"
                ).is_file():
                    raise RuntimeError("Candidate is missing the boot recovery service")
                self.pin("candidate", candidate)
            except Exception as error:
                self.save(record, "build-failed", error=str(error))
                self.cleanup()
                raise
            # Durable intent precedes timer setup; boot recovery can handle a
            # crash even between starting the timer and invoking activation.
            self.save(record, "armed", candidate=candidate)
            try:
                self.command(
                    [
                        "systemd-run",
                        "--collect",
                        "--unit",
                        f"fred-rollback-{ident}",
                        "--on-active=10min",
                        "--timer-property=AccuracySec=1s",
                        "--timer-property=RemainAfterElapse=no",
                        "--property=Type=exec",
                        "--property=RuntimeMaxSec=12min",
                        "--property=TimeoutStopSec=10s",
                        self.runner,
                        "recover",
                        "--id",
                        ident,
                    ],
                    timeout=20,
                )
                self.save(record, "activating")
                self.command(
                    [candidate + "/bin/switch-to-configuration", "test"], timeout=120
                )
                self.wait_healthy()
                self.save(record, "committing")
                self.command(
                    ["nix-env", "--profile", str(self.profile), "--set", candidate],
                    timeout=30,
                )
                self.command(
                    [candidate + "/bin/switch-to-configuration", "boot"], timeout=120
                )
                self.healthy()
                self.save(record, "committed")
            except Exception as error:
                # Timer owns recovery after SIGKILL; ordinary errors recover now.
                print(f"Deployment failed: {error}", flush=True)
                record["error"] = str(error)
                self.rollback_locked(record)
                raise RuntimeError(
                    "Deployment failed; previous runtime and boot configuration restored"
                ) from error
            try:
                self.disarm(record)
            except Exception as error:
                print(
                    f"Timer cleanup failed (committed state is durable): {error}",
                    flush=True,
                )
            self.cleanup()
            recorder = Path("/run/current-system/sw/bin/record-update")
            if recorder.exists():
                try:
                    self.command(
                        [str(recorder), record["old_running"], candidate], timeout=30
                    )
                except Exception as error:
                    print(
                        f"Deployment committed; history recording failed: {error}",
                        flush=True,
                    )
            if reboot:
                components = ("kernel", "initrd", "kernel-modules")
                if any(
                    Path("/run/booted-system", p).resolve()
                    != Path(candidate, p).resolve()
                    for p in components
                ):
                    self.command(["systemctl", "reboot", "--no-block"])


def main():
    parser = argparse.ArgumentParser(
        prog="fred-deploy",
        description="Build and safely deploy the mediaserver configuration.",
        usage="%(prog)s [--flake SOURCE] [--reboot]\n       %(prog)s recover",
    )
    parser.add_argument(
        "action",
        choices=["launch", "launch-recovery", "deploy", "recover"],
        help=argparse.SUPPRESS,
    )
    parser.add_argument(
        "--flake",
        default="git+https://forg.gregersen.it/rope/nixos",
        help="Flake source without a #fragment; defaults to the Forgejo repository",
    )
    parser.add_argument("--id", help=argparse.SUPPRESS)
    parser.add_argument(
        "--reboot",
        action="store_true",
        help="Reboot after success if boot components changed",
    )
    parser.add_argument("--boot", action="store_true", help=argparse.SUPPRESS)
    args = parser.parse_args()
    if os.geteuid() != 0:
        parser.error("run this command with sudo")
    if "#" in args.flake:
        parser.error("pass a flake source without a #fragment")
    if args.id and not re.fullmatch(r"[0-9a-f]{32}", args.id):
        parser.error("invalid transaction ID")
    runner = os.environ["FRED_DEPLOY_RUNNER"]
    if args.action in {"launch", "launch-recovery"}:
        ident = uuid.uuid4().hex
        recovering = args.action == "launch-recovery"
        unit = f"fred-recover-{ident}" if recovering else f"fred-deploy-{ident}"
        print(f"Deployment logs: journalctl -fu {unit}", flush=True)
        command = [
            "systemd-run",
            "--wait",
            "--collect",
            "--unit",
            unit,
            "--property=Type=exec",
            "--property=RuntimeMaxSec=100min",
            "--property=TimeoutStopSec=10s",
            "--property=WorkingDirectory=/",
            runner,
        ]
        command += (
            ["recover"]
            if recovering
            else ["deploy", "--id", ident, "--flake", args.flake]
        )
        if args.reboot:
            command.append("--reboot")
        result = subprocess.run(command)
        subprocess.run(["journalctl", "--no-pager", "-o", "cat", "-u", unit])
        return result.returncode
    deployment = Deployment(runner)
    if args.action == "recover":
        deployment.recover(args.id, args.boot)
    else:
        if not args.id:
            parser.error("deploy requires a transaction ID")
        deployment.deploy(args.id, args.flake, args.reboot)
    return 0


if __name__ == "__main__":
    try:
        sys.exit(main())
    except Exception as error:
        print(str(error), file=sys.stderr, flush=True)
        sys.exit(1)