#!/usr/bin/env python3
# Fetch NIST ACVP ML-KEM-768 ENCAPSULATION vectors and emit them in the flat
# format MlKemAcvpKat.java consumes:  tcId|ek_hex|m_hex|c_hex|k_hex
#
# Source: usnistgov/ACVP-Server, ML-KEM-encapDecap-FIPS203 internalProjection.json
# (the "internal projection" carries both the inputs ek/m and expected outputs c/k).
# Encapsulation AFT test groups give ek + m and expected ciphertext c + shared key k.

import json
import sys
import urllib.request

URLS = [
    "https://raw.githubusercontent.com/usnistgov/ACVP-Server/master/gen-val/json-files/ML-KEM-encapDecap-FIPS203/internalProjection.json",
    "https://raw.githubusercontent.com/usnistgov/ACVP-Server/master/gen-val/json-files/ML-KEM-encapDecap-FIPS203/expectedResults.json",
]


def fetch(url):
    req = urllib.request.Request(url, headers={"User-Agent": "aere-kat/1.0"})
    with urllib.request.urlopen(req, timeout=60) as r:
        return json.loads(r.read().decode())


def get(d, *names):
    for n in names:
        for k in d:
            if k.lower() == n.lower():
                return d[k]
    return None


def main():
    out = sys.argv[1] if len(sys.argv) > 1 else "mlkem768_acvp.txt"
    data = None
    used = None
    for u in URLS:
        try:
            data = fetch(u)
            used = u
            break
        except Exception as e:  # noqa
            print(f"fetch failed {u}: {e}", file=sys.stderr)
    if data is None:
        print("ACVP_FETCH_FAILED", file=sys.stderr)
        sys.exit(2)

    groups = get(data, "testGroups") or []
    lines = []
    for g in groups:
        pset = str(get(g, "parameterSet") or "")
        func = str(get(g, "function") or "")
        if "768" not in pset:
            continue
        if "encap" not in func.lower():
            continue
        for t in get(g, "tests") or []:
            ek = get(t, "ek")
            m = get(t, "m")
            c = get(t, "c")
            k = get(t, "k")
            tc = get(t, "tcId")
            if ek and m and c and k:
                lines.append(f"{tc}|{ek}|{m}|{c}|{k}")

    if not lines:
        print("ACVP_NO_ENCAP_VECTORS (schema mismatch) source=" + str(used), file=sys.stderr)
        sys.exit(2)

    with open(out, "w") as f:
        f.write("# NIST ACVP ML-KEM-768 encapsulation vectors (ek|m|c|k)\n")
        f.write("# source: " + used + "\n")
        f.write("\n".join(lines) + "\n")
    print(f"ACVP_OK wrote {len(lines)} ML-KEM-768 encapsulation vectors -> {out}")


if __name__ == "__main__":
    main()