#!/usr/bin/env python3
"""felhom-release-copy — copy the released images and agent packages off DooPlex to ep0, nightly (D33, R-1227).

Runs on DooPlex as root from felhom-release-copy.timer. Runbook: documentation/runbooks/total-loss-of-dooplex.md
(„Restore the release copy"). Pinned by test_release_copy.py.

WHAT. The newest KEEP release tags (X.Y.Z only) of each image in REPOS, and the newest AGENT_KEEP agent versions (the
binary and its config bundle). After a DooPlex loss these are what every box update, reinstall and new install pulls.

HOW, and why skopeo, not `docker save`. Each image tag is copied with `skopeo copy --all --preserve-digests` into
skopeo's `dir:` layout: the registry's own manifest bytes and compressed layer blobs, unchanged. Pushed back into any
registry, the image has THE SAME DIGEST as the release (a `docker save` tar is re-compressed on the next push, so the
digest — what the „a version is pushed once" rule and the hub's spoiled-tag list are about — would differ). It also
needs no Docker daemon on DooPlex.

CHECKS, before anything leaves: the sha256 of the copied top manifest equals the registry's digest; every blob and
child manifest hashes to its own name; every blob a manifest names is present; each agent file's sha256 equals the
registry's record. The staged tree must be ≤ CAP_MB (the cap — the sender enforces it; ep0 has no quota). Then
`rsync --delete` to ep0 (the limited user `felhom-relcopy`, rrsync-confined to one directory), and a second
`rsync --checksum --dry-run` must report NO difference. Only then is the success signal written (CLAUDE.md „presence
is not success"). Versions beyond KEEP / AGENT_KEEP are deleted locally, and so on ep0 by `--delete`.
"""
import base64
import hashlib
import json
import os
import re
import shutil
import subprocess
import sys
import tempfile
import time
import urllib.request

CONF = os.environ.get("FELHOM_RC_CONF", "/etc/felhom-release-copy")
STATE = os.environ.get("FELHOM_RC_STATE", "/var/lib/felhom-release-copy")
TEXTFILE_DIR = os.environ.get("FELHOM_RC_TEXTFILE_DIR", "/var/lib/node_exporter/textfile_collector")
REGISTRY = "gitea.dooplex.hu"
OWNER = "admin"
SPOILED = {"felhom-controller": {"0.315.0", "0.316.3"}}  # withdrawn builds (R-1085) — never offered, never copied
AGENT_FILES = ("felhom-agent", "felhom-config-bundle.json")
HEX64 = re.compile(r"^[0-9a-f]{64}$")
RELEASE = re.compile(r"^\d+\.\d+\.\d+$")


class Refused(Exception):
    pass


def log(msg):
    print("felhom-release-copy: " + msg, flush=True)


def load_env(path):
    env = {}
    for line in open(path, encoding="utf-8"):
        line = line.strip()
        if line and not line.startswith("#") and "=" in line:
            k, v = line.split("=", 1)
            env[k.strip()] = v.strip().strip("'\"")
    return env


def newest(versions, n, spoiled=()):
    rel = [v for v in versions if RELEASE.match(v) and v not in spoiled]
    rel.sort(key=lambda v: tuple(int(x) for x in v.split(".")))
    return rel[-n:] if n > 0 else []


def sha256_file(path):
    h = hashlib.sha256()
    with open(path, "rb") as f:
        for chunk in iter(lambda: f.read(1 << 20), b""):
            h.update(chunk)
    return h.hexdigest()


def verify_image_dir(d, want_digest):
    """Refuses unless the dir: copy in d is exactly the image want_digest names (see the module docstring)."""
    top = os.path.join(d, "manifest.json")
    if not os.path.isfile(top):
        raise Refused(f"{d}: no manifest.json")
    got = "sha256:" + sha256_file(top)
    if got != want_digest:
        raise Refused(f"{d}: manifest digest {got} != registry {want_digest}")
    names = set(os.listdir(d))
    manifests = [top]
    for n in sorted(names):
        p = os.path.join(d, n)
        if HEX64.match(n):
            if sha256_file(p) != n:
                raise Refused(f"{d}: blob {n} does not hash to its name")
        elif n.endswith(".manifest.json") and HEX64.match(n[: -len(".manifest.json")]):
            if sha256_file(p) != n[: -len(".manifest.json")]:
                raise Refused(f"{d}: child manifest {n} does not hash to its name")
            manifests.append(p)
    for m in manifests:
        doc = json.load(open(m, encoding="utf-8"))
        for ref in doc.get("manifests", []):
            hx = ref["digest"].split(":", 1)[1]
            if hx + ".manifest.json" not in names:
                raise Refused(f"{d}: child manifest {hx[:12]} named by the index is missing")
        for ref in ([doc["config"]] if "config" in doc else []) + doc.get("layers", []):
            hx = ref["digest"].split(":", 1)[1]
            if hx not in names:
                raise Refused(f"{d}: blob {hx[:12]} named by a manifest is missing")


