#!/usr/bin/env python3
"""
verify_lean.py — 5-check verification for Lean formalizations.

Adapted from Peng et al. (2026) pipeline-math verify.sh.
Standalone — no LLM, just verification.

Usage:
  python3 verify_lean.py [--target LeanModule] [--file path/to/Foo.lean]
"""
import os, sys, re, hashlib, subprocess, json, argparse
from pathlib import Path

SILVERSIGHT_DIR = os.path.expanduser("~/SilverSight")
ALLOWED_AXIOMS = {"propext", "Classical.choice", "Quot.sound"}
BANNED_KEYWORDS = ["sorryAx", "native_decide", "admit", "unsafe",
                   "implemented_by", "ofReduceBool"]

def sha256_of(path):
    h = hashlib.sha256()
    with open(path, "rb") as f:
        h.update(f.read())
    return h.hexdigest()

def strip_comments(s):
    out = []
    i, n = 0, len(s)
    depth = 0
    while i < n:
        two = s[i:i+2]
        if depth == 0 and two == "--":
            j = s.find("\n", i)
            i = j if j != -1 else n
        elif two == "/-":
            depth += 1
            i += 2
        elif depth > 0 and two == "-/":
            depth -= 1
            i += 2
        elif depth > 0:
            i += 1
        else:
            out.append(s[i])
            i += 1
    return "".join(out)

def lake_build(target=None, timeout=120):
    cmd = ["lake", "build"]
    if target:
        cmd.append(target)
    try:
        result = subprocess.run(cmd, cwd=SILVERSIGHT_DIR, capture_output=True,
                                text=True, timeout=timeout)
        output = result.stdout + result.stderr
        errors = output.count("error:")
        warnings = output.count("warning:") - output.count("declaration uses .sorry.")
        return errors, warnings, output
    except subprocess.TimeoutExpired:
        return 999, 0, "TIMEOUT"
    except Exception as e:
        return 999, 0, str(e)

# ── Checks ──────────────────────────────────────────────────────────

def check_sha_pins():
    pins_path = Path(SILVERSIGHT_DIR) / "scripts" / "frozen.sha256"
    if not pins_path.exists():
        return True, "no pins file"
    ok = True
    for line in pins_path.read_text().splitlines():
        line = line.strip()
        if not line or line.startswith("#"):
            continue
        parts = line.split(None, 1)
        if len(parts) != 2:
            continue
        pinned, rel = parts[0], parts[1]
        fpath = os.path.join(SILVERSIGHT_DIR, rel)
        if not os.path.exists(fpath):
            print(f"  FAIL: pinned file missing: {rel}")
            ok = False
            continue
        actual = sha256_of(fpath)
        if pinned != actual:
            print(f"  FAIL: {rel} SHA mismatch")
            print(f"    pinned: {pinned}")
            print(f"    actual: {actual}")
            ok = False
        else:
            print(f"  PASS: {rel} pin OK")
    return ok, "SHA pin check complete"

def check_banned(source_dirs):
    ok = True
    for d in source_dirs:
        dpath = os.path.join(SILVERSIGHT_DIR, d)
        if not os.path.isdir(dpath):
            continue
        for root, dirs, files in os.walk(dpath):
            for f in files:
                if not f.endswith(".lean"):
                    continue
                fpath = os.path.join(root, f)
                code = strip_comments(open(fpath).read())
                for kw in BANNED_KEYWORDS:
                    if re.search(r'\b' + re.escape(kw) + r'\b', code):
                        print(f"  FAIL: {os.path.relpath(fpath, SILVERSIGHT_DIR)} contains `{kw}`")
                        ok = False
    return ok, "banned keyword check complete"

def check_build(target=None):
    err, warn, out = lake_build(target=target)
    if err == 0 and warn == 0:
        print(f"  PASS: build clean ({err} errors, {warn} warnings)")
        return True, "build OK"
    print(f"  FAIL: {err} errors, {warn} unexpected warnings")
    for line in out.splitlines():
        if "error:" in line or ("warning:" in line and "sorry" not in line):
            print(f"    {line.strip()}")
    return False, f"build FAIL ({err} err, {warn} warn)"

def check_axioms(module_name):
    ax_file = os.path.join(SILVERSIGHT_DIR, ".lake", "tmp_ax.lean")
    os.makedirs(os.path.dirname(ax_file), exist_ok=True)
    with open(ax_file, "w") as f:
        f.write(f"import {module_name}\n#print axioms {module_name}\n")
    try:
        r = subprocess.run(["lake", "env", "lean", ax_file], cwd=SILVERSIGHT_DIR,
                           capture_output=True, text=True, timeout=60)
        os.remove(ax_file)
        m = re.search(r"'([^']+)' depends on axioms:\s*\[([^\]]+)\]", r.stdout + r.stderr)
        if m:
            axioms = set(a.strip() for a in m.group(2).split(",") if a.strip())
            bad = axioms - ALLOWED_AXIOMS
            if bad:
                print(f"  FAIL: unexpected axioms: {bad}")
                return False, "unexpected axioms"
            print(f"  PASS: axioms = {axioms}")
            return True, "axioms OK"
        print("  PASS: no axiom info")
        return True, "no axioms"
    except Exception as e:
        return False, str(e)

def check_discharge():
    dpath = os.path.join(SILVERSIGHT_DIR, "formal", "CoreFormalism", "Discharge.lean")
    if not os.path.exists(dpath):
        return True, "no Discharge.lean (SKIP)"
    err, warn, out = lake_build(target="CoreFormalism.Discharge", timeout=60)
    if err == 0:
        print("  PASS: Discharge gate compiles")
        return True, "Discharge OK"
    print(f"  FAIL: Discharge gate ({err} errors)")
    return False, "Discharge FAIL"

def main():
    parser = argparse.ArgumentParser(description="5-check Lean verification")
    parser.add_argument("--target", help="Lake build target")
    parser.add_argument("--source-dirs", nargs="*", default=["formal"],
                        help="Source dirs to check for banned keywords")
    parser.add_argument("--module", help="Module for #print axioms check")
    args = parser.parse_args()

    checks = [
        ("1. SHA pins", lambda: check_sha_pins()),
        ("2. Banned keywords", lambda: check_banned(args.source_dirs)),
        ("3. lake build", lambda: check_build(args.target)),
    ]
    if args.module:
        checks.append(("4. #print axioms", lambda: check_axioms(args.module)))
    checks.append(("5. Discharge gate", check_discharge))

    errors = 0
    for name, fn in checks:
        print(f"\n─── {name} ───")
        ok, msg = fn()
        if not ok:
            errors += 1

    print(f"\n{'='*50}")
    print(f"RESULT: {'PASS' if errors == 0 else 'FAIL'} ({errors} check(s) failed)")
    return errors

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