#!/usr/bin/env python3
"""Generate the extras of the RSC certificate check (the exact c_Z and the stage sizes).

Usage:
    python3 formal/scripts/export_rsc_extras.py <rsc-certificate.json.gz> <RSCData dir> [--selfcheck]

Writes <RSCData dir>/Extras/Shards.lean, Extras/All.lean and Extras/MANIFEST.sha256. No new data
enters Lean: the six shard theorems check run A's exported `Z` and stages, and the literal arguments
written here are not trusted (a wrong literal makes a shard or a fixed proof fail).

The script aborts (non-zero exit, message on stderr) unless
  1. the gz and decompressed SHA-256 equal the pinned values;
  2. the RSCData/*.lean files present (non-recursive, so Extras/ is excluded) are exactly the 467 files
     of run A's pinned list (the non-comment lines of RSCData/MANIFEST.sha256, whose SHA-256 is pinned
     below), each with its listed SHA-256;
  3. run A's exporter export_rsc_cert.py (imported read-only for build_plan, nbl, enc_pt, fin_n,
     lt_lex) has the pinned SHA-256.
It never writes to run A's files. Output is deterministic. Standard library only.
--selfcheck re-parses the two generated files and run A's Z.lean and St_k_pts.lean, and re-evaluates
checkCZ and checkStages (RSCCheckExtras.lean) in exact integer arithmetic.
"""
import argparse
import ast
import gzip
import hashlib
import importlib.util
import json
import os
import re
import sys
from fractions import Fraction

GZ_SHA = "9465bd0f7f74dc3a6376724645bfbde97b417035583c89ba787e347813177651"
JSON_SHA = "188c5b7206fe3ba487513caa859ae68777efdae15d184f94092dd8904a55b48b"
RUNA_EXPORTER_SHA = "9191995423a44b6d882afb9076aff1fa9775ae272a2ad4f1985c144344c8386b"
RUNA_LIST_SHA = "055df297a7473f49da8a1ea794e085eab05bd84cae0af2a84769f2720dd5ef72"
RUNA_NFILES = 467
NS = "MaxLeafFormal.RSCData.Extras"
DATA_NS = "MaxLeafFormal.RSCData"
METHOD = "m2"   # Serge, 2026-10-29: "use native_decide"; applied to the six new shards (PD5)
HERE = os.path.dirname(os.path.abspath(__file__))


class ExtrasError(Exception):
    pass


def fail(msg):
    raise ExtrasError(msg)


def sha256_bytes(b):
    return hashlib.sha256(b).hexdigest()


def sha256_file(path):
    with open(path, "rb") as fh:
        return sha256_bytes(fh.read())


# ----------------------------------------------------------------------------------------------
# checks of the inputs

def check_runa_tree(datadir):
    man = os.path.join(datadir, "MANIFEST.sha256")
    with open(man, "rb") as fh:
        lines = fh.read().decode("utf-8").splitlines()
    listed = [l for l in lines if not l.startswith("#")]
    lst = "".join(l + "\n" for l in listed).encode("utf-8")
    if sha256_bytes(lst) != RUNA_LIST_SHA:
        fail("SHA-256 of run A's file list (non-comment lines of %s) is %s, expected %s"
             % (man, sha256_bytes(lst), RUNA_LIST_SHA))
    want = {}
    for l in listed:
        m = re.match(r"^([0-9a-f]{64})  (\S+)$", l)
        if m is None:
            fail("bad manifest line %r" % l)
        want[m.group(2)] = m.group(1)
    if len(want) != RUNA_NFILES:
        fail("run A's list has %d files, expected %d" % (len(want), RUNA_NFILES))
    present = sorted(f for f in os.listdir(datadir)
                     if f.endswith(".lean") and os.path.isfile(os.path.join(datadir, f)))
    if sorted(want) != present:
        missing = sorted(set(want) - set(present))
        extra = sorted(set(present) - set(want))
        fail("RSCData/*.lean files differ from run A's list: missing %s, extra %s" % (missing[:5], extra[:5]))
    bad = [f for f in present if sha256_file(os.path.join(datadir, f)) != want[f]]
    if bad:
        fail("%d RSCData files differ from run A's list, e.g. %s" % (len(bad), bad[:5]))
    return len(present)


