#!/usr/bin/env python3
"""Token and form audit of the RSC extras (the exact c_Z and the stage sizes). Standard library only.

  python3 formal/scripts/rsc_extras_audit.py

1. Runs run A's unchanged rsc_data_audit.py --scan (SHA-256 pinned) and requires run A's numbers:
   469 files, 466 native_decide occurrences all in one-line shard theorems, 477 one-line shard
   theorems, 3 term-mode aggregations, 0 exceptions, PASS.
2. Scans RSCStages.lean, RSCCheckExtras.lean, RSCFinalExtras.lean and RSCData/Extras/*.lean for the
   parent's full forbidden list (tokens, the identifiers implemented_by and extern in any syntax, the
   commands) plus `decide +native` and the identifiers ofReduceBool, reduceBool and trustCompiler;
   applies the added items and the parent's commands to RSCCheck.lean, RSCFinal.lean and run A's 467
   RSCData/*.lean files. Comments and strings are stripped with run A's stripper (imported read-only).
3. Extras/Shards.lean: after its header, exactly the six theorem lines of the specification, byte for
   byte, and otherwise only run A's framing lines.
4. Extras/All.lean: exactly the five term-mode theorems facts_0 ... facts_4 (no `by`, no
   native_decide), importing exactly Extras.Shards and the 83 St_k_i_r modules.
5. native_decide occurs nowhere else among the new files; the total is 466 + 6 = 472.
6. Every theorem of RSCStages.lean and RSCCheckExtras.lean that is not private, and IsStages, is
   printed in scripts/AxiomsRSCCheckExtras.lean.
Exit status 0 on PASS, 1 otherwise.
"""
import glob
import hashlib
import importlib.util
import os
import re
import subprocess
import sys

RUNA_AUDIT_SHA = "85f3872ab63fb3e90dee0ada2027f9c85027d350e3a6c7b239bdde07b1465b6c"
HERE = os.path.dirname(os.path.abspath(__file__))
FORMAL = os.path.normpath(os.path.join(HERE, os.pardir))
LEAN_DIR = os.path.join(FORMAL, "MaxLeafFormal")
DATA_DIR = os.path.join(LEAN_DIR, "RSCData")
EXTRAS_DIR = os.path.join(DATA_DIR, "Extras")

SHARD_LINES = [
    "theorem cZ_ok : RSCCheck.checkCZ Z 298 123019910197327428728743 202985322574630920671369 = true := by native_decide",
    "theorem sz_0 : RSCCheck.checkStages St_0_pts.stages [62] = true := by native_decide",
    "theorem sz_1 : RSCCheck.checkStages St_1_pts.stages [62, 304] = true := by native_decide",
    "theorem sz_2 : RSCCheck.checkStages St_2_pts.stages [61] = true := by native_decide",
    "theorem sz_3 : RSCCheck.checkStages St_3_pts.stages [62] = true := by native_decide",
    "theorem sz_4 : RSCCheck.checkStages St_4_pts.stages [62] = true := by native_decide",
]
SHARD_FRAMING = {
    "set_option maxRecDepth 1000000", "set_option maxHeartbeats 0",
    "namespace MaxLeafFormal.RSCData.Extras", "noncomputable section", "open MaxLeafFormal.RSCCheck",
    "end", "end MaxLeafFormal.RSCData.Extras", "",
}
# stage blocks per chain: chain k, stage i, blocks r (run A's nbl of [1] + sizes)
CHAIN_BLOCKS = [[1, 8], [1, 8, 38], [1, 8], [1, 8], [1, 8]]
ALL_IMPORTS = ["MaxLeafFormal.RSCData.Extras.Shards"] + [
    "MaxLeafFormal.RSCData.St_%d_%d_%03d" % (k, i + 1, r)
    for k, bl in enumerate(CHAIN_BLOCKS) for i, nb in enumerate(bl) for r in range(nb)]

# the parent's list (Method and trust), as whole identifiers after stripping comments and strings
BASE_TOKENS = ["axiom", "sorry", "admit", "unsafe", "debug.skipKernelTC", "implemented_by", "extern"]
COMMANDS = ["run_cmd", "run_elab", "run_meta", "elab", "elab_rules", "macro", "macro_rules", "syntax",
            "initialize", "builtin_initialize", "#eval"]
ADDED_IDENTS = ["ofReduceBool", "reduceBool", "trustCompiler"]


def ident_re(names):
    alts = "|".join(re.escape(n) for n in sorted(names, key=len, reverse=True))
    return re.compile(r"(?<![\w'!?#])(" + alts + r")(?![\w'!?])")


FULL_RE = ident_re(BASE_TOKENS + COMMANDS + ADDED_IDENTS)
ADDED_RE = ident_re(COMMANDS + ADDED_IDENTS)
DECIDE_NATIVE_RE = re.compile(r"(?<![\w'!?])decide\s*\+\s*native(?![\w'!?])")
NATIVE_RE = re.compile(r"(?<![\w'!?])native_decide(?![\w'!?])")
THEOREM_RE = re.compile(r"^\s*(?:@\[[^\]]*\]\s*)?(private\s+|protected\s+)?(theorem|lemma)\s+([^\s:({\[]+)")
DEF_RE = re.compile(r"^\s*(?:@\[[^\]]*\]\s*)?(private\s+|protected\s+)?(noncomputable\s+)?(def|abbrev)\s+([^\s:({\[]+)")


