#!/usr/bin/env python3
"""RSC certificate Lean check: token audit and provenance comparison (standard library only).

  python3 formal/scripts/rsc_data_audit.py --scan
      Scans formal/MaxLeafFormal/RSCCheck.lean, RSCFinal.lean and RSCData/*.lean for forbidden
      tokens outside comments, counts native_decide (required to be zero), and checks the theorem
      forms in RSCData (one-line `... = true := by decide +kernel`, or term-mode aggregation in
      All.lean).

  python3 formal/scripts/rsc_data_audit.py --compare <rsc-certificate.json.gz>
      Parses RSCData/Z.lean and RSCData/St_*_pts.lean independently of the exporter and checks
      exact equality (values and order) with the JSON certificate. Also checks the gz SHA-256.

Exit status: 0 on success, 1 on any finding or mismatch, 2 on usage error.
"""

import argparse
import glob
import gzip
import hashlib
import json
import os
import re
import sys
from fractions import Fraction

GZ_SHA256 = "9465bd0f7f74dc3a6376724645bfbde97b417035583c89ba787e347813177651"
REPO = os.path.normpath(os.path.join(os.path.dirname(os.path.abspath(__file__)), os.pardir, os.pardir))
LEAN_DIR = os.path.join(REPO, "formal", "MaxLeafFormal")
DATA_DIR = os.path.join(LEAN_DIR, "RSCData")

FORBIDDEN_TOKENS = ["axiom", "sorry", "admit", "unsafe", "@[extern", "@[implemented_by",
                    "implemented_by", "debug.skipKernelTC", "native_decide"]
FORBIDDEN_RE = re.compile(
    r"(?<![\w'!?])(axiom|sorry|admit|unsafe|implemented_by|debug\.skipKernelTC|native_decide)(?![\w'!?])"
    r"|@\[\s*extern\b")
SPECIAL_RE = re.compile(r"[-/\"']")
CHARLIT_RE = re.compile(r"'(?:\\.|[^\\'\n])'")
THEOREM_RE = re.compile(r"(?<![\w'!?])theorem(?![\w'!?])")
FORM1_RE = re.compile(r"^\s*theorem\s+[^\s:]+\s*:.*=\s*true\s*:=\s*by\s+decide\s+\+kernel\s*$")
# Method M2 (Serge, 2026-10-29): the same one-line form, decided by `native_decide`; allowed only in RSCData files.
FORM1N_RE = re.compile(r"^\s*theorem\s+[^\s:]+\s*:.*=\s*true\s*:=\s*by\s+native_decide\s*$")
METHOD = {"m": "m2"}
BY_RE = re.compile(r"(?<![\w'!?])by(?![\w'!?])")
DECL_RE = re.compile(r"^(def|theorem|lemma|end|namespace|section|abbrev|instance|open|import|"
                     r"noncomputable|set_option)\b")
LIST_TOK_RE = re.compile(r"\[|\]|,|[0-9]+|\s+")
KEY_RE = re.compile(r"\(\s*([0-9]+)\s*,\s*([0-9]+)\s*,\s*([0-9]+)\s*\)")


class Fail(Exception):
    pass


def strip_line(line, depth, in_str):
    """Blank out comments and string literals in one physical line.

    Returns (code, block_depth, in_string). Lean block comments nest; `--` ends a line.
    """
    if depth == 0 and not in_str and SPECIAL_RE.search(line) is None:
        return line, 0, False
    out = []
    i, n = 0, len(line)
    while i < n:
        if depth:
            if line.startswith("/-", i):
                depth += 1
                i += 2
            elif line.startswith("-/", i):
                depth -= 1
                i += 2
            else:
                i += 1
        elif in_str:
            if line[i] == "\\":
                i += 2
            else:
                if line[i] == '"':
                    in_str = False
                i += 1
        elif line.startswith("--", i):
            break
        elif line.startswith("/-", i):
            depth = 1
            out.append(" ")
            i += 2
        elif line[i] == '"':
            in_str = True
            out.append(" ")
            i += 1
        else:
            m = CHARLIT_RE.match(line, i) if line[i] == "'" else None
            if m:
                out.append(" ")
                i = m.end()
            else:
                out.append(line[i])
                i += 1
    return "".join(out), depth, in_str


