#!/usr/bin/env python3
"""Compile the pinned official FPS layer and instantiate both macros.

The archived Basic.lean imports Mathlib, but the file itself uses only the
Lean core/elaborator API.  The temporary Mathlib module below is therefore a
dependency shim, not a rewritten copy of Basic.lean: the source file is
extracted byte-for-byte and compiled as the imported module.
"""

from __future__ import annotations

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


ARCHIVE = Path("source/official-current.tar.gz")
SOURCE_SUFFIX = "/data/formal_problem_solving/FormalProblemSolving/Basic.lean"


AUDIT = r'''import FormalProblemSolving.Basic

open Lean Elab Command Term Meta Tactic
open scoped FPS DFPS

set_option autoImplicit false
set_option relaxedAutoImplicit false

problem source_level_direct
  find (n : Nat) s.t.
  : n = 2 := solve
  exact rfl

example : source_level_direct.Answer = 2 := by rfl

problem source_level_find_all
  find_all (answer_predicate : Prop) iff
  (hH : True)
  : ((True : Prop) ↔ answer_predicate) := solve
  case Forward => exact _h_conclusion
  case Backward => exact _h_answer

example : (source_level_find_all True).Answer := by
  exact ((source_level_find_all True).Proof True.intro).mp Iff.rfl
example : (True : Prop) ↔ (source_level_find_all True).Answer := by
  simpa [ProblemTarget] using (source_level_find_all True).Proof True.intro

theorem constructor_soundness {α : Type} {P : α → Prop}
    (s : ProblemSol P) : P s.Answer := s.Proof

theorem find_all_has_both_directions {H C : Prop}
    (s : ProblemSol (fun answer : Prop => H → (C ↔ answer)))
    (hH : H) :
    (C → s.Answer) ∧ (s.Answer → C) := by
  constructor
  · intro hC
    exact (s.Proof hH).mp hC
  · intro hA
    exact (s.Proof hH).mpr hA

example : (2 : Nat) = 2 :=
  constructor_soundness source_level_direct

example : (True : Prop) := by
  have h := find_all_has_both_directions (source_level_find_all True) True.intro
  have hA : (source_level_find_all True).Answer :=
    ((source_level_find_all True).Proof True.intro).mp Iff.rfl
  exact (h.2 hA).mp True.intro
'''


def extract_basic(destination: Path) -> tuple[bytes, str]:
    with tarfile.open(ARCHIVE, "r:gz") as archive:
        member = next(
            item for item in archive.getmembers() if item.name.endswith(SOURCE_SUFFIX)
        )
        handle = archive.extractfile(member)
        if handle is None:
            raise RuntimeError(member.name)
        raw = handle.read()
    target = destination / "FormalProblemSolving" / "Basic.lean"
    target.parent.mkdir(parents=True)
    target.write_bytes(raw)
    return raw, member.name


def main() -> None:
    lean = shutil.which("lean") or "/Users/sshpro/.elan/bin/lean"
    if not Path(lean).exists():
        raise RuntimeError("Lean 4.15.0 executable is unavailable")
    with tempfile.TemporaryDirectory(prefix="fps-source-audit-") as temp:
        root = Path(temp)
        raw, member_name = extract_basic(root)
        # Basic.lean imports Mathlib, but all of its declarations use Lean's
        # bundled API.  This shim makes that dependency boundary explicit.
        (root / "Mathlib.lean").write_text("import Lean\n", encoding="utf-8")
        audit = root / "official_audit.lean"
        audit.write_text(AUDIT, encoding="utf-8")
        environment = os.environ.copy()
        environment["LEAN_PATH"] = str(root)
        compile_shim = subprocess.run(
            [lean, "-o", str(root / "Mathlib.olean"), str(root / "Mathlib.lean")],
            cwd=root,
            env=environment,
            capture_output=True,
            text=True,
            check=False,
        )
        if compile_shim.returncode:
            raise RuntimeError(
                "Mathlib dependency shim compilation failed\n"
                + compile_shim.stdout
                + compile_shim.stderr
            )
        basic = root / "FormalProblemSolving" / "Basic.lean"
        compile_basic = subprocess.run(
            [lean,
             "-o",
             str(root / "FormalProblemSolving" / "Basic.olean"),
             str(basic)],
            cwd=root,
            env=environment,
            capture_output=True,
            text=True,
            check=False,
        )
        if compile_basic.returncode:
            raise RuntimeError(
                "official Basic.lean compilation failed\n"
                + compile_basic.stdout
                + compile_basic.stderr
            )
        command = [lean, str(audit)]
        completed = subprocess.run(
            command,
            cwd=root,
            env=environment,
            capture_output=True,
            text=True,
            check=False,
        )
        if completed.returncode:
            raise RuntimeError(
                "official Basic.lean audit failed\n"
                + completed.stdout
                + completed.stderr
            )
        version = subprocess.run(
            [lean, "--version"], capture_output=True, text=True, check=True
        ).stdout.strip()
        print(
            json.dumps(
                {
                    "basic_member": member_name,
                    "basic_sha256": hashlib.sha256(raw).hexdigest(),
                    "lean": version,
                    "exit_status": completed.returncode,
                    "tests": [
                        "official Basic.lean imported byte-for-byte",
                        "source-level FPS problem/solve macro",
                        "source-level DFPS problem/find_all/solve macro",
                        "kernel-checked ProblemSol soundness projection",
                        "kernel-checked forward and backward find-all directions",
                    ],
                },
                sort_keys=True,
            )
        )


if __name__ == "__main__":
    main()