#!/usr/bin/env python
"""Stream candidate Lean statement sets into the P2 Problem JSONL schema (plan v3, Phase 0a item 4).

Sources (docs/workbook/data/Datasets.json): FormalMATH-All, NuminaMath-LEAN-Sol,
Nemotron-Math-Proofs-v1, RL-Goedel-Pset (levels 3-5 / 4-5), Lean4-FormalConjectures.

Properties
  * Streaming and memory-bounded: JSONL files are read line by line over HTTP (Range-resumable,
    `align.data.hub.iter_lines`), parquet-only datasets through the datasets-server JSON /rows API
    (no pyarrow). Nothing is ever held beyond one row plus the per-source dedup set of statement shas.
  * Every row is pinned to the dataset revision sha (HF API, resolved once at start) and carries
    `provenance = {source_id, revision, row_index, license, ...}` plus an `ancestry` record
    (`root` = own problem_id, `parents` = []), the field `align.data.decontam` propagates to mutation
    children.
  * Output rows are PROBLEM_FIELDS (+ `provenance`, `ancestry`): `formal_statement` is
    `GOEDEL_HEADER + ... := by sorry` (the header is REWRITTEN when the source's differs, recorded in
    `provenance.header_rewritten` / `original_header`; whether the rewrite still elaborates is the
    elaboration check's job, `elab_hash.py` / `build_problems.py --phase check`).
  * Nothing is dropped silently: every drop is counted by reason in `ingest_manifest.json` and the first
    few examples per reason are kept in `<source>.drops.sample.json`.

    HF_HOME=/apps/development/rlhf/hf_cache HF_HUB_DISABLE_XET=1 \
      .venv/bin/python scripts/p2/supply/ingest.py --out-dir .tmp/p2/supply --workers 2
"""

from __future__ import annotations

import argparse
import importlib.util
import json
import logging
import os
import re
import sys
import time
import urllib.request
from collections import Counter
from concurrent.futures import ThreadPoolExecutor
from pathlib import Path
from typing import Any, Callable, Iterator

REPO = Path(__file__).resolve().parents[3]
sys.path.insert(0, str(REPO))
os.environ.setdefault("HF_HOME", "/apps/development/rlhf/hf_cache")
os.environ.setdefault("HF_HUB_DISABLE_XET", "1")


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


C = _load_common()
from align.data import hub  # noqa: E402  (after sys.path)

log = logging.getLogger("p2.supply.ingest")
MAX_STATEMENT_CHARS = 4000  # same cap as build_problems.py
DEFAULT_OUT = Path(os.environ.get("ALIGN_P2_TMP", "/apps/development/rlhf/align/.tmp/p2")) / "supply"
_BY_SORRY_TAIL = re.compile(r":=\s*sorry\s*$")
_BY_TAIL = re.compile(r":=\s*by\s*$")


# --------------------------------------------------------------------------- source registry


class Source:
    def __init__(self, key, repo_id, license_, iterate, split="train", problem_source=None, note=""):
        self.key, self.repo_id, self.license, self.iterate = key, repo_id, license_, iterate
        self.split, self.problem_source, self.note = split, problem_source or key, note


def _it_formalmath(src: Source, rev: str, args) -> Iterator[dict]:
    """FormalMATH_all.json is one 13 MB JSON array (the repo has no JSONL); read once, yield rows."""
    req = urllib.request.Request(hub.resolve_url(src.repo_id, "FormalMATH_all.json", rev))
    with urllib.request.urlopen(req, timeout=120) as r:
        data = json.loads(r.read())
    for i, row in enumerate(data):
        yield {"row_index": i, "statement": row.get("autoformalization"), "header": None,
               "informal": row.get("refined_statement"), "name": row.get("theorem_names"),
               "extra": {"domain": row.get("domain"), "orig_source": row.get("source")}}


