#!/usr/bin/env python3
"""Аудит Lean4-доказательств Balansis для сайта.

Скрипт делает ровно то, что обещает раздел «Доказательства» на сайте:

1. Разбирает исходники ``formal/ACT/*.lean``, собирая полные имена всех теорем.
2. Проверяет, что в исходниках нет ``sorry`` / ``axiom`` / ``admit``.
3. Собирает библиотеку (``lake build``) и для каждой теоремы спрашивает у Lean
   ``#print axioms`` — то есть на какие аксиомы она фактически опирается после
   проверки ядром, а не по утверждению README.
4. Пишет ``data/lean-audit.json`` с результатом, версией toolchain, ревизией
   Mathlib и кодами возврата каждой команды.

Скрипт не изменяет репозиторий Balansis: временный Lean-файл создаётся вне
дерева, а артефакты сборки лежат в игнорируемом ``formal/.lake``.

Запуск:
    BALANSIS_REPO=/path/to/Balansis python scripts/generate-lean-audit.py \
        --out data/lean-audit.json
"""
from __future__ import annotations

import argparse
import json
import os
import re
import subprocess
import sys
import tempfile
from datetime import datetime, timezone
from pathlib import Path
from typing import Dict, List, Optional

DECL_RE = re.compile(r"^(?P<kw>theorem|lemma)\s+(?P<name>[^\s({\[:]+)")
NS_OPEN_RE = re.compile(r"^namespace\s+(\S+)")
NS_CLOSE_RE = re.compile(r"^end\s+(\S+)")
FORBIDDEN_RE = re.compile(r"\b(sorry|admit)\b|^\s*axiom\s")
AXIOMS_RE = re.compile(r"^'(?P<name>[^']+)' depends on axioms: \[(?P<axioms>[^\]]*)\]")
NO_AXIOMS_RE = re.compile(r"^'(?P<name>[^']+)' does not depend on any axioms")
TRUSTED = {"propext", "Classical.choice", "Quot.sound"}


def run(cmd: List[str], cwd: Path, timeout: int = 3600) -> Dict[str, object]:
    started = datetime.now(timezone.utc)
    proc = subprocess.run(cmd, cwd=str(cwd), capture_output=True, text=True, timeout=timeout)
    # Абсолютные пути рабочей машины в артефакт не попадают: команда
    # публикуется в обобщённом виде.
    printable = " ".join(cmd)
    for part in cmd:
        if part.startswith("/") and part.endswith(".lean") and "/tmp" in part:
            printable = printable.replace(part, f"<tmp>/{Path(part).name}")
    return {
        "command": printable,
        "exit_code": proc.returncode,
        "started_at_utc": started.isoformat(),
        "duration_seconds": (datetime.now(timezone.utc) - started).total_seconds(),
        # полный вывод нужен разбору #print axioms; в JSON уходит только хвост
        "stdout": proc.stdout,
        "stdout_tail": proc.stdout[-4000:],
        "stderr_tail": proc.stderr[-4000:],
    }


def collect_theorems(lean_file: Path) -> List[Dict[str, object]]:
    """Полные имена теорем файла с учётом стека namespace."""
    stack: List[str] = []
    found: List[Dict[str, object]] = []
    for lineno, raw in enumerate(lean_file.read_text(encoding="utf-8").splitlines(), 1):
        line = raw.strip()
        m = NS_OPEN_RE.match(line)
        if m:
            stack.append(m.group(1))
            continue
        m = NS_CLOSE_RE.match(line)
        if m:
            if stack and stack[-1] == m.group(1):
                stack.pop()
            continue
        m = DECL_RE.match(line)
        if m:
            found.append({
                "name": ".".join(stack + [m.group("name")]),
                "short_name": m.group("name"),
                "kind": m.group("kw"),
                "file": lean_file.name,
                "line": lineno,
                "statement": raw.strip(),
            })
    return found


def scan_forbidden(root: Path) -> Dict[str, object]:
    hits = []
    for path in sorted(root.rglob("*.lean")):
        if ".lake" in path.parts:
            continue
        for lineno, line in enumerate(path.read_text(encoding="utf-8").splitlines(), 1):
            if FORBIDDEN_RE.search(line):
                hits.append({"file": str(path.relative_to(root)), "line": lineno, "text": line.strip()})
    return {"scanned_files": len([p for p in root.rglob("*.lean") if ".lake" not in p.parts]),
            "hits": hits, "clean": not hits}


