#!/usr/bin/env python3
"""resolve-images — what version is each of a recipe's images on, and what is newest?

An abra-independent version resolver. `abra recipe upgrade` is the normal path, but it has a hard
failure mode: an image pinned with BOTH a tag and a digest makes it FATA and abandon the WHOLE
recipe — even images it already parsed. immich pins two that way:

    ghcr.io/immich-app/postgres:14-vectorchord0.4.3-pgvectors0.2.0@sha256:bcf6…
    docker.io/valkey/valkey:9@sha256:3acc…

so immich contributes NO version data at all and silently drops out of every survey. That is
indistinguishable from "up to date" unless a human notices the missing row — which is exactly how it
kept getting skipped, and why a CVE sweep reported it as unknown.

This reads the compose files directly and queries the registries itself, so a digest pin is just a
digest pin. Output is JSON (default) or a table.

    resolve-images.py <recipe> [--ssh HOST] [--recipe-dir DIR] [--table] [--only IMAGE]

The cc-ci host has no python3, so `--ssh cc-ci` reads the compose files from that host's checkout
over ssh and does the resolving locally. That keeps the source of truth the SAME tree abra and CI
use, rather than a second copy that can drift.

TAG SHAPES. Registries mix wildly different tag conventions in one repo, so "newest" is meaningless
without a shape. Each tag is reduced to a signature by replacing digit runs with '#':

    v3.1.0                              -> v#.#.#
    1.27.1-rootless                     -> #.#.#-rootless
    8.10-alpine                         -> #.#-alpine
    14-vectorchord0.4.3-pgvectors0.2.0  -> #-vectorchord#.#.#-pgvectors#.#.#

Only tags sharing the CURRENT pin's shape are candidates. That keeps `-alpine` on `-alpine`, and
stops a `latest`/`release`/`sha-…` tag from ever being proposed as an upgrade.

TWO ANSWERS, NOT ONE. It reports `newest_same_shape` AND `newest_within_major` (same leading number).
For a plain app image they usually agree. For a compatibility-pinned sidecar they do not, and taking
the max would be wrong: immich's postgres tag encodes the pg major plus the vectorchord/pgvectors
versions that immich-server is built against, so jumping pg major because a newer tag exists breaks
the deployment. The caller picks; this tool refuses to guess and shows both.
"""

from __future__ import annotations

import argparse
import glob
import shlex
import subprocess
import gzip
import json
import os
import re
import sys
import time
import urllib.error
import urllib.parse
import urllib.request

UA = "cc-ci-resolve-images (+https://git.autonomic.zone/recipe-maintainers/cc-ci)"
TIMEOUT = int(os.environ.get("RESOLVE_IMAGES_TIMEOUT", "45"))
RECIPE_DIR = os.environ.get("ABRA_RECIPE_DIR", os.path.expanduser("~/.abra/recipes"))
MAX_TAG_PAGES = int(os.environ.get("RESOLVE_IMAGES_MAX_PAGES", "40"))

IMAGE_RE = re.compile(r"""^\s*image:\s*["']?([^"'\s]+)["']?\s*$""", re.M)


RETRIES = int(os.environ.get("RESOLVE_IMAGES_RETRIES", "4"))


def _fetch(url: str, headers: dict | None = None) -> bytes:
    """GET with backoff on rate limits.

    Docker Hub throttles anonymous clients hard, and a sweep re-reads the same popular repos
    (nginx, redis, postgres) for recipe after recipe. A 429 mid-sweep used to surface as
    'unresolved', which is indistinguishable from a real lookup failure — so retry, and let the
    per-repo cache below remove most of the requests entirely."""
    h = {"User-Agent": UA, "Accept-Encoding": "gzip"}
    h.update(headers or {})
    delay = 2.0
    for attempt in range(RETRIES):
        try:
            with urllib.request.urlopen(urllib.request.Request(url, headers=h), timeout=TIMEOUT) as r:
                raw = r.read()
                if r.headers.get("Content-Encoding") == "gzip":
                    raw = gzip.decompress(raw)
                return raw
        except urllib.error.HTTPError as e:
            if e.code in (429, 503) and attempt < RETRIES - 1:
                time.sleep(delay)
                delay *= 2
                continue
            raise
    raise RuntimeError("unreachable")