def iter_code(path):
    """Yield (lineno, raw_line, code_line) for a Lean file, streaming line by line."""
    depth, in_str = 0, False
    with open(path, encoding="utf-8", newline="") as f:
        for lineno, raw in enumerate(f, 1):
            line = raw.rstrip("\r\n")
            code, depth, in_str = strip_line(line, depth, in_str)
            yield lineno, line, code


def audit_theorems(lines, is_all):
    """Classify the theorem declarations of one RSCData file.

    Returns (form1_count, form2_count, exceptions). Form 1 is a one-line
    `theorem N : e = true := by decide +kernel`. Form 2 (All.lean only) is a term-mode
    aggregation whose body is a term and contains no `by`.
    """
    form1 = form2 = 0
    exc = []
    n = len(lines)
    i = 0
    while i < n:
        lineno, raw, code = lines[i]
        if THEOREM_RE.search(code) is None:
            i += 1
            continue
        if FORM1_RE.match(code) or (METHOD["m"] == "m2" and FORM1N_RE.match(code)):
            form1 += 1
            i += 1
            continue
        if not is_all:
            exc.append((lineno, raw.strip(), "not the one-line decide +kernel form"))
            i += 1
            continue
        j, acc = i, code
        while ":=" not in acc and j + 1 < n:
            j += 1
            acc += " " + lines[j][2]
        if ":=" not in acc:
            exc.append((lineno, raw.strip(), "no ':=' found"))
            i += 1
            continue
        body = acc.split(":=", 1)[1]
        k = j
        while k + 1 < n and (lines[k + 1][2].strip() == "" or lines[k + 1][2][:1].isspace()):
            k += 1
            body += " " + lines[k][2]
        stripped = body.strip()
        if not stripped:
            exc.append((lineno, raw.strip(), "empty term-mode body"))
        elif stripped.split()[0] == "by" or BY_RE.search(body):
            exc.append((lineno, raw.strip(), "tactic block in theorem body"))
        else:
            form2 += 1
        i = k + 1
    return form1, form2, exc


def short(text, width=150):
    return text if len(text) <= width else text[:width - 3] + "..."


def run_scan():
    print("lean dir: " + LEAN_DIR)
    scan_files = [os.path.join(LEAN_DIR, "RSCCheck.lean"), os.path.join(LEAN_DIR, "RSCFinal.lean")]
    data_files = sorted(glob.glob(os.path.join(DATA_DIR, "*.lean")))
    for p in scan_files:
        if not os.path.isfile(p):
            raise Fail("missing " + p)
    if not data_files:
        raise Fail("no RSCData/*.lean files under " + DATA_DIR)

    token_counts = {t: 0 for t in FORBIDDEN_TOKENS}
    token_hits = []
    native_hits = []
    native_shard = []
    n_lines = n_bytes = 0
    form1 = form2 = 0
    exceptions = []

    for path in scan_files + data_files:
        rel = os.path.relpath(path, REPO).replace(os.sep, "/")
        n_bytes += os.path.getsize(path)
        lines = list(iter_code(path))
        n_lines += len(lines)
        for lineno, raw, code in lines:
            for m in FORBIDDEN_RE.finditer(code):
                tok = m.group(1) if m.group(1) else "@[extern"
                token_counts[tok] = token_counts.get(tok, 0) + 1
                entry = (rel, lineno, tok, short(raw.strip()))
                if tok == "native_decide":
                    # M2: allowed only in a one-line shard theorem of an RSCData file
                    if (METHOD["m"] == "m2" and os.path.dirname(path) == DATA_DIR
                            and FORM1N_RE.match(code)):
                        native_shard.append(entry)
                    else:
                        native_hits.append(entry)
                else:
                    token_hits.append(entry)
        if os.path.dirname(path) == DATA_DIR:
            f1, f2, exc = audit_theorems(lines, os.path.basename(path) == "All.lean")
            form1 += f1
            form2 += f2
            for lineno, text, reason in exc:
                exceptions.append((rel, lineno, text, reason))

    print("files scanned: %d (RSCCheck.lean, RSCFinal.lean, %d RSCData files)"
          % (len(scan_files) + len(data_files), len(data_files)))
    print("physical lines: %d, bytes: %d" % (n_lines, n_bytes))
    print("per-token counts (outside comments and strings):")
    for t in FORBIDDEN_TOKENS:
        print("  %-22s %d" % (t, token_counts.get(t, 0)))
    print("forbidden-token hits (excluding native_decide): %d" % len(token_hits))
    for rel, lineno, tok, text in token_hits:
        print("  %s:%d: %s: %s" % (rel, lineno, tok, text))
    print("method: %s" % METHOD["m"])
    print("native_decide in RSCData one-line shard theorems (allowed under M2): %d" % len(native_shard))
    print("native_decide occurrences elsewhere: %d (required: 0)" % len(native_hits))
    for rel, lineno, tok, text in native_hits:
        print("  %s:%d: %s" % (rel, lineno, text))
    print("RSCData theorem forms: one-line decide +kernel or native_decide (M2)=%d, term-mode aggregation (All.lean)=%d, exceptions=%d"
          % (form1, form2, len(exceptions)))
    for rel, lineno, text, reason in exceptions:
        print("  EXCEPTION %s:%d: %s  [%s]" % (rel, lineno, short(text), reason))
    ok = not token_hits and not native_hits and not exceptions
    print("SCAN RESULT: " + ("PASS" if ok else "FAIL"))
    return ok


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 decl_text(path, name):
    """Return the code text after `:=` of the single top-level `def name`."""
    code_lines = [code for _, _, code in iter_code(path)]
    pat = re.compile(r"^def\s+" + re.escape(name) + r"\b")
    starts = [i for i, c in enumerate(code_lines) if pat.match(c)]
    if len(starts) != 1:
        raise Fail("%s: expected one 'def %s', found %d" % (path, name, len(starts)))
    out = [code_lines[starts[0]]]
    for c in code_lines[starts[0] + 1:]:
        if DECL_RE.match(c):
            break
        out.append(c)
    text = "\n".join(out)
    if ":=" not in text:
        raise Fail("%s: 'def %s' has no ':='" % (path, name))
    return text.split(":=", 1)[1]