def tree_bytes(root):
    total = 0
    for dp, _, fs in os.walk(root):
        for f in fs:
            total += os.lstat(os.path.join(dp, f)).st_size
    return total


def prune(root, keep):
    """Removes every directory under root whose name is not in keep; returns the removed names."""
    gone = []
    if os.path.isdir(root):
        for n in sorted(os.listdir(root)):
            if n not in keep:
                shutil.rmtree(os.path.join(root, n))
                gone.append(n)
    return gone


class Registry:
    def __init__(self, user, token):
        self.auth = "Basic " + base64.b64encode(f"{user}:{token}".encode()).decode()

    def get(self, url, raw=False):
        r = urllib.request.urlopen(urllib.request.Request(url, headers={"Authorization": self.auth}), timeout=120)
        return r.read() if raw else json.load(r)

    def agent_versions(self):
        out, page = [], 1
        while True:
            batch = self.get(f"https://{REGISTRY}/api/v1/packages/{OWNER}?type=generic&q=felhom-agent&limit=50&page={page}")
            if not batch:
                return out
            out += [p["version"] for p in batch if p.get("name") == "felhom-agent"]
            page += 1

    def agent_files(self, ver):
        return {f["name"]: f["sha256"] for f in self.get(
            f"https://{REGISTRY}/api/v1/packages/{OWNER}/generic/felhom-agent/{ver}/files")}

    def agent_download(self, ver, name, dest):
        with open(dest, "wb") as f:
            f.write(self.get(f"https://{REGISTRY}/api/packages/{OWNER}/generic/felhom-agent/{ver}/{name}", raw=True))


def run(cmd, **kw):
    return subprocess.run(cmd, check=True, text=True, capture_output=True, **kw)


def skopeo(authfile, *args):
    return run(["skopeo", *args[:1], "--authfile", authfile, *args[1:]]).stdout


def copy_images(reg_authfile, stage, repos, keep):
    images = os.path.join(stage, "images")
    os.makedirs(images, exist_ok=True)
    copied = []
    for repo in repos:
        tags = json.loads(skopeo(reg_authfile, "list-tags", f"docker://{REGISTRY}/{OWNER}/{repo}"))["Tags"]
        want = newest(tags, keep, SPOILED.get(repo, ()))
        if not want:
            raise Refused(f"{repo}: no release tag")
        rdir = os.path.join(images, repo)
        os.makedirs(rdir, exist_ok=True)
        for tag in want:
            ref = f"docker://{REGISTRY}/{OWNER}/{repo}:{tag}"
            digest = "sha256:" + hashlib.sha256(
                skopeo(reg_authfile, "inspect", "--raw", ref).encode()).hexdigest()
            d = os.path.join(rdir, tag)
            if os.path.isfile(os.path.join(d, "DIGEST")) and open(os.path.join(d, "DIGEST")).read().strip() == digest:
                verify_image_dir(d, digest)  # already here — re-checked every night, not trusted
                continue
            tmp = d + ".tmp"
            shutil.rmtree(tmp, ignore_errors=True)
            shutil.rmtree(d, ignore_errors=True)
            skopeo(reg_authfile, "copy", "--all", "--preserve-digests", "--quiet", ref, f"dir:{tmp}")
            verify_image_dir(tmp, digest)
            with open(os.path.join(tmp, "DIGEST"), "w") as f:
                f.write(digest + "\n")
            os.rename(tmp, d)
            copied.append(f"{repo}:{tag}")
        gone = prune(rdir, set(want))
        log(f"{repo}: keeping {', '.join(want)}" + (f"; removed {', '.join(gone)}" if gone else ""))
    prune(images, set(repos))
    return copied


def copy_agents(reg, stage, keep):
    adir = os.path.join(stage, "agent")
    os.makedirs(adir, exist_ok=True)
    want = newest(reg.agent_versions(), keep)
    if not want:
        raise Refused("felhom-agent: no release version")
    for ver in want:
        d = os.path.join(adir, ver)
        os.makedirs(d, exist_ok=True)
        recorded = reg.agent_files(ver)
        for name in AGENT_FILES:
            if name not in recorded:
                raise Refused(f"felhom-agent {ver}: the registry has no {name}")
            p = os.path.join(d, name)
            if not (os.path.isfile(p) and sha256_file(p) == recorded[name]):
                reg.agent_download(ver, name, p + ".tmp")
                if sha256_file(p + ".tmp") != recorded[name]:
                    os.remove(p + ".tmp")
                    raise Refused(f"felhom-agent {ver}/{name}: the download does not match the registry's sha256")
                os.rename(p + ".tmp", p)
            with open(os.path.join(d, name + ".sha256"), "w") as f:
                f.write(f"{recorded[name]}  {name}\n")
    gone = prune(adir, set(want))
    log(f"felhom-agent: keeping {', '.join(want)}" + (f"; removed {', '.join(gone)}" if gone else ""))


