#!/usr/bin/env python
"""Elaborated-type hash of formal statements (plan v3 Phase 0a item 5, first bullet).

Why: token/alpha hashes (p2_common.statement_sha / alpha_sha) are blind to hypothesis order, notation
(`x ∈ Finset.range n` vs `x < n`-style rewrites that elaborate to the same Expr), implicit coercion spelling, and
`open` differences. The elaborated type of `theorem _align_ref <sig> := by sorry` in the pinned Lean env is
not. Expr uses de Bruijn indices for bound variables, so binder names never reach the hash.

What is hashed (all decided in Lean, `LEAN_DUMP`, then canonicalised and hashed in pure Python):
  * the statement type is split with `forallTelescope` into binders and a conclusion;
  * instance-implicit binders are DROPPED, and so are instance-implicit ARGUMENTS of every application
    (`getFunInfo` paramInfo), so `@HAdd.hAdd ℝ ℝ ℝ inst a b` is `HAdd.hAdd ℝ ℝ ℝ a b`;
  * data binders (non-Prop) keep their order (later types may depend on earlier ones) and are named by that
    ordinal `v0, v1, ...`; Prop binders (hypotheses) are serialised with the variables as `v_k`, and the
    hypothesis list is SORTED, so reordering hypotheses does not change the hash;
  * universe levels are dropped; binder names/infos inside terms are dropped; `mdata` is stripped; numerals
    keep their type argument (`OfNat.ofNat ℝ 2` differs from `OfNat.ofNat ℕ 2`).
Known blind spots (by design, documented): permuting the order of data variables that are interchangeable is not
canonicalised; hypotheses that mention other hypotheses (proof-dependent types) are serialised with `?h`.

    # unit-testable pure parts: parse_dump, canonical_form, elab_sha, build_lean_prelude
    # <=50 statements validation (2 Lean workers max):
    .venv/bin/python scripts/p2/supply/elab_hash.py --in .tmp/p2/problems/minif2f.jsonl --out .tmp/p2/supply/elab_minif2f.jsonl --limit 50 --workers 2
    # FULL RUN (GPU-free window only; ~0.5-3 s/statement/worker after header warm-up; resumable by --out):
    .venv/bin/python scripts/p2/supply/elab_hash.py --in <set>.clean.jsonl --out <set>.elab.jsonl --workers 2
"""

from __future__ import annotations

import argparse
import hashlib
import importlib.util
import json
import os
import re
import sys
import time
from pathlib import Path
from typing import Any, Iterable, Optional

REPO = Path(__file__).resolve().parents[3]
sys.path.insert(0, str(REPO))

MARK = "ALIGN_ELAB_HASH "
NS = "AlignElabHash"
REF = "_align_elab_ref"

#: Trusted Lean prelude, elaborated once per worker BEFORE the header's `open`s (pristine namespace).
LEAN_DUMP = r'''
namespace AlignElabHash
open Lean Meta

partial def ser (nm : FVarId → Option String) (e : Expr) : MetaM String := do
  match e with
  | .bvar i => return s!"#{i}"
  | .fvar id => return (nm id).getD "?fv"
  | .mvar _ => return "?m"
  | .sort l => return s!"Sort({toString l})"
  | .const n _ => return n.toString
  | .lit (.natVal n) => return toString n
  | .lit (.strVal s) => return s!"\"{s}\""
  | .mdata _ b => ser nm b
  | .proj s i b => return s!"{s}.{i}({← ser nm b})"
  | .lam _ t b _ => return s!"(fun {← ser nm t} => {← ser nm b})"
  | .forallE _ t b _ => return s!"(all {← ser nm t} , {← ser nm b})"
  | .letE _ t v b _ => return s!"(let {← ser nm t} := {← ser nm v} in {← ser nm b})"
  | .app .. =>
    let f := e.getAppFn
    let args := e.getAppArgs
    let mut keep : Array Bool := args.map fun _ => true
    try
      let info ← getFunInfoNArgs f args.size
      for i in [0:args.size] do
        if i < info.paramInfo.size && info.paramInfo[i]!.isInstImplicit then
          keep := keep.set! i false
    catch _ => pure ()
    let mut out := #[← ser nm f]
    for i in [0:args.size] do
      if keep[i]! then out := out.push (← ser nm args[i]!)
    return "(" ++ " ".intercalate out.toList ++ ")"

def dump (n : Name) : Elab.Command.CommandElabM Unit := do
  let some ci := (← getEnv).find? n | throwError "no such declaration {n}"
  let js ← Elab.Command.liftTermElabM do
    forallTelescope ci.type fun xs concl => do
      let mut names : List (FVarId × String) := []
      let mut vars : Array String := #[]
      let mut hyps : Array String := #[]
      let mut k := 0
      for x in xs do
        let d ← x.fvarId!.getDecl
        if d.binderInfo.isInstImplicit then continue
        let nmf := fun (id : FVarId) => (names.find? (fun p => p.1 == id)).map (·.2)
        let ty ← ser nmf d.type
        if (← isProp d.type) then
          hyps := hyps.push ty
        else
          vars := vars.push ty
          names := (x.fvarId!, s!"v{k}") :: names
          k := k + 1
      let nmf := fun (id : FVarId) => (names.find? (fun p => p.1 == id)).map (·.2)
      let c ← ser nmf concl
      let hs := hyps.toList.map fun h => Json.str h
      return Json.compress (Json.mkObj [("vars", Json.arr (vars.toList.map Json.str).toArray), ("hyps", Json.arr hs.toArray), ("concl", Json.str c)])
  logInfo m!"ALIGN_ELAB_HASH {js}"

end AlignElabHash
'''


# --------------------------------------------------------------------------- pure parts


def build_lean_prelude() -> str:
    return LEAN_DUMP


def build_ref_command(sig_text: str) -> str:
    return f"theorem {REF} {sig_text} := by\n  sorry"


DUMP_COMMAND = f"#eval {NS}.dump `{REF}"


def parse_dump(messages: Iterable[dict]) -> dict:
    """Extract the JSON the Lean dump logged. Raises ValueError with the Lean error otherwise."""
    errs = []
    for m in messages:
        if not isinstance(m, dict):
            continue
        data = str(m.get("data", ""))
        if m.get("severity") == "error":
            errs.append(data)
        elif MARK in data:
            return json.loads(data.split(MARK, 1)[1].strip())
    raise ValueError("no dump in output" + (": " + errs[0][:200] if errs else ""))


def canonical_form(dump: dict) -> str:
    """Hypotheses sorted (multiset order-invariance); data binders and conclusion kept in order."""
    if not {"vars", "hyps", "concl"} <= set(dump):
        raise ValueError(f"dump lacks fields: {sorted(dump)}")
    return json.dumps({"vars": list(dump["vars"]), "hyps": sorted(dump["hyps"]), "concl": dump["concl"]},
                      ensure_ascii=False, separators=(",", ":"), sort_keys=True)


def elab_sha(dump: dict) -> str:
    return hashlib.sha256(canonical_form(dump).encode("utf-8")).hexdigest()


# --------------------------------------------------------------------------- Lean side


def _load_common():
    if "p2_common" in sys.modules:
        return sys.modules["p2_common"]
    spec = importlib.util.spec_from_file_location("p2_common", REPO / "scripts" / "p2" / "p2_common.py")
    mod = importlib.util.module_from_spec(spec)
    sys.modules["p2_common"] = mod
    spec.loader.exec_module(mod)
    return mod


def hash_statement(worker, statement: str, timeout_s: float = 120.0) -> dict:
    """Elaborate one statement in the pinned env and return {elab_sha, dump} or {error}."""
    from align.verify.lean.repl_pool import ReplError, ReplProtocolError
    from align.verify.lean.strict import parse_statement

    ref = parse_statement(statement)
    if not ref.ok:
        return {"error": f"unsupported_statement: {ref.problem}"}
    try:
        env = worker.header_env(ref.header_text, LEAN_DUMP)
    except ReplProtocolError as e:
        return {"error": "header_error: " + " ".join(str(e).split())[:200]}
    except ReplError as e:
        return {"error": e.code if hasattr(e, "code") else type(e).__name__}
    try:
        r1 = worker.command(build_ref_command(ref.sig_text), env, timeout_s)
        errs = [m for m in r1.get("messages", []) if isinstance(m, dict) and m.get("severity") == "error"]
        if errs or "env" not in r1:
            return {"error": "elab_error: " + " ".join(str(errs[0].get("data", r1) if errs else r1).split())[:200]}
        r2 = worker.command(DUMP_COMMAND, r1["env"], timeout_s)
        dump = parse_dump(r2.get("messages", []))
    except ReplError as e:
        return {"error": e.code if hasattr(e, "code") else type(e).__name__}
    except ValueError as e:
        return {"error": f"dump_error: {e}"}
    return {"elab_sha": elab_sha(dump), "dump": dump}


def main(argv=None) -> int:
    ap = argparse.ArgumentParser(description=__doc__.split("\n")[0])
    ap.add_argument("--in", dest="inp", required=True)
    ap.add_argument("--out", required=True, help="JSONL {problem_id, elab_sha | error}; existing ids are skipped (resume)")
    ap.add_argument("--limit", type=int, default=0)
    ap.add_argument("--workers", type=int, default=2)
    ap.add_argument("--timeout", type=float, default=120.0)
    a = ap.parse_args(argv)
    if not 1 <= a.workers <= 2:
        raise SystemExit("--workers must be 1 or 2 (>2 needs the heavy-job lock; do the full run in a GPU-free window)")
    C = _load_common()
    from align.utils import guard
    from align.verify.lean import ReplPool, find_toolchain
    from align.verify.lean.repl_pool import ReplLimits

    guard.preflight_host_check(strict=True)
    done = set()
    if os.path.exists(a.out):
        done = {r["problem_id"] for r in C.iter_jsonl(a.out)}
    rows = []
    for r in C.iter_jsonl(a.inp):
        if r["problem_id"] not in done:
            rows.append(r)
        if a.limit and len(rows) >= a.limit:
            break
    tc = find_toolchain()
    pool = ReplPool.from_toolchain(tc, n_workers=a.workers, limits=ReplLimits(), preflight=True,
                                   scratch_dir=str(Path(os.environ.get("TMPDIR", str(REPO / ".tmp"))) / "p2_elab_pool"))
    t0 = time.time()
    ok = 0
    with C.AppendJsonl(a.out) as out:
        for i in range(0, len(rows), 16):
            chunk = rows[i : i + 16]
            res = pool.map(lambda w, r: hash_statement(w, r["formal_statement"], a.timeout), chunk)
            for r, x in zip(chunk, res):
                if isinstance(x, BaseException):
                    raise x
                out.append({"problem_id": r["problem_id"], **({"elab_sha": x["elab_sha"]} if "elab_sha" in x else {"error": x["error"]})})
                ok += "elab_sha" in x
    pool.close()
    print(json.dumps({"statements": len(rows), "hashed": ok, "errors": len(rows) - ok, "seconds": round(time.time() - t0, 1)}))
    return 0


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