#!/usr/bin/env python3
"""Run the official FPS/DFPS sources under their pinned Mathlib environment.

The audit deliberately uses the archived project and its Lake manifest.  It
does not replace Mathlib with a shim or reimplement the examples.  With no
argument it creates a temporary checkout, fetches the manifest-pinned
dependencies, builds them, and runs the three official Lean files.  The
``--project`` option is useful for rerunning the same checks in an already
built checkout without rebuilding it.
"""

from __future__ import annotations

import argparse
import hashlib
import json
import os
import shutil
import subprocess
import tarfile
import tempfile
from pathlib import Path


REPO = Path(__file__).resolve().parents[1]
ARCHIVE = REPO / "source" / "official-current.tar.gz"
EXPECTED_BASIC_SHA256 = "a908c060d1031506ac5844b838b819391a001474eee88a5df184fa9bdac97dc4"
FILES = (
    "FormalProblemSolving/Basic.lean",
    "FormalProblemSolving/FPS_Example.lean",
    "FormalProblemSolving/DFPS_Example.lean",
)


def run(command: list[str], cwd: Path, env: dict[str, str]) -> subprocess.CompletedProcess[str]:
    return subprocess.run(
        command,
        cwd=cwd,
        env=env,
        capture_output=True,
        text=True,
        check=False,
    )


def project_from_archive(destination: Path) -> Path:
    with tarfile.open(ARCHIVE, "r:gz") as archive:
        archive.extractall(destination)
        top = archive.getnames()[0].split("/", 1)[0]
    project = destination / top / "data" / "formal_problem_solving"
    packages = project / ".lake" / "packages"
    if packages.is_symlink():
        packages.unlink()
    packages.mkdir(parents=True, exist_ok=True)
    return project


def main() -> None:
    parser = argparse.ArgumentParser()
    parser.add_argument(
        "--project",
        type=Path,
        help="already prepared formal_problem_solving project; skips extraction and Lake update/build",
    )
    args = parser.parse_args()
    lake = os.environ.get("LAKE", "/Users/sshpro/.elan/bin/lake")
    if not Path(lake).exists():
        lake = shutil.which("lake") or lake
    if not Path(lake).exists():
        raise SystemExit("lake executable is unavailable")

    basic_bytes = None
    project = args.project.resolve() if args.project else None
    with tempfile.TemporaryDirectory(prefix="fps-full-audit-") as temporary:
        if project is None:
            project = project_from_archive(Path(temporary))
            update = run([lake, "update", "-v"], project, os.environ.copy())
            if update.returncode:
                raise SystemExit(update.stdout + update.stderr)
            build = run([lake, "build"], project, os.environ.copy())
            if build.returncode:
                raise SystemExit(build.stdout + build.stderr)
        basic = project / FILES[0]
        basic_bytes = basic.read_bytes()
        if hashlib.sha256(basic_bytes).hexdigest() != EXPECTED_BASIC_SHA256:
            raise SystemExit("archived Basic.lean hash does not match the pinned source")
        results = {}
        environment = os.environ.copy()
        for relative in FILES:
            completed = run([lake, "env", "lean", relative], project, environment)
            results[relative] = {
                "exit_status": completed.returncode,
                "stdout": completed.stdout,
                "stderr": completed.stderr,
            }
            if completed.returncode:
                raise SystemExit(json.dumps(results, indent=2, sort_keys=True))
        print(
            json.dumps(
                {
                    "basic_sha256": hashlib.sha256(basic_bytes).hexdigest(),
                    "lake_project": str(project),
                    "manifest": "lake-manifest.json",
                    "official_files": results,
                    "full_dependency_environment": True,
                },
                indent=2,
                sort_keys=True,
            )
        )


if __name__ == "__main__":
    main()