class Fail(Exception):
    pass


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


def load_runa_audit():
    path = os.path.join(HERE, "rsc_data_audit.py")
    h = sha256_file(path)
    if h != RUNA_AUDIT_SHA:
        raise Fail("SHA-256 of %s is %s, expected %s" % (path, h, RUNA_AUDIT_SHA))
    spec = importlib.util.spec_from_file_location("rsc_data_audit_runA", path)
    mod = importlib.util.module_from_spec(spec)
    spec.loader.exec_module(mod)
    return mod, path


def rel(p):
    return os.path.relpath(p, FORMAL).replace(os.sep, "/")


def step1(path):
    r = subprocess.run([sys.executable, path, "--scan"], capture_output=True, text=True, encoding="utf-8")
    out = r.stdout
    print("---- run A's rsc_data_audit.py --scan (exit %d) ----" % r.returncode)
    print(out.rstrip())
    print("---- end ----")
    want = [
        (r"^files scanned: (\d+)", "469"),
        (r"^native_decide in RSCData one-line shard theorems \(allowed under M2\): (\d+)", "466"),
        (r"^native_decide occurrences elsewhere: (\d+)", "0"),
        (r"one-line decide \+kernel or native_decide \(M2\)=(\d+)", "477"),
        (r"term-mode aggregation \(All\.lean\)=(\d+)", "3"),
        (r"exceptions=(\d+)", "0"),
        (r"^SCAN RESULT: (\w+)", "PASS"),
    ]
    probs = []
    for pat, val in want:
        m = re.search(pat, out, re.M)
        got = m.group(1) if m else None
        if got != val:
            probs.append("run A scan: %r gave %r, expected %r" % (pat, got, val))
    if r.returncode != 0:
        probs.append("run A scan exit status %d" % r.returncode)
    return probs


def code_lines(audit, path):
    return list(audit.iter_code(path))