def parse_nested(text):
    """Parse a Lean list literal of naturals (nested lists) into Python lists of ints."""
    toks, pos = [], 0
    while pos < len(text):
        m = LIST_TOK_RE.match(text, pos)
        if m is None:
            raise Fail("unexpected text in list literal: %r" % text[pos:pos + 40])
        if not m.group(0).isspace():
            toks.append(m.group(0))
        pos = m.end()
    root, stack = None, []
    for tok in toks:
        if tok == "[":
            new = []
            if stack:
                stack[-1].append(new)
            elif root is None:
                root = new
            else:
                raise Fail("more than one top-level list literal")
            stack.append(new)
        elif tok == "]":
            if not stack:
                raise Fail("unbalanced ']'")
            stack.pop()
        elif tok == ",":
            continue
        else:
            if not stack:
                raise Fail("number outside a list")
            stack[-1].append(int(tok))
    if root is None or stack:
        raise Fail("unbalanced or empty list literal")
    return root


def lean_point(row):
    """A 9-number Lean row (8 numerators, common denominator D) as 8 exact Fractions."""
    if len(row) != 9 or any(type(x) is not int for x in row):
        raise Fail("Lean row is not 9 naturals: %r" % (row,))
    den = row[8]
    if den <= 0:
        raise Fail("Lean row with non-positive denominator")
    return tuple(Fraction(row[i], den) for i in range(8))


def json_point(pt, where):
    """A JSON point: 8 coordinates, each [num, den] with integer entries and den > 0."""
    if not isinstance(pt, list) or len(pt) != 8:
        raise Fail("%s: point is not a list of 8 coordinates" % where)
    out = []
    for c in pt:
        if not (isinstance(c, list) and len(c) == 2 and all(type(x) is int for x in c) and c[1] > 0):
            raise Fail("%s: bad coordinate %r" % (where, c))
        out.append(Fraction(c[0], c[1]))
    return tuple(out)


def compare_seq(expected, actual, what, problems):
    if len(expected) != len(actual):
        problems.append("%s: length %d (Lean) vs %d (JSON)" % (what, len(expected), len(actual)))
        return False
    bad = [i for i, (x, y) in enumerate(zip(expected, actual)) if x != y]
    for i in bad[:5]:
        problems.append("%s[%d]: Lean %s != JSON %s" % (what, i, expected[i], actual[i]))
    if len(bad) > 5:
        problems.append("%s: %d differences in total" % (what, len(bad)))
    return not bad