def write_sums(stage):
    lines = []
    for dp, _, fs in os.walk(stage):
        for f in fs:
            if f == "SHA256SUMS":
                continue
            p = os.path.join(dp, f)
            lines.append(f"{sha256_file(p)}  {os.path.relpath(p, stage)}")
    with open(os.path.join(stage, "SHA256SUMS"), "w") as f:
        f.write("\n".join(sorted(lines, key=lambda l: l[66:])) + "\n")


def rsync_to(env, stage, dry_check=False):
    ssh = (f"ssh -i {CONF}/id_ed25519 -o UserKnownHostsFile={CONF}/known_hosts -o StrictHostKeyChecking=yes "
           f"-o BatchMode=yes -o ConnectTimeout=30")
    cmd = ["rsync", "-a", "--delete", "-e", ssh]
    if dry_check:
        cmd += ["--checksum", "--dry-run", "--itemize-changes"]
    return run(cmd + [stage.rstrip("/") + "/", env["REMOTE"] + ":./"]).stdout


def signal(nbytes):
    tmp = os.path.join(TEXTFILE_DIR, f"felhom_release_copy.prom.{os.getpid()}")
    with open(tmp, "w") as f:
        f.write("# HELP felhom_release_copy_last_success_timestamp_seconds Last verified copy of the releases to ep0 (D33).\n"
                "# TYPE felhom_release_copy_last_success_timestamp_seconds gauge\n"
                f"felhom_release_copy_last_success_timestamp_seconds {int(time.time())}\n"
                f"felhom_release_copy_last_success_bytes {nbytes}\n")
    os.chmod(tmp, 0o644)
    os.rename(tmp, os.path.join(TEXTFILE_DIR, "felhom_release_copy.prom"))


def main():
    env = load_env(os.path.join(CONF, "env"))
    keep, akeep, cap = int(env.get("KEEP", "3")), int(env.get("AGENT_KEEP", "5")), int(env["CAP_MB"])
    repos = env["REPOS"].split()
    token = open(os.path.join(CONF, "registry-token"), encoding="utf-8").read().strip()
    user = env.get("REGISTRY_USER", OWNER)
    stage = os.path.join(STATE, "stage")
    os.makedirs(stage, exist_ok=True)
    os.chmod(STATE, 0o700)
    with tempfile.TemporaryDirectory(dir=STATE) as td:
        authfile = os.path.join(td, "auth.json")
        fd = os.open(authfile, os.O_WRONLY | os.O_CREAT | os.O_EXCL, 0o600)
        with os.fdopen(fd, "w") as f:
            json.dump({"auths": {REGISTRY: {"auth": base64.b64encode(f"{user}:{token}".encode()).decode()}}}, f)
        start = time.time()
        copied = copy_images(authfile, stage, repos, keep)
    copy_agents(Registry(user, token), stage, akeep)
    for n in os.listdir(stage):
        if n not in ("images", "agent", "SHA256SUMS", "COPIED_AT"):
            p = os.path.join(stage, n)
            shutil.rmtree(p) if os.path.isdir(p) else os.remove(p)
    with open(os.path.join(stage, "COPIED_AT"), "w") as f:
        f.write(time.strftime("%Y-%m-%dT%H:%M:%SZ", time.gmtime()) + "\n")
    write_sums(stage)
    nbytes = tree_bytes(stage)
    if nbytes > cap * 1024 * 1024:
        raise Refused(f"the copy is {nbytes // (1024 * 1024)} MB, above the cap {cap} MB — nothing sent")
    log(f"staged and verified: {nbytes // (1024 * 1024)} MB, new this run: {', '.join(copied) or 'none'}")
    rsync_to(env, stage)
    diff = rsync_to(env, stage, dry_check=True).strip()
    if diff:
        raise Refused("ep0's copy differs from the staged one after the push: " + diff.splitlines()[0])
    log(f"pushed to ep0 and compared by checksum: no difference ({int(time.time() - start)} s)")
    signal(nbytes)
    log("success signal written")


if __name__ == "__main__":
    try:
        main()
    except Refused as e:
        print(f"felhom-release-copy: FAILED: {e}", file=sys.stderr)
        sys.exit(1)
    except subprocess.CalledProcessError as e:
        print(f"felhom-release-copy: FAILED: {e.cmd[0]} {e.cmd[1] if len(e.cmd) > 1 else ''} rc={e.returncode}: "
              f"{(e.stderr or '').strip()[:300]}", file=sys.stderr)
        sys.exit(1)