_HUB_JWT: list = []


def _hub_auth() -> dict:
    """Authenticated Docker Hub calls get a far higher rate limit than anonymous ones.

    Credentials come from /srv/cc-ci/.testenv (DOCKERHUB_USERNAME / DOCKERHUB_TOKEN), the same pair
    the CI host already uses. Absent creds are fine — the sweep just runs anonymous and slower."""
    if _HUB_JWT:
        return _HUB_JWT[0]
    env = {}
    try:
        for ln in open(os.environ.get("CCCI_TESTENV", "/srv/cc-ci/.testenv")):
            if "=" in ln and not ln.strip().startswith("#"):
                k, v = ln.strip().split("=", 1)
                env[k] = v.strip().strip("\"'")
    except OSError:
        pass
    u = os.environ.get("DOCKERHUB_USERNAME") or env.get("DOCKERHUB_USERNAME")
    t = os.environ.get("DOCKERHUB_TOKEN") or env.get("DOCKERHUB_TOKEN")
    hdrs = {}
    if u and t:
        try:
            body = json.dumps({"username": u, "password": t}).encode()
            req = urllib.request.Request("https://hub.docker.com/v2/users/login",
                                         data=body, method="POST",
                                         headers={"Content-Type": "application/json", "User-Agent": UA})
            with urllib.request.urlopen(req, timeout=TIMEOUT) as r:
                tokj = json.load(r).get("token")
            if tokj:
                hdrs = {"Authorization": f"JWT {tokj}"}
        except Exception:  # noqa: BLE001 — anonymous is a valid fallback
            hdrs = {}
    _HUB_JWT.append(hdrs)
    return hdrs


def _json(url: str, headers: dict | None = None):
    return json.loads(_fetch(url, headers))


def shape(tag: str) -> str:
    """Signature of a tag with every digit run replaced by '#'. See module docstring."""
    return re.sub(r"\d+", "#", tag)


def vkey(tag: str) -> tuple:
    """Ordering key: every number in the tag, in order. '1.27.10' > '1.27.9'; text ignored."""
    return tuple(int(x) for x in re.findall(r"\d+", tag))


def parse_ref(ref: str) -> dict:
    """Split an image reference into registry / repo / tag / digest."""
    digest = None
    if "@" in ref:
        ref, _, digest = ref.partition("@")
    host, repo, tag = "docker.io", ref, "latest"
    # A leading component is a REGISTRY only when there is a path after it. Without the slash test,
    # a bare `postgres:15.18` looks like host "postgres:15.18" because of the tag's colon — which
    # silently sent every library image to a nonexistent registry.
    if "/" in ref:
        first = ref.split("/")[0]
        if "." in first or ":" in first or first == "localhost":
            host, _, repo = ref.partition("/")
    if ":" in repo.split("/")[-1]:
        repo, _, tag = repo.rpartition(":")
    if host == "docker.io" and "/" not in repo:
        repo = f"library/{repo}"          # bare `redis` is really `library/redis`
    return {"registry": host, "repo": repo, "tag": tag, "digest": digest}


HUB_RECENT_PAGES = int(os.environ.get("RESOLVE_IMAGES_HUB_PAGES", "10"))


def _hub_tag_exists(repo: str, tag: str) -> bool:
    try:
        _json(f"https://hub.docker.com/v2/repositories/{repo}/tags/{tag}", _hub_auth())
        return True
    except Exception:  # noqa: BLE001
        return False


def _hub_tags(repo: str) -> list[str]:
    """Recently-pushed tags, newest first.

    Popular Docker Hub repos carry many thousands of tags, so a full enumeration is impractical —
    but it is also unnecessary: a tag NEWER than the one we run must have been pushed AFTER it, so
    ordering by last_updated and reading a bounded recent window is sufficient to find any upgrade.
    (ghcr offers no ordering, which is why that path needs a different strategy.)
    """
    tags, url = [], (f"https://hub.docker.com/v2/repositories/{repo}/tags"
                     f"?page_size=100&ordering=last_updated")
    auth = _hub_auth()
    for _ in range(HUB_RECENT_PAGES):
        d = _json(url, auth)
        tags += [r["name"] for r in d.get("results", [])]
        url = d.get("next")
        if not url:
            break
    return tags