def load_runa_exporter():
    path = os.path.join(HERE, "export_rsc_cert.py")
    h = sha256_file(path)
    if h != RUNA_EXPORTER_SHA:
        fail("SHA-256 of %s is %s, expected %s" % (path, h, RUNA_EXPORTER_SHA))
    spec = importlib.util.spec_from_file_location("export_rsc_cert_runA", path)
    mod = importlib.util.module_from_spec(spec)
    spec.loader.exec_module(mod)
    return mod


# ----------------------------------------------------------------------------------------------
# Lean text

class Ctx:
    script_sha = ""


def lean_file(imports, body):
    head = [
        "/-",
        "Generated by formal/scripts/export_rsc_extras.py -- do not edit.",
        "Data, not proof: Lean re-checks everything.",
        "input rsc-certificate.json.gz sha256: " + GZ_SHA,
        "input rsc-certificate.json sha256:    " + JSON_SHA,
        "script sha256:                        " + Ctx.script_sha,
        "run A exporter sha256:                " + RUNA_EXPORTER_SHA,
        "run A file list sha256:               " + RUNA_LIST_SHA,
        "method:                               " + METHOD,
        "-/",
    ]
    head += ["import " + i for i in imports]
    head += ["", "set_option maxRecDepth 1000000", "set_option maxHeartbeats 0", "",
             "namespace " + NS, "", "noncomputable section", "", "open MaxLeafFormal.RSCCheck", ""]
    tail = ["", "end", "end " + NS, ""]
    return "\n".join(head) + "\n".join(body) + "\n".join(tail)


def shard_lines(kmin, a, b, sizes):
    tac = "native_decide"
    out = ["theorem cZ_ok : RSCCheck.checkCZ Z %d %d %d = true := by %s" % (kmin, a, b, tac)]
    for k, s in enumerate(sizes):
        out.append("theorem sz_%d : RSCCheck.checkStages St_%d_pts.stages [%s] = true := by %s"
                   % (k, k, ", ".join(map(str, s)), tac))
    return out


def facts_tuple(ex, k, sizes_x):
    """Run A's rule (export_rsc_cert.py, gen_all) for the stage facts of chain k.

    sizes_x = [|X_1|, ..., |X_j|] = [1] + stage sizes."""
    def forall_tuple(items):
        return items[0] if len(items) == 1 else "⟨%s⟩" % ", ".join(items)
    facts = []
    for i in range(1, len(sizes_x) + 1):
        nb = ex.nbl(sizes_x[i - 1])
        blocks = forall_tuple(["⟨_, St_%d_%d_%03d.ok⟩" % (k, i, r) for r in range(nb)])
        facts.append("stage_of_blocks Z_ok (stage_X_ok St_%d_pts.pts_ok St_%d_pts.key %d) "
                     "(stage_Yt_ok Z_ok St_%d_pts.pts_ok %d) %d St_%d_pts.nb_%d %s" %
                     (k, k, i - 1, k, i - 1, nb, k, i, blocks))
    return forall_tuple(facts) if len(facts) == 1 else "⟨%s⟩" % ",\n      ".join(facts)


def facts_stmt(k):
    K = "St_%d_pts" % k
    return ("theorem facts_%d : List.Forall (fun i => StageFact Z (([initI %s.key.1 %s.key.2.1 %s.key.2.2]\n"
            "    :: %s.stages).getD i []) ((%s.stages ++ [Z]).getD i []))\n"
            "    (List.range (%s.stages.length + 1)) :=" % (k, K, K, K, K, K, K))


# ----------------------------------------------------------------------------------------------
# exact computation

def compute(ex, cert):
    P = ex.build_plan(cert)
    Z = [ex.decode_pt(r) for r in P.Zraw]
    fins = [Fraction(ex.fin_n(N), 2000 * D) for N, D in Z]
    m = min(fins)
    kmin = fins.index(m)
    st = []
    for (key, j, sizes) in P.st_info:
        st.append((key, j, sizes))
    return P, Z, kmin, m, st