def main() -> int:
    ap = argparse.ArgumentParser(description=__doc__)
    # Путь к формальному слою берётся из BALANSIS_REPO либо ищется рядом с
    # репозиторием сайта: абсолютных путей конкретной машины в скрипте нет.
    site_root = Path(__file__).resolve().parent.parent
    env_repo = os.environ.get("BALANSIS_REPO", "")
    default_formal = next(
        (str(c / "formal") for c in (
            Path(env_repo) if env_repo else None,
            site_root / "Balansis",
            site_root.parent / "Balansis",
        ) if c and (c / "formal" / "lakefile.lean").exists()),
        "Balansis/formal",
    )
    ap.add_argument("--formal", default=default_formal,
                    help="каталог formal/ репозитория Balansis (или переменная BALANSIS_REPO)")
    ap.add_argument("--out", default="data/lean-audit.json")
    ap.add_argument("--skip-build", action="store_true",
                    help="не запускать lake build (использовать готовые .olean)")
    args = ap.parse_args()

    formal = Path(args.formal).resolve()
    if not (formal / "lakefile.lean").exists():
        print(f"не найден lakefile.lean в {formal}", file=sys.stderr)
        return 2

    toolchain = (formal / "lean-toolchain").read_text().strip()
    manifest = json.loads((formal / "lake-manifest.json").read_text())
    mathlib = next((p for p in manifest.get("packages", []) if p.get("name") == "mathlib"), {})

    theorems: List[Dict[str, object]] = []
    for lean_file in sorted((formal / "ACT").glob("*.lean")):
        theorems.extend(collect_theorems(lean_file))

    steps: List[Dict[str, object]] = []
    if not args.skip_build:
        steps.append(run(["lake", "build", "BalansisFormal"], formal))
        steps.append(run(["lake", "build", "ACT"], formal))
    steps.append(run(["lake", "env", "lean", "FormalAudit.lean"], formal))

    # #print axioms по каждой теореме — вне дерева репозитория
    with tempfile.TemporaryDirectory() as tmp:
        probe = Path(tmp) / "SiteAxiomAudit.lean"
        body = ["import ACT", ""] + [f"#print axioms {t['name']}" for t in theorems]
        probe.write_text("\n".join(body) + "\n", encoding="utf-8")
        axioms_step = run(["lake", "env", "lean", str(probe)], formal)
        steps.append(axioms_step)
        raw_axioms = axioms_step["stdout"]

    by_name: Dict[str, List[str]] = {}
    for line in str(raw_axioms).splitlines():
        stripped = line.strip()
        m = AXIOMS_RE.match(stripped)
        if m:
            axioms = [a.strip() for a in m.group("axioms").split(",") if a.strip()]
            by_name[m.group("name")] = axioms
            continue
        m = NO_AXIOMS_RE.match(stripped)
        if m:
            by_name[m.group("name")] = []

    untrusted: List[Dict[str, object]] = []
    unresolved: List[str] = []
    for t in theorems:
        axioms = by_name.get(str(t["name"]))
        if axioms is None:
            t["axioms"] = None
            t["axioms_trusted"] = None
            unresolved.append(str(t["name"]))
            continue
        extra = [a for a in axioms if a not in TRUSTED]
        t["axioms"] = axioms
        t["axioms_trusted"] = not extra
        if extra:
            untrusted.append({"theorem": t["name"], "extra_axioms": extra})

    forbidden = scan_forbidden(formal)

    for step in steps:
        step.pop("stdout", None)   # в артефакт идёт только хвост, чтобы файл оставался читаемым

    payload = {
        "schema": "balansis-site/lean-audit/1",
        "run": {
            "generated_at_utc": datetime.now(timezone.utc).isoformat(),
            "generator": "scripts/generate-lean-audit.py",
            "formal_dir": "<Balansis>/formal",
        },
        "toolchain": {
            "lean_toolchain": toolchain,
            "lake_version": str(run(["lake", "--version"], formal)["stdout_tail"]).strip(),
            "mathlib_rev": mathlib.get("rev"),
            "mathlib_input_rev": mathlib.get("inputRev"),
            "mathlib_url": mathlib.get("url"),
        },
        "steps": steps,
        "summary": {
            "theorems_declared": len(theorems),
            "theorems_with_axiom_report": len(by_name),
            "unresolved": unresolved,
            "all_axioms_trusted": not untrusted and not unresolved,
            "trusted_axiom_set": sorted(TRUSTED),
            "untrusted": untrusted,
            "forbidden_tokens": forbidden,
            "all_steps_succeeded": all(s["exit_code"] == 0 for s in steps),
        },
        "theorems": theorems,
    }

    out = Path(args.out)
    out.parent.mkdir(parents=True, exist_ok=True)
    out.write_text(json.dumps(payload, ensure_ascii=False, indent=2) + "\n", encoding="utf-8")
    print(f"written: {out}; theorems: {len(theorems)}; "
          f"axiom reports: {len(by_name)}; steps ok: {payload['summary']['all_steps_succeeded']}")
    return 0


if __name__ == "__main__":
    raise SystemExit(main())