def run_compare(cert_path):
    problems = []
    digest = sha256_file(cert_path)
    print("certificate: " + cert_path)
    print("gz sha256: " + digest)
    print("gz sha256 expected: " + GZ_SHA256 + (" (match)" if digest == GZ_SHA256 else " (MISMATCH)"))
    if digest != GZ_SHA256:
        problems.append("gz SHA-256 differs from the pinned value")

    with gzip.open(cert_path, "rb") as f:
        cert = json.load(f)
    if not isinstance(cert, dict) or "Z" not in cert or "internal" not in cert:
        raise Fail("JSON lacks 'Z' or 'internal'")

    # Z: points in file order.
    z_rows = parse_nested(decl_text(os.path.join(DATA_DIR, "Z.lean"), "Zraw"))
    z_lean = [lean_point(r) for r in z_rows]
    z_json = [json_point(p, "Z[%d]" % i) for i, p in enumerate(cert["Z"])]
    z_ok = compare_seq(z_lean, z_json, "Z", problems)
    print("Z: Lean %d points, JSON %d points, exact match in order: %s"
          % (len(z_lean), len(z_json), "yes" if z_ok else "NO"))

    # Stage data St_k_pts.lean, k = 0..K-1.
    idx_re = re.compile(r"^St_([0-9]+)_pts\.lean$")
    st_files = []
    for p in glob.glob(os.path.join(DATA_DIR, "St_*_pts.lean")):
        m = idx_re.match(os.path.basename(p))
        if m:
            st_files.append((int(m.group(1)), p))
    st_files.sort()
    if [k for k, _ in st_files] != list(range(len(st_files))):
        problems.append("St_k_pts.lean indices are not 0..%d" % (len(st_files) - 1))
    lean_st = []
    lean_by_key = {}
    for k, path in st_files:
        mk = KEY_RE.search(decl_text(path, "key"))
        if mk is None:
            raise Fail("%s: cannot parse 'def key'" % path)
        key = tuple(int(x) for x in mk.groups())
        if key in lean_by_key:
            problems.append("duplicate key %s in Lean stage files" % (key,))
        lean_by_key[key] = k
        stages = parse_nested(decl_text(path, "stagesRaw"))
        lean_st.append((key, [[lean_point(p) for p in stage] for stage in stages]))

    json_st = [e for e in cert["internal"] if int(e["j"]) >= 2]
    print("stage entries: Lean %d St files, JSON %d internal entries with j >= 2"
          % (len(lean_st), len(json_st)))
    if len(lean_st) != len(json_st):
        problems.append("stage entry count: Lean %d vs JSON %d" % (len(lean_st), len(json_st)))

    n_stages = n_points = matched = 0
    for t, e in enumerate(json_st):
        key = (int(e["m"]), int(e["c2"]), int(e["c3"]))
        if key not in lean_by_key:
            problems.append("JSON entry %d key %s has no St_k_pts.lean" % (t, key))
            continue
        matched += 1
        k = lean_by_key[key]
        if k != t:
            problems.append("order: JSON entry %d has key %s but it is St_%d" % (t, key, k))
        _, lean_stages = lean_st[k]
        json_stages = [[json_point(p, "%s stage %d point %d" % (key, si, pi)) for pi, p in enumerate(Y)]
                       for si, Y in enumerate(e["stages"])]
        if compare_seq(lean_stages, json_stages, "stages of %s" % (key,), problems):
            n_stages += len(json_stages)
            n_points += sum(len(Y) for Y in json_stages)
    print("stages compared: %d entries matched by key, %d stages, %d points, %d coordinates"
          % (matched, n_stages, n_points, 8 * n_points))

    for p in problems:
        print("MISMATCH: " + p)
    ok = not problems
    print("COMPARE RESULT: " + ("PASS (exact equality with the JSON certificate)" if ok else "FAIL"))
    return ok


def main(argv=None):
    try:
        sys.stdout.reconfigure(errors="backslashreplace")
    except AttributeError:
        pass
    ap = argparse.ArgumentParser(description="RSC certificate Lean check: token audit and provenance comparison.")
    ap.add_argument("--scan", action="store_true",
                    help="token and theorem-form audit of RSCCheck, RSCFinal and RSCData")
    ap.add_argument("--compare", metavar="CERT_GZ",
                    help="compare RSCData Z and stage data with the JSON certificate")
    args = ap.parse_args(argv)
    if not args.scan and not args.compare:
        ap.print_usage(sys.stderr)
        print("error: give --scan and/or --compare CERT_GZ", file=sys.stderr)
        return 2
    ok = True
    try:
        if args.scan:
            ok = run_scan() and ok
        if args.compare:
            ok = run_compare(args.compare) and ok
    except Fail as e:
        print("ERROR: %s" % e)
        return 1
    return 0 if ok else 1


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