def _oci_bearer(host: str, repo: str) -> dict:
    """Token for an OCI registry, discovered from its own auth challenge.

    Registries do NOT share a token endpoint. ghcr answers at /token?scope=…&service=ghcr.io, but
    lscr.io and dock.mau.dev advertise different realms, and assuming ghcr's shape made both 401 —
    which then read as "could not resolve" rather than "asked the wrong URL". The registry tells us
    where to go in its WWW-Authenticate header; use that."""
    try:
        urllib.request.urlopen(
            urllib.request.Request(f"https://{host}/v2/{repo}/tags/list?n=1",
                                   headers={"User-Agent": UA}), timeout=TIMEOUT)
        return {}                                     # no auth needed
    except urllib.error.HTTPError as e:
        if e.code != 401:
            return {}
        chal = e.headers.get("WWW-Authenticate", "") or ""
    except Exception:  # noqa: BLE001
        return {}
    if not chal.lower().startswith("bearer"):
        return {}
    parts = dict(re.findall(r'(\w+)="([^"]*)"', chal))
    realm = parts.get("realm")
    if not realm:
        return {}
    q = {"service": parts.get("service", host), "scope": parts.get("scope", f"repository:{repo}:pull")}
    url = realm + ("&" if "?" in realm else "?") + urllib.parse.urlencode(q)
    try:
        tok = (_json(url) or {}).get("token") or (_json(url) or {}).get("access_token")
        return {"Authorization": f"Bearer {tok}"} if tok else {}
    except Exception:  # noqa: BLE001
        return {}


def _oci_tags(host: str, repo: str) -> list[str]:
    """Tags from any OCI/v2 registry, with challenge-derived auth and Link pagination.

    ghcr paginates hard — immich-server has >40,000 tags — and a truncated listing silently hides
    the newest release line, so follow the cursor and let the caller's integrity check catch a read
    that never reached the current pin."""
    hdrs = _oci_bearer(host, repo)
    tags, url = [], f"https://{host}/v2/{repo}/tags/list?n=1000"
    for _ in range(MAX_TAG_PAGES):
        req = urllib.request.Request(url, headers={"User-Agent": UA, **hdrs})
        with urllib.request.urlopen(req, timeout=TIMEOUT) as r:
            tags += (json.load(r) or {}).get("tags") or []
            link = r.headers.get("Link", "") or ""
        m = re.search(r'<([^>]+)>;\s*rel="next"', link)
        if not m:
            break
        nxt = m.group(1)
        url = f"https://{host}{nxt}" if nxt.startswith("/") else nxt
    return tags


def _gh_token() -> str | None:
    tok = os.environ.get("GITHUB_TOKEN")
    if tok:
        return tok.strip()
    try:
        return open(os.environ.get("GITHUB_TOKEN_FILE", "/srv/cc-ci/.github-token")).read().strip() or None
    except OSError:
        return None


def github_release_tags(owner: str, repo: str, max_pages: int = 4) -> list[str]:
    """Release tag names for a GitHub repo, newest first.

    FALLBACK for registries whose tag listing cannot be enumerated. ghcr has no ordering and no
    server-side filter, and immich-machine-learning carries >40,000 tags — a full read is impractical
    and a partial read silently hides the newest release line. The project's RELEASES are ordered,
    small, and authoritative: container tags track them. (The GitHub Packages API would answer this
    directly but needs a scoped token; this scan's token deliberately has none.)
    """
    hdrs = {"Accept": "application/vnd.github+json"}
    tok = _gh_token()
    if tok:
        hdrs["Authorization"] = f"Bearer {tok}"
    out = []
    for page in range(1, max_pages + 1):
        try:
            rows = _json(f"https://api.github.com/repos/{owner}/{repo}/releases"
                         f"?per_page=100&page={page}", hdrs)
        except Exception:  # noqa: BLE001
            break
        if not rows:
            break
        out += [r.get("tag_name") or "" for r in rows]
    return [t for t in out if t]