# ----------------------------------------------------------------------------------------------
# self-check

def lean_list_after(text, name):
    m = re.search(r"^def %s\b[^\n]*:= " % re.escape(name), text, re.M)
    if m is None:
        fail("no 'def %s' found" % name)
    rest = text[m.end():]
    end = re.search(r"\n\n", rest)
    return ast.literal_eval(rest[:end.start()] if end else rest)


def check_cz(Z, k, a, b, ex):
    if not ex.ok_pts(Z) or not k < len(Z) or not 0 < b:
        return False
    N, D = Z[k]
    if b * ex.fin_n(N) != a * (2000 * D):
        return False
    return all(a * (2000 * D2) <= b * ex.fin_n(N2) for N2, D2 in Z)


def check_stages(stages, ns, ex):
    if len(stages) != len(ns):
        return False
    for Y, n in zip(stages, ns):
        if len(Y) != n or not ex.ok_pts(Y):
            return False
        if not all(ex.lt_lex(P[1], Q[1], P[0], Q[0]) for P, Q in zip(Y, Y[1:])):
            return False
    return True


def selfcheck(ex, datadir, nchains):
    exdir = os.path.join(datadir, "Extras")
    with open(os.path.join(exdir, "Shards.lean"), encoding="utf-8") as fh:
        sh = fh.read()
    with open(os.path.join(datadir, "Z.lean"), encoding="utf-8") as fh:
        Z = [ex.decode_pt(r) for r in lean_list_after(fh.read(), "Zraw")]
    m = re.search(r"^theorem cZ_ok : RSCCheck\.checkCZ Z ([0-9]+) ([0-9]+) ([0-9]+) = true := by native_decide$",
                  sh, re.M)
    if m is None:
        fail("selfcheck: cZ_ok not found in Shards.lean")
    k, a, b = map(int, m.groups())
    if not check_cz(Z, k, a, b, ex):
        fail("selfcheck: checkCZ Z %d %d %d is false" % (k, a, b))
    print("selfcheck: checkCZ Z %d %d %d = true (c_Z = %s)" % (k, a, b, Fraction(a, b)))
    for c in range(nchains):
        m = re.search(r"^theorem sz_%d : RSCCheck\.checkStages St_%d_pts\.stages \[([0-9, ]+)\] = true := by native_decide$"
                      % (c, c), sh, re.M)
        if m is None:
            fail("selfcheck: sz_%d not found" % c)
        ns = [int(t) for t in m.group(1).split(",")]
        with open(os.path.join(datadir, "St_%d_pts.lean" % c), encoding="utf-8") as fh:
            raw = lean_list_after(fh.read(), "stagesRaw")
        stages = [[ex.decode_pt(p) for p in Y] for Y in raw]
        if not check_stages(stages, ns, ex):
            fail("selfcheck: checkStages St_%d_pts.stages %s is false" % (c, ns))
        print("selfcheck: checkStages St_%d_pts.stages %s = true" % (c, ns))
    with open(os.path.join(exdir, "All.lean"), encoding="utf-8") as fh:
        al = fh.read()
    with open(os.path.join(datadir, "All.lean"), encoding="utf-8") as fh:
        runa_all = fh.read()
    for c in range(nchains):
        mm = re.search(r"^theorem facts_%d : .*?:=\n  (.*?)\n(?:\n|end\n)" % c, al, re.M | re.S)
        if mm is None:
            fail("selfcheck: facts_%d not found in Extras/All.lean" % c)
        tup = mm.group(1)
        needle = "chain_stages_table Z_ok St_%d_pts.key St_%d_pts.stages St_%d_pts.pts_ok %s" % (c, c, c, tup)
        if needle not in runa_all:
            fail("selfcheck: the proof term of facts_%d is not the tuple of run A's RSCData/All.lean" % c)
    print("selfcheck: the %d facts_k proof terms are the tuples of run A's RSCData/All.lean" % nchains)
    print("selfcheck passed")


# ----------------------------------------------------------------------------------------------

