"""Tiny Lean runner for the 03xx notebooks. Goes ONLY through align.verify.lean.

Rules (house + task): <= 2 workers (we use 1), nice 15, MemAvailable >= 45 GB checked
before the pool starts, never the heavy-job lock, never touches a GPU.
"""
from __future__ import annotations
import os, sys, re
os.environ["CUDA_VISIBLE_DEVICES"] = ""
sys.path.insert(0, "/apps/development/rlhf/align")
from align.verify.lean import ReplPool, ReplLimits, find_toolchain
from align.verify.lean.repl_pool import meminfo_available_gb

HEADER = "import Mathlib\nimport Aesop\nset_option maxHeartbeats 400000\nopen Nat\n"

class Lean:
    def __init__(self, header: str = HEADER):
        avail = meminfo_available_gb()
        if avail < 45.0:
            raise RuntimeError(f"MemAvailable {avail:.1f} GB < 45 GB: refusing to start Lean")
        self.tc = find_toolchain()
        self.pool = ReplPool.from_toolchain(self.tc, n_workers=1, limits=ReplLimits(nice=15))
        self.header = header
        self._cm = self.pool.acquire()
        self.w = self._cm.__enter__()
        self.env = self.w.header_env(header)

    def run(self, code: str, timeout: float = 120.0, keep: bool = False):
        """Run `code` on top of the header; return (messages, new_env)."""
        resp = self.w.command(code, self.env, timeout)
        if keep and "env" in resp:
            self.env = resp["env"]
        return resp

    def show(self, code: str, timeout: float = 120.0, keep: bool = False) -> dict:
        r = self.run(code, timeout, keep)
        for m in r.get("messages", []):
            print(f"[{m['severity']}] line {m['pos']['line']}: {m['data']}")
        for s in r.get("sorries", []):
            print("[sorries]", s.get("goal", "")[:200])
        if not r.get("messages") and not r.get("sorries"):
            print("[ok] no messages")
        return r

    def close(self):
        try: self._cm.__exit__(None, None, None)
        finally: self.pool.close()