def _release_fallback_repos(registry: str, repo: str) -> list[tuple[str, str]]:
    """Candidate GitHub repos whose releases track this image's tags."""
    if "ghcr.io" not in registry:
        return []
    parts = repo.split("/")
    if len(parts) < 2:
        return []
    owner, name = parts[0], parts[-1]
    cands = [(owner, name)]
    # ghcr.io/immich-app/immich-machine-learning is built from immich-app/immich.
    if name.startswith(owner.split("-")[0]):
        cands.append((owner, owner.split("-")[0]))
    return cands


_TAG_CACHE: dict[tuple[str, str], tuple[list[str], str | None]] = {}


def list_tags(registry: str, repo: str) -> tuple[list[str], str | None]:
    if (registry, repo) in _TAG_CACHE:
        return _TAG_CACHE[(registry, repo)]
    res = _list_tags_uncached(registry, repo)
    _TAG_CACHE[(registry, repo)] = res
    return res


def _list_tags_uncached(registry: str, repo: str) -> tuple[list[str], str | None]:
    try:
        return (_hub_tags(repo) if registry in ("docker.io", "registry-1.docker.io")
                else _oci_tags(registry, repo)), None
    except urllib.error.HTTPError as e:
        return [], f"HTTP {e.code}"
    except Exception as e:  # noqa: BLE001
        return [], f"{type(e).__name__}: {e}"


def resolve(ref: str) -> dict:
    """Current pin -> newest same-shape tag, and newest within the current major."""
    if "${" in ref or "$(" in ref:
        # The tag is a compose variable (ghost pins `ghost:${IMAGE_VERSION}-alpine`). Its real value
        # lives in .env, not here. Report it as skipped, never as a failed lookup.
        return {**parse_ref(ref), "ref": ref, "shape": None, "candidates": 0,
                "newest_same_shape": None, "newest_within_major": None,
                "upgrade_available": False,
                "status": "skipped: templated ref (tag comes from a compose variable)"}
    info = parse_ref(ref)
    out = {**info, "ref": ref, "shape": shape(info["tag"]), "status": "ok",
           "newest_same_shape": None, "newest_within_major": None, "candidates": 0,
           "upgrade_available": False}
    tags, err = list_tags(info["registry"], info["repo"])
    if err:
        out["status"] = f"error: {err}"
        return out
    out["tags_seen"] = len(set(tags))
    # INTEGRITY CHECK: the tag we are currently running MUST appear in the listing. If it does not,
    # the listing is incomplete and any "newest" derived from it is a guess — ghcr paginates to tens
    # of thousands of tags and a truncated read silently hides whole release lines. immich's
    # machine-learning image is pinned v3.1.0, which EXISTS, yet a short read reported v1.134.0 as
    # newest; without this check that becomes a confident, wrong answer.
    if info["tag"] not in set(tags):
        # Docker Hub: the window is recency-ordered, so the pin being outside it just means the pin
        # is old — which is fine, because anything NEWER is necessarily inside the window. Confirm
        # the pin genuinely exists (so a typo is still caught) and carry on.
        if info["registry"] in ("docker.io", "registry-1.docker.io") and _hub_tag_exists(info["repo"], info["tag"]):
            out["source"] = f"docker-hub:recent-{HUB_RECENT_PAGES * 100}"
            tags = list(tags) + [info["tag"]]
        else:
            for owner, name in _release_fallback_repos(info["registry"], info["repo"]):
                rel = github_release_tags(owner, name)
                if info["tag"] in rel:
                    tags = rel
                    out["source"] = f"github-releases:{owner}/{name}"
                    out["tags_seen"] = len(set(rel))
                    break
            else:
                out["status"] = ("error: tag listing incomplete — the current pin "
                                 f"{info['tag']!r} is absent from {len(set(tags))} registry tags "
                                 f"and from the project's GitHub releases")
                return out
    want, cur = out["shape"], vkey(info["tag"])
    same = [t for t in set(tags) if shape(t) == want and vkey(t)]
    out["candidates"] = len(same)
    if not same:
        # Not a failure: digest-only pins and `latest`/`stable` have no comparable siblings.
        out["status"] = "no comparable tags (shape has no numeric siblings)"
        return out
    newest = max(same, key=vkey)
    out["newest_same_shape"] = newest
    if cur:
        within = [t for t in same if vkey(t)[:1] == cur[:1]]
        if within:
            out["newest_within_major"] = max(within, key=vkey)
    out["upgrade_available"] = bool(cur and vkey(newest) > cur)
    return out


def compose_images_ssh(recipe: str, host: str, recipe_dir: str) -> list[str]:
    """Same as compose_images, but the recipe tree lives on another host (cc-ci has no python3)."""
    # NB: no shell-quoting of the directory — it may legitimately start with ~ or $HOME, and
    # quoting it stops the remote shell expanding it, which yields an empty (and silent) result.
    d = f"{recipe_dir}/{shlex.quote(recipe)}".replace("~", "$HOME")
    cmd = (f'for f in {d}/compose*.yml; do case "$f" in *compose.ccci.yml) continue;; esac; '
           f'[ -f "$f" ] && {{ cat "$f"; echo; }}; done; exit 0')
    out = subprocess.run(["ssh", host, cmd], capture_output=True, text=True, timeout=120)
    if out.returncode != 0:
        raise RuntimeError(f"ssh {host}: {(out.stderr.strip() or 'no output')[:200]}")
    if not out.stdout.strip():
        raise RuntimeError(f"ssh {host}: no compose files found under {d}")
    refs = []
    for m in IMAGE_RE.finditer(out.stdout):
        if m.group(1) not in refs:
            refs.append(m.group(1))
    return refs


def compose_images(recipe: str, recipe_dir: str) -> list[str]:
    """Every `image:` ref in the recipe's own compose files (the cc-ci overlay is NOT the recipe)."""
    refs, base = [], os.path.join(recipe_dir, recipe)
    for path in sorted(glob.glob(os.path.join(base, "compose*.yml"))):
        if os.path.basename(path) == "compose.ccci.yml":
            continue
        try:
            for m in IMAGE_RE.finditer(open(path).read()):
                if m.group(1) not in refs:
                    refs.append(m.group(1))
        except OSError:
            continue
    return refs


def main() -> int:
    ap = argparse.ArgumentParser(description=__doc__,
                                 formatter_class=argparse.RawDescriptionHelpFormatter)
    ap.add_argument("recipe")
    ap.add_argument("--ssh", default=None, metavar="HOST",
                    help="read the recipe's compose files from HOST over ssh (e.g. --ssh cc-ci); "
                         "resolving still happens locally")
    ap.add_argument("--recipe-dir", default=RECIPE_DIR)
    ap.add_argument("--table", action="store_true", help="human-readable table instead of JSON")
    ap.add_argument("--only", default=None, help="resolve just the images whose ref contains this")
    a = ap.parse_args()

    rdir = a.recipe_dir if a.recipe_dir != RECIPE_DIR or not a.ssh else "~/.abra/recipes"
    refs = (compose_images_ssh(a.recipe, a.ssh, rdir) if a.ssh
            else compose_images(a.recipe, a.recipe_dir))
    if a.only:
        refs = [r for r in refs if a.only in r]
    results = [resolve(r) for r in refs]
    report = {
        "recipe": a.recipe,
        "images": results,
        "upgrades_available": [r["ref"] for r in results if r["upgrade_available"]],
        "unresolved": [r["ref"] for r in results if r["status"].startswith("error")],
        # The whole point: distinguish "checked, current" from "could not check".
        "all_resolved": not any(r["status"].startswith("error") for r in results),
    }
    if not a.table:
        print(json.dumps(report, indent=2))
        return 0
    print(f"{a.recipe} — {len(results)} images")
    for r in results:
        flag = "UPGRADE" if r["upgrade_available"] else ("ERROR" if r["status"].startswith("error") else "current")
        print(f"  [{flag:7}] {r['repo']}:{r['tag']}" + ("  (digest-pinned)" if r["digest"] else ""))
        print(f"            shape={r['shape']}  candidates={r['candidates']}"
              f"  newest_same_shape={r['newest_same_shape']}"
              f"  newest_within_major={r['newest_within_major']}")
        if r["status"] != "ok":
            print(f"            status: {r['status']}")
    return 0


if __name__ == "__main__":
    sys.exit(main())