#!/usr/bin/env python3
"""Say which followed engines have published a version newer than the one llmgw installs.

The versions llmgw installs are DefaultEngines in internal/config/config.go. An
engine followed here is one whose releases are taken up as they come: each new
one is checked against the flags llmgw gives it (llmgw ctl engine-check, and a
replica on a card it runs on) and then moved to in DefaultEngines.

Exits 1 when one is behind, so the scheduled workflow's failed run is the
notice; 0 when all are current. No token needed: PyPI's JSON API is public.

Usage:
    scripts/engine_releases.py [--config internal/config/config.go]
"""

import argparse
import json
import re
import sys
import urllib.request

# Engine name in DefaultEngines -> (PyPI project, where its changes are read).
FOLLOWED = {
    "freetoken": ("freetoken", "https://github.com/FlashML-org/FreeToken/releases"),
}


def pinned(config_path):
    """The DefaultEngines map, as written in the Go source."""
    src = open(config_path, encoding="utf-8").read()
    block = re.search(r"var DefaultEngines = map\[string\]string\{(.*?)\n\}", src, re.S)
    if not block:
        sys.exit(f"no DefaultEngines in {config_path}")
    return dict(re.findall(r'"([^"]+)":\s*"([^"]+)"', block.group(1)))


def key(version):
    """A version as numbers to compare; a pre-release sorts before its release."""
    nums = [int(n) for n in re.findall(r"\d+", version.split("+")[0])]
    pre = bool(re.search(r"(a|b|rc|dev)\d*$", version))
    return (nums, 0 if pre else 1)


def latest(project):
    """The newest version PyPI has of a project that is not yanked or a pre-release."""
    with urllib.request.urlopen(f"https://pypi.org/pypi/{project}/json", timeout=30) as r:
        data = json.load(r)
    versions = [
        v for v, files in data["releases"].items()
        if files and not all(f.get("yanked") for f in files) and not re.search(r"(a|b|rc|dev)\d*$", v)
    ]
    return max(versions, key=key) if versions else data["info"]["version"]


def main():
    ap = argparse.ArgumentParser(description=__doc__.splitlines()[0])
    ap.add_argument("--config", default="internal/config/config.go")
    args = ap.parse_args()
    pins = pinned(args.config)
    behind = False
    for engine, (project, changes) in FOLLOWED.items():
        have = pins.get(engine)
        if have is None:
            print(f"{engine}: not in DefaultEngines")
            behind = True
            continue
        new = latest(project)
        if key(new) > key(have):
            print(f"{engine}: {new} is out, llmgw installs {have} — {changes}")
            behind = True
        else:
            print(f"{engine}: {have} is current")
    return 1 if behind else 0


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