def main():
    try:
        sys.stdout.reconfigure(encoding="utf-8", errors="backslashreplace")
    except Exception:
        pass
    print("rsc_extras_audit.py sha256 %s ; python %s" % (sha256_file(os.path.abspath(__file__)),
                                                         sys.version.split()[0]))
    audit, audit_path = load_runa_audit()
    print("run A audit %s sha256 %s (pinned)" % (rel(audit_path), RUNA_AUDIT_SHA))
    probs = step1(audit_path)

    new_files = [os.path.join(LEAN_DIR, f) for f in ("RSCStages.lean", "RSCCheckExtras.lean",
                                                       "RSCFinalExtras.lean")]
    extras = sorted(glob.glob(os.path.join(EXTRAS_DIR, "*.lean")))
    if [os.path.basename(p) for p in extras] != ["All.lean", "Shards.lean"]:
        probs.append("RSCData/Extras/*.lean is %s, expected exactly All.lean and Shards.lean"
                     % [os.path.basename(p) for p in extras])
    new_files += extras
    old_files = [os.path.join(LEAN_DIR, "RSCCheck.lean"), os.path.join(LEAN_DIR, "RSCFinal.lean")] + \
        sorted(glob.glob(os.path.join(DATA_DIR, "*.lean")))
    for p in new_files + old_files:
        if not os.path.isfile(p):
            probs.append("missing " + rel(p))

    # step 2 and 5: tokens
    native_new = []
    hits = []
    for p in new_files:
        for lineno, raw, code in code_lines(audit, p):
            for m in FULL_RE.finditer(code):
                hits.append("%s:%d: %s" % (rel(p), lineno, m.group(1)))
            if DECIDE_NATIVE_RE.search(code):
                hits.append("%s:%d: decide +native" % (rel(p), lineno))
            for m in NATIVE_RE.finditer(code):
                native_new.append((p, lineno, raw))
    n_old_native = 0
    for p in old_files:
        for lineno, raw, code in code_lines(audit, p):
            for m in ADDED_RE.finditer(code):
                hits.append("%s:%d: %s" % (rel(p), lineno, m.group(1)))
            if DECIDE_NATIVE_RE.search(code):
                hits.append("%s:%d: decide +native" % (rel(p), lineno))
            n_old_native += len(NATIVE_RE.findall(code))
    print("new files scanned (%d): %s" % (len(new_files), ", ".join(rel(p) for p in new_files)))
    print("old files scanned for the added items and commands: %d (RSCCheck.lean, RSCFinal.lean, %d RSCData files)"
          % (len(old_files), len(old_files) - 2))
    print("forbidden tokens, identifiers, commands and `decide +native`: %d hits" % len(hits))
    for h in hits:
        print("  " + h)
    probs += ["forbidden: " + h for h in hits]
    shards = os.path.join(EXTRAS_DIR, "Shards.lean")
    bad_native = [(p, l, r) for p, l, r in native_new
                  if not (p == shards and r in SHARD_LINES)]
    print("native_decide in new files: %d (in the six Shards.lean lines: %d, elsewhere: %d)"
          % (len(native_new), len(native_new) - len(bad_native), len(bad_native)))
    for p, l, r in bad_native:
        probs.append("native_decide outside the six shard lines: %s:%d: %s" % (rel(p), l, r.strip()))
    total = n_old_native + len(native_new)
    print("native_decide occurrences in total (run A's files + new files): %d + %d = %d (required 472)"
          % (n_old_native, len(native_new), total))
    if n_old_native != 466 or len(native_new) != 6:
        probs.append("native_decide count %d + %d, expected 466 + 6" % (n_old_native, len(native_new)))

    # step 3: Shards.lean
    if os.path.isfile(shards):
        with open(shards, encoding="utf-8", newline="") as fh:
            text = fh.read()
        if "\r" in text:
            probs.append("Shards.lean has CR characters")
        lines = text.split("\n")
        if not lines or lines[0] != "/-":
            probs.append("Shards.lean does not start with the header comment")
            body = lines
        else:
            try:
                close = lines.index("-/")
            except ValueError:
                close = len(lines)
                probs.append("Shards.lean header comment not closed")
            hdr = lines[1:close]
            if any("/-" in h or "-/" in h for h in hdr):
                probs.append("Shards.lean header nests comments")
            body = lines[close + 1:]
        thm = [l for l in body if l.startswith("theorem ")]
        if thm != SHARD_LINES:
            probs.append("Shards.lean theorem lines differ from the specification")
        other = [l for l in body if l not in SHARD_LINES and not l.startswith("import ")
                 and l not in SHARD_FRAMING]
        if other:
            probs.append("Shards.lean has lines outside run A's framing: %s" % other[:5])
        print("Shards.lean: %d theorem lines equal to the specification's six: %s; other lines: %d"
              % (len(thm), "yes" if thm == SHARD_LINES else "NO", len(other)))

    # step 4: All.lean
    allp = os.path.join(EXTRAS_DIR, "All.lean")
    if os.path.isfile(allp):
        lines = code_lines(audit, allp)
        imports = [c.strip()[len("import "):] for _, _, c in lines if c.startswith("import ")]
        if imports != ALL_IMPORTS:
            probs.append("Extras/All.lean imports differ from Extras.Shards + the 83 St_k_i_r modules")
        f1, f2, exc = audit.audit_theorems(lines, True)
        names = [m.group(3) for _, _, c in lines for m in [THEOREM_RE.match(c)] if m]
        defs = [m.group(4) for _, _, c in lines for m in [DEF_RE.match(c)] if m]
        if names != ["facts_%d" % k for k in range(5)] or f1 != 0 or f2 != 5 or exc or defs:
            probs.append("Extras/All.lean theorems %s (one-line %d, term-mode %d, exceptions %s, defs %s)"
                         % (names, f1, f2, exc, defs))
        nby = sum(1 for _, _, c in lines if audit.BY_RE.search(c))
        if nby:
            probs.append("Extras/All.lean contains `by` (%d lines)" % nby)
        print("Extras/All.lean: imports = Extras.Shards + %d St_k_i_r modules: %s; theorems %s; term-mode %d, "
              "one-line %d, exceptions %d, `by` lines %d"
              % (len(ALL_IMPORTS) - 1, "yes" if imports == ALL_IMPORTS else "NO", names, f2, f1, len(exc), nby))

    # step 6: the axioms script covers every non-private theorem of RSCStages and RSCCheckExtras
    ax = os.path.join(HERE, "AxiomsRSCCheckExtras.lean")
    with open(ax, encoding="utf-8") as fh:
        printed = set(re.findall(r"^#print axioms (\S+)\s*$", fh.read(), re.M))
    need = {"MaxLeafFormal.RSCCert.IsStages"}
    for fn, ns in (("RSCStages.lean", "MaxLeafFormal.RSCCert"), ("RSCCheckExtras.lean", "MaxLeafFormal.RSCCheck")):
        for _, _, c in code_lines(audit, os.path.join(LEAN_DIR, fn)):
            m = THEOREM_RE.match(c)
            if m and not m.group(1):
                need.add(ns + "." + m.group(3))
            if m and m.group(1) and m.group(1).strip() == "protected":
                need.add(ns + "." + m.group(3))
    missing = sorted(need - printed)
    print("AxiomsRSCCheckExtras.lean prints %d declarations; required %d (IsStages and every non-private "
          "theorem of RSCStages, RSCCheckExtras); missing: %s" % (len(printed), len(need), missing or "none"))
    if missing:
        probs.append("AxiomsRSCCheckExtras.lean misses %s" % missing)

    for p in probs:
        print("FINDING: " + p)
    ok = not probs
    print("EXTRAS AUDIT RESULT: " + ("PASS" if ok else "FAIL"))
    return 0 if ok else 1


if __name__ == "__main__":
    try:
        sys.exit(main())
    except Fail as e:
        print("ERROR: %s" % e)
        sys.exit(1)