def main():
    try:
        sys.stdout.reconfigure(encoding="utf-8")
    except Exception:
        pass
    ap = argparse.ArgumentParser()
    ap.add_argument("gz")
    ap.add_argument("datadir")
    ap.add_argument("--selfcheck", action="store_true")
    a = ap.parse_args()
    Ctx.script_sha = sha256_file(os.path.abspath(__file__))
    with open(a.gz, "rb") as fh:
        gzb = fh.read()
    if sha256_bytes(gzb) != GZ_SHA:
        fail("SHA-256 of the gz file is %s, expected %s" % (sha256_bytes(gzb), GZ_SHA))
    raw = gzip.decompress(gzb)
    if sha256_bytes(raw) != JSON_SHA:
        fail("SHA-256 of the decompressed JSON is %s, expected %s" % (sha256_bytes(raw), JSON_SHA))
    datadir = os.path.abspath(a.datadir)
    n = check_runa_tree(datadir)
    ex = load_runa_exporter()
    print("script sha256 %s ; python %s" % (Ctx.script_sha, sys.version.split()[0]))
    print("input hashes match; run A tree: %d files match the pinned list (sha256 %s); run A exporter %s"
          % (n, RUNA_LIST_SHA, RUNA_EXPORTER_SHA))
    cert = json.loads(raw)
    del raw, gzb
    P, Z, kmin, m, st = compute(ex, cert)
    del cert
    print("c_Z = %d/%d = %.10f at least index %d (of %d points)"
          % (m.numerator, m.denominator, float(m), kmin, len(Z)))
    sizes = [s[1:] for (_, _, s) in st]
    print("chains with j >= 2 (file order): %s" % [(key, j, s[1:]) for key, j, s in st])
    # Shards.lean
    imports = ["MaxLeafFormal.RSCCheckExtras", DATA_NS + ".Z"] + \
              [DATA_NS + ".St_%d_pts" % k for k in range(len(st))]
    shards = lean_file(imports, shard_lines(kmin, m.numerator, m.denominator, sizes))
    # All.lean
    mods = []
    body = []
    for k, (key, j, sx) in enumerate(st):
        for i in range(1, j + 1):
            for r in range(ex.nbl(sx[i - 1])):
                mods.append(DATA_NS + ".St_%d_%d_%03d" % (k, i, r))
        if k:
            body.append("")
        body.append(facts_stmt(k))
        body.append("  " + facts_tuple(ex, k, sx))
    print("stage-block modules imported by Extras/All.lean: %d (per chain %s)"
          % (len(mods), [sum(ex.nbl(n) for n in sx) for _, _, sx in st]))
    allf = lean_file([NS + ".Shards"] + mods, body)
    allf = allf.replace("open MaxLeafFormal.RSCCheck\n",
                        "open MaxLeafFormal.RSCCheck\nopen MaxLeafFormal.RSCCert MaxLeafFormal.RSCParams\n", 1)
    exdir = os.path.join(datadir, "Extras")
    os.makedirs(exdir, exist_ok=True)
    out = {}
    for name, txt in (("Shards.lean", shards), ("All.lean", allf)):
        data = txt.encode("utf-8")
        with open(os.path.join(exdir, name), "wb") as fh:
            fh.write(data)
        out[name] = sha256_bytes(data)
    lines = ["# input-gz sha256 %s" % GZ_SHA, "# input-json sha256 %s" % JSON_SHA,
             "# script sha256 %s" % Ctx.script_sha, "# run-A exporter sha256 %s" % RUNA_EXPORTER_SHA,
             "# run-A file list sha256 %s" % RUNA_LIST_SHA,
             "# python %d.%d.%d" % sys.version_info[:3]]
    for name in sorted(out):
        lines.append("%s  %s" % (out[name], name))
    with open(os.path.join(exdir, "MANIFEST.sha256"), "wb") as fh:
        fh.write(("\n".join(lines) + "\n").encode("utf-8"))
    for name in sorted(out):
        print("wrote Extras/%s sha256 %s" % (name, out[name]))
    if a.selfcheck:
        selfcheck(ex, datadir, len(st))


if __name__ == "__main__":
    try:
        main()
    except ExtrasError as e:
        print("EXTRAS ERROR: %s" % e, file=sys.stderr)
        sys.exit(1)