def _it_numina(src: Source, rev: str, args) -> Iterator[dict]:
    for i, raw in enumerate(hub.iter_lines(hub.resolve_url(src.repo_id, "train.jsonl", rev))):
        row = json.loads(raw) if raw.strip() else None
        if row is None:
            continue
        yield {"row_index": i, "statement": row.get("formal_statement"), "header": None,
               "informal": row.get("problem"), "name": row.get("uuid"),
               "english": str(row.get("is_english")).lower() != "false",
               "extra": {"uuid": row.get("uuid"), "orig_source": row.get("source"),
                         "statement_source": row.get("statement_source")}}


def _it_nemotron(src: Source, rev: str, args) -> Iterator[dict]:
    """29.5 GB file (whole trajectories): never downloaded in full. `--nemotron-stripes` evenly spaced
    Range windows of `--nemotron-rows/stripes` lines each, so the sample is not biased to the file head.
    Row locator = byte offset of the line (stable for a pinned revision)."""
    url = hub.resolve_url(src.repo_id, "data/lean.jsonl", rev)
    req = urllib.request.Request(url, method="HEAD")
    with hub._urlopen(req) as r:  # follows the redirect
        size = int(r.headers["Content-Length"])
    stripes = max(1, args.nemotron_stripes)
    per = max(1, args.nemotron_rows // stripes)
    for s in range(stripes):
        start = (size // stripes) * s
        pos = start
        n = 0
        first = start > 0  # the first line after a mid-file Range is a fragment: skip it
        for raw in hub.iter_lines(url, start_byte=start):
            here = pos
            pos += len(raw) + 1
            if first:
                first = False
                continue
            if not raw.strip():
                continue
            try:
                row = json.loads(raw)
            except json.JSONDecodeError:
                continue  # truncated last line of the file only
            yield {"row_index": here, "statement": row.get("formal_statement"), "header": row.get("lean_header"),
                   "informal": row.get("problem"), "name": row.get("uuid"),
                   "extra": {"uuid": row.get("uuid"), "orig_source": row.get("source"), "byte_offset": here,
                             "stripe": s, "row_license": row.get("license")}}
            n += 1
            if n >= per:
                break


def _parquet_or_server_rows(repo_id: str, rev: str, max_rows: int | None) -> Iterator[dict]:
    """Rows of the single `train` split. With pyarrow (run under .venv-rollout; the import lives in
    align.data.ingest_parquet, the one whitelisted module) the pinned parquet shard(s) are fetched once to
    HF_HOME and streamed in row batches. Without it: the datasets-server JSON /rows API, paced at 1 page/s
    with a long backoff on HTTP 429 (the server rate-limits sequential paging)."""
    from align.data import ingest_parquet as ip

    try:
        ip._require_pyarrow()
        have_pa = True
    except ip.IngestUnavailable:
        have_pa = False
    if have_pa:
        dest = Path(os.environ["HF_HOME"]) / "align_streams" / "supply" / repo_id.replace("/", "__") / rev
        dest.mkdir(parents=True, exist_ok=True)
        files = [e for e in hub.list_repo_files(repo_id, rev, "data") if e.get("type") == "file" and e["path"].endswith(".parquet")]
        emitted = 0
        for e in sorted(files, key=lambda e: e["path"]):
            target = dest / Path(e["path"]).name
            want = (e.get("lfs") or {}).get("oid")
            if not (target.exists() and target.stat().st_size == e["size"]):
                tmp = target.with_name(target.name + f".tmp{os.getpid()}")
                try:
                    with urllib.request.urlopen(hub.resolve_url(repo_id, e["path"], rev), timeout=300) as r, open(tmp, "wb") as f:
                        while chunk := r.read(4 << 20):
                            f.write(chunk)
                    if want and C.sha256_file(tmp) != want:
                        raise RuntimeError(f"{e['path']}: sha256 differs from the hub's LFS oid")
                    os.replace(tmp, target)
                finally:
                    if tmp.exists():
                        tmp.unlink()
            for row in ip.iter_parquet_rows(target, prefer_ingested=False):
                if max_rows and emitted >= max_rows:
                    return
                emitted += 1
                yield row
        return
    offset, emitted = 0, 0
    while not max_rows or emitted < max_rows:
        for attempt in range(8):
            try:
                page = hub.fetch_rows_page(repo_id, "default", "train", offset, 100, use_cache=False)
                break
            except hub.HubError as err:
                if "429" not in str(err) or attempt == 7:
                    raise
                time.sleep(60)
        rows = page["rows"]
        if not rows:
            return
        for r in rows:
            if r.get("truncated_cells"):
                raise hub.HubError(f"truncated cells at row {r.get('row_idx')}")
            if max_rows and emitted >= max_rows:
                return
            emitted += 1
            yield r["row"]
        offset += len(rows)
        if offset >= page.get("num_rows_total", offset + 1):
            return
        time.sleep(1.0)


def _it_rows(src: Source, rev: str, args) -> Iterator[dict]:
    for i, row in enumerate(_parquet_or_server_rows(src.repo_id, rev, args.max_rows)):
        yield {"row_index": i, "statement": row.get("formal_statement"), "header": None,
               "informal": row.get("informal_statement"), "name": row.get("problem_id"),
               "extra": {"orig_problem_id": row.get("problem_id"), "orig_source": row.get("source")}}


def _it_conjectures(src: Source, rev: str, args) -> Iterator[dict]:
    """Rows are declarations: statement = everything after the name (binders + type)."""
    for i, row in enumerate(_parquet_or_server_rows(src.repo_id, rev, args.max_rows)):
        kind = (row.get("type") or "").strip()
        imports = row.get("imports") or "[]"
        yield {"row_index": i, "statement": None, "header": None, "informal": row.get("docstring") or None,
               "name": row.get("symbolic_name"), "decl_kind": kind, "decl_tail": row.get("statement"),
               "imports": imports,
               "extra": {"library": row.get("library"), "filename": row.get("filename"),
                         "source_url": row.get("source_url"), "upstream_commit": row.get("commit")}}


SOURCES: dict[str, Source] = {s.key: s for s in [
    Source("formalmath-all", "SphereLab/FormalMATH-All", "mit", _it_formalmath),
    Source("numinamath-lean-sol", "iiis-lean/NuminaMath-LEAN-Sol", "apache-2.0", _it_numina),
    Source("nemotron-math-proofs-v1", "nvidia/Nemotron-Math-Proofs-v1", "cc-by-sa-4.0", _it_nemotron,
           note="sampled in Range stripes; the full file is 29.5 GB of trajectories"),
    Source("rl-goedel-pset-l45", "HangHor/RL-Goedel-Pset-Level-4-5", "mit", _it_rows,
           note="levels 4-5 proper (46,290 rows)"),
    Source("rl-goedel-pset-l35", "HangHor/RL-Goedel-Pset-Level-3-5", "mit", _it_rows,
           note="levels 3-5 superset of l45; not in the default set"),
    Source("lean4-formalconjectures", "phanerozoic/Lean4-FormalConjectures", "apache-2.0", _it_conjectures),
]}
DEFAULT_SOURCES = ["formalmath-all", "numinamath-lean-sol", "nemotron-math-proofs-v1", "rl-goedel-pset-l45",
                   "lean4-formalconjectures"]


def repo_revision(repo_id: str) -> str:
    with urllib.request.urlopen(f"https://huggingface.co/api/datasets/{repo_id}", timeout=60) as r:
        return json.loads(r.read())["sha"]


# --------------------------------------------------------------------------- pure conversion


def assemble_statement(rec: dict) -> tuple[str, dict]:
    """-> (candidate full source text, extra provenance). Raises ContractError('drop:<reason>')."""
    extra: dict[str, Any] = {}
    if rec.get("decl_kind") is not None:  # FormalConjectures declaration
        kind = rec["decl_kind"]
        if kind not in ("theorem", "lemma"):
            raise C.ContractError("drop:not_theorem_or_lemma")
        imps = str(rec.get("imports", "")).lower()
        if "formalconjectures" in imps or "formal_conjectures" in imps:
            raise C.ContractError("drop:needs_nonmathlib_imports")
        if not rec.get("name") or not rec.get("decl_tail"):
            raise C.ContractError("drop:empty_statement")
        raw = f"{kind} {rec['name']} {rec['decl_tail'].strip()} := by sorry"
        return C.GOEDEL_HEADER + raw, {"constructed_from_declaration": True}
    raw = rec.get("statement")
    if not isinstance(raw, str) or not raw.strip():
        raise C.ContractError("drop:empty_statement")
    raw = raw.strip("\n")
    if rec.get("header"):
        raw = rec["header"].rstrip() + "\n\n" + raw.lstrip()
    if _BY_SORRY_TAIL.search(raw):  # `:= sorry` term-mode placeholder
        raw = _BY_SORRY_TAIL.sub(":= by sorry", raw)
        extra["sorry_term_rewritten"] = True
    elif _BY_TAIL.search(raw):  # statement stored without the placeholder
        raw = raw.rstrip() + " sorry"
        extra["sorry_appended"] = True
    elif ":= by" not in raw and re.search(r"\bsorry\b", raw) is None:
        raw = raw.rstrip() + " := by sorry"  # statement stored without any body
        extra["body_appended"] = True
    return raw, extra


def _open_names(header: str) -> set[str]:
    out: set[str] = set()
    for ln in header.splitlines():
        ln = ln.strip()
        if ln.startswith("open "):
            out.update(ln[5:].split())
    return out


def to_problem(rec: dict, source: Source, revision: str, pid: str) -> tuple[dict, dict]:
    """Full conversion of one raw row -> (Problem record, side info). Raises ContractError('drop:<reason>')."""
    raw, extra = assemble_statement(rec)
    if len(raw) > MAX_STATEMENT_CHARS:
        raise C.ContractError("drop:too_long")
    try:
        stmt = C.normalize_formal_statement(raw)
    except C.ContractError:
        raise C.ContractError("drop:not_single_by_sorry")
    try:
        name, header, _sig = C.statement_key(stmt)
    except C.ContractError as e:
        why = re.sub(r"`[^`]*`", "`X`", str(e).replace("unsupported statement: ", ""))
        raise C.ContractError(f"drop:unsupported_shape: {why}"[:120])
    prov_header = {}
    if header.rstrip() != C.GOEDEL_HEADER.rstrip():
        parsed = C._strict().parse_statement(stmt)
        body = stmt[len(parsed.header_text):].lstrip("\n")
        stmt = C.GOEDEL_HEADER + body
        prov_header = {"header_rewritten": True, "original_header": header.strip()[:600]}
        extra_open = sorted(_open_names(header) - _open_names(C.GOEDEL_HEADER))
        if extra_open:  # the rewrite drops namespaces the statement may rely on: elaboration decides
            prov_header["header_extra_opens"] = extra_open
        try:  # the rewritten statement must still be a supported shape
            C.statement_key(stmt)
        except C.ContractError:
            raise C.ContractError("drop:unsupported_after_header_rewrite")
    prob = {
        "problem_id": pid, "source": source.problem_source, "split": "train", "formal_statement": stmt,
        "informal_prefix": rec.get("informal") if isinstance(rec.get("informal"), str) and rec["informal"].strip() else None,
        "statement_sha": C.statement_sha(stmt),
        "provenance": {"source_id": source.repo_id, "revision": revision, "row_index": rec["row_index"],
                       "license": source.license, "theorem_name": name, **extra, **prov_header, **rec.get("extra", {})},
        "ancestry": {"root": pid, "parents": []},
    }
    C.validate_problem(prob)
    return prob, {"header_rewritten": bool(prov_header)}


def ingest_source(src: Source, revision: str, out_dir: Path, args) -> dict:
    counts: Counter = Counter()
    drop_examples: dict[str, list] = {}
    seen: set[str] = set()
    t0 = time.time()
    out_path = out_dir / f"{src.key}.jsonl"
    tmp = out_path.with_name(out_path.name + f".tmp{os.getpid()}")
    try:
        with open(tmp, "w", encoding="utf-8") as f:
            for rec in src.iterate(src, revision, args):
                counts["rows_seen"] += 1
                if rec.get("english") is False:
                    counts["drop:not_english"] += 1
                    continue
                pid = f"{src.key}-{rec['row_index']}"
                try:
                    prob, side = to_problem(rec, src, revision, pid)
                except C.ContractError as e:
                    why = str(e)
                    if not why.startswith("drop:"):
                        raise
                    counts[why] += 1
                    ex = drop_examples.setdefault(why, [])
                    if len(ex) < 3:
                        ex.append({"row_index": rec["row_index"], "name": rec.get("name")})
                    continue
                if prob["statement_sha"] in seen:
                    counts["drop:duplicate_statement"] += 1
                    continue
                seen.add(prob["statement_sha"])
                f.write(json.dumps(prob, ensure_ascii=False) + "\n")
                counts["kept"] += 1
                counts["header_rewritten"] += int(side["header_rewritten"])
                counts["header_extra_opens"] += int("header_extra_opens" in prob["provenance"])
                if counts["rows_seen"] % 20000 == 0:
                    log.info("%s: %d rows, %d kept", src.key, counts["rows_seen"], counts["kept"])
            f.flush()
            os.fsync(f.fileno())
        os.replace(tmp, out_path)
    finally:
        if tmp.exists():
            tmp.unlink()
    (out_dir / f"{src.key}.drops.sample.json").write_text(json.dumps(drop_examples, indent=1, ensure_ascii=False))
    return {"source_id": src.repo_id, "revision": revision, "license": src.license, "note": src.note,
            "counts": dict(counts), "seconds": round(time.time() - t0, 1), "output": str(out_path),
            "output_sha256": C.sha256_file(out_path)}


def main(argv=None) -> int:
    ap = argparse.ArgumentParser(description=__doc__.split("\n")[0])
    ap.add_argument("--out-dir", default=str(DEFAULT_OUT))
    ap.add_argument("--sources", nargs="+", default=DEFAULT_SOURCES, choices=sorted(SOURCES))
    ap.add_argument("--workers", type=int, default=2, help="parallel sources; hard-capped at 2")
    ap.add_argument("--max-rows", type=int, default=0, help="cap rows per datasets-server source (0 = all)")
    ap.add_argument("--nemotron-rows", type=int, default=60000)
    ap.add_argument("--nemotron-stripes", type=int, default=8)
    a = ap.parse_args(argv)
    logging.basicConfig(level=logging.INFO, format="%(asctime)s %(name)s %(message)s")
    out_dir = Path(a.out_dir)
    out_dir.mkdir(parents=True, exist_ok=True)
    revs = {k: repo_revision(SOURCES[k].repo_id) for k in a.sources}
    log.info("revisions: %s", revs)
    workers = max(1, min(2, a.workers))
    results: dict[str, Any] = {}
    with ThreadPoolExecutor(max_workers=workers) as ex:
        futs = {k: ex.submit(ingest_source, SOURCES[k], revs[k], out_dir, a) for k in a.sources}
        for k, fu in futs.items():
            try:
                results[k] = fu.result()
                log.info("%s: %s", k, results[k]["counts"])
            except Exception as e:  # one failed source must not discard the others' results
                log.error("%s FAILED: %r", k, e)
                results[k] = {"source_id": SOURCES[k].repo_id, "revision": revs[k], "error": repr(e)}
    man = out_dir / "ingest_manifest.json"
    prev = json.loads(man.read_text()) if man.exists() else {}
    prev.update(results)
    man.write_text(json.dumps(prev, indent=1, ensure_ascii=False))
    return 0


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