"""Shared helpers for the P2 data / verification / evaluation scripts (stdlib + numpy).

Loaded BY PATH from the sibling scripts (`scripts/` is not a package and the dependency
firewall rejects an import it cannot place in stdlib / allowed / align), the same way
`scripts/lean_ceiling_test.py` loads its case tables.

Contents: the three JSONL record schemas of the P2 contract with strict validators, an
append-only JSONL writer that survives a SIGKILL mid-line, and the statement / proof
normalisations used for decontamination and for "distinct proof" selection.
"""

from __future__ import annotations

import hashlib
import importlib.util
import json
import os
import re
import threading
from pathlib import Path
from typing import Any, Iterable, Iterator, Optional

REPO = Path(__file__).resolve().parents[2]

#: The header Goedel's DeepSeekCoTHandler.extrac_code prepends to every extracted proof
#: (reference-repos/Goedel-Prover-V2/src/utils.py). A pool statement with any other header
#: can never be matched by a generated proof (strict verifier: header token-equality), so it
#: is unwinnable and is dropped at build time instead of becoming a silent 0.
GOEDEL_HEADER = "import Mathlib\nimport Aesop\n\nset_option maxHeartbeats 0\n\nopen BigOperators Real Nat Topology Rat\n\n"

PROBLEM_FIELDS = ("problem_id", "source", "split", "formal_statement", "informal_prefix", "statement_sha")
GENERATION_FIELDS = (
    "problem_id", "sample_idx", "round", "cap_tokens", "prompt_ids", "completion_ids", "prompt_text",
    "completion_text", "finish_reason", "lean_code", "messages", "model_path", "model_rev", "sampling",
    "parent",
)
KEY_FIELDS = ("problem_id", "sample_idx", "round", "cap_tokens")


class ContractError(ValueError):
    """A record violates the P2 shared contract. Always loud: a lenient reader is how a
    cap-mislabelled or stale-verdict row ends up in a table."""


# --------------------------------------------------------------------------- hashing


def sha256_text(s: str) -> str:
    return hashlib.sha256(s.encode("utf-8")).hexdigest()


def sha256_file(path: str | os.PathLike, chunk: int = 4 << 20) -> str:
    h = hashlib.sha256()
    with open(path, "rb") as f:
        while True:
            b = f.read(chunk)
            if not b:
                break
            h.update(b)
    return h.hexdigest()


# --------------------------------------------------------------------------- jsonl io


def iter_jsonl(path: str | os.PathLike) -> Iterator[dict]:
    """Strict JSONL reader: a malformed line raises with its line number (no skipping)."""
    with open(path, encoding="utf-8") as f:
        for i, line in enumerate(f, 1):
            if not line.strip():
                continue
            try:
                obj = json.loads(line)
            except json.JSONDecodeError as e:
                raise ContractError(f"{path}:{i}: malformed JSON ({e.msg})") from e
            if not isinstance(obj, dict):
                raise ContractError(f"{path}:{i}: expected a JSON object, got {type(obj).__name__}")
            yield obj


def write_jsonl_atomic(path: str | os.PathLike, rows: Iterable[dict]) -> int:
    """tmp + fsync + rename. A half-written problem file looks complete; never leave one."""
    p = Path(path)
    p.parent.mkdir(parents=True, exist_ok=True)
    tmp = p.with_name(p.name + f".tmp{os.getpid()}")
    n = 0
    try:
        with open(tmp, "w", encoding="utf-8") as f:
            for r in rows:
                f.write(json.dumps(r, ensure_ascii=False) + "\n")
                n += 1
            f.flush()
            os.fsync(f.fileno())
        os.replace(tmp, p)
    finally:
        if tmp.exists():
            tmp.unlink()
    return n


def repair_torn_tail(path: str | os.PathLike) -> int:
    """Make an append-only JSONL safe to resume: if the last line was cut off by a kill, drop
    it. Returns the number of bytes removed. A MIDDLE line that does not parse is not a torn
    write (appends only ever damage the tail), so that raises when the file is read."""
    p = Path(path)
    if not p.exists() or p.stat().st_size == 0:
        return 0
    with open(p, "rb") as f:
        data = f.read()
    if data.endswith(b"\n"):
        return 0
    cut = data.rfind(b"\n") + 1  # 0 when the file has a single unterminated line
    tail = data[cut:]
    try:
        json.loads(tail.decode("utf-8"))
        with open(p, "ab") as f:  # complete object that lost only its newline
            f.write(b"\n")
        return 0
    except (json.JSONDecodeError, UnicodeDecodeError):
        with open(p, "r+b") as f:
            f.truncate(cut)
            f.flush()
            os.fsync(f.fileno())
        return len(tail)


class AppendJsonl:
    """Thread-safe append-only JSONL writer; every `append` is flushed and fsynced so a
    SIGKILL loses at most the record being written (which `repair_torn_tail` discards)."""

    def __init__(self, path: str | os.PathLike) -> None:
        self.path = Path(path)
        self.path.parent.mkdir(parents=True, exist_ok=True)
        self.torn_bytes_removed = repair_torn_tail(self.path)
        self._f = open(self.path, "ab")
        self._lock = threading.Lock()

    def append(self, rec: dict) -> None:
        line = (json.dumps(rec, ensure_ascii=False) + "\n").encode("utf-8")
        with self._lock:
            self._f.write(line)
            self._f.flush()
            os.fsync(self._f.fileno())

    def close(self) -> None:
        self._f.close()

    def __enter__(self) -> "AppendJsonl":
        return self

    def __exit__(self, *exc) -> None:
        self.close()


# --------------------------------------------------------------------------- contract


def gen_key(rec: dict) -> tuple:
    return (rec["problem_id"], int(rec["sample_idx"]), int(rec["round"]), int(rec["cap_tokens"]))


def validate_problem(rec: dict) -> dict:
    missing = [k for k in PROBLEM_FIELDS if k not in rec]
    if missing:
        raise ContractError(f"problem {rec.get('problem_id')!r} lacks fields {missing}")
    for k in ("problem_id", "source", "split", "formal_statement", "statement_sha"):
        if not isinstance(rec[k], str) or not rec[k].strip():
            raise ContractError(f"problem field {k!r} must be a non-empty string, got {rec[k]!r}")
    if rec["informal_prefix"] is not None and not isinstance(rec["informal_prefix"], str):
        raise ContractError("informal_prefix must be a string or null")
    return rec


def effective_max_tokens(rec: dict) -> int:
    """What the sampler was actually allowed to emit. `cap_tokens` is the NOMINAL cap every
    number is labelled with; `sampling.max_tokens` is smaller only when the context window,
    not the cap, ended the sample (`context_clipped`): at cap 40960 with a 40960-token window
    every sample loses its prompt length (~250-600 tokens)."""
    mt = (rec.get("sampling") or {}).get("max_tokens")
    return int(mt) if mt is not None else int(rec["cap_tokens"])


def validate_generation(rec: dict) -> dict:
    """Shape checks on one Generation record. Cross-record checks live in the callers."""
    missing = [k for k in GENERATION_FIELDS if k not in rec]
    pid = rec.get("problem_id")
    if missing:
        raise ContractError(f"generation {pid!r} lacks fields {missing}")
    for k in ("sample_idx", "round", "cap_tokens"):
        if not isinstance(rec[k], int) or isinstance(rec[k], bool) or rec[k] < 0:
            raise ContractError(f"generation {pid!r}: {k} must be a non-negative int, got {rec[k]!r}")
    if rec["cap_tokens"] < 1:
        raise ContractError(f"generation {pid!r}: cap_tokens must be >= 1")
    if rec["finish_reason"] not in ("stop", "length"):
        raise ContractError(f"generation {pid!r}: finish_reason {rec['finish_reason']!r} not in stop|length")
    for k in ("prompt_ids", "completion_ids"):
        v = rec[k]
        if not isinstance(v, list) or not v or not all(isinstance(t, int) and not isinstance(t, bool) for t in v):
            raise ContractError(f"generation {pid!r}: {k} must be a non-empty list of ints")
    eff = effective_max_tokens(rec)
    if eff > rec["cap_tokens"]:
        raise ContractError(
            f"generation {pid!r}: sampling.max_tokens={eff} exceeds the stated cap_tokens={rec['cap_tokens']} -- the "
            "cap label understates what the sampler was allowed, and every number must state its true cap"
        )
    if len(rec["completion_ids"]) > eff:
        raise ContractError(
            f"generation {pid!r}: {len(rec['completion_ids'])} completion tokens exceed the sampler's max_tokens={eff} "
            f"(cap_tokens={rec['cap_tokens']}) -- the cap label is wrong, and every number must state its true cap"
        )
    if rec["finish_reason"] == "length" and len(rec["completion_ids"]) != eff:
        raise ContractError(
            f"generation {pid!r}: finish_reason=length but {len(rec['completion_ids'])} tokens != max_tokens {eff} "
            "(the run was cut by a different limit than the one it is labelled with)"
        )
    if rec["lean_code"] is not None and not isinstance(rec["lean_code"], str):
        raise ContractError(f"generation {pid!r}: lean_code must be a string or null")
    parent = rec["parent"]
    if rec["round"] == 0 and parent is not None:
        raise ContractError(f"generation {pid!r}: round 0 cannot have a parent")
    if rec["round"] > 0:
        if not (isinstance(parent, dict) and {"problem_id", "sample_idx", "round"} <= set(parent)):
            raise ContractError(f"generation {pid!r}: round {rec['round']} needs a parent {{problem_id,sample_idx,round}}")
        if parent["problem_id"] != pid or parent["round"] != rec["round"] - 1:
            raise ContractError(f"generation {pid!r}: parent {parent} is not the previous round of the same problem")
    return rec


def load_problems(path: str | os.PathLike) -> dict[str, dict]:
    out: dict[str, dict] = {}
    for rec in iter_jsonl(path):
        validate_problem(rec)
        if rec["problem_id"] in out:
            raise ContractError(f"{path}: duplicate problem_id {rec['problem_id']!r}")
        out[rec["problem_id"]] = rec
    if not out:
        raise ContractError(f"{path}: no problems")
    return out


def code_sha(lean_code: Optional[str]) -> Optional[str]:
    return None if lean_code is None else sha256_text(lean_code)


# --------------------------------------------------------------------------- statements


_BY = re.compile(r":=\s*by\b")


def normalize_formal_statement(raw: str) -> str:
    """Goedel's statement normalisation (`lean4_code.split(":= by")[0] + ":= by sorry"`), made
    safe. Goedel's split silently TRUNCATES a statement that contains `:= by` earlier (for
    example a `(h : x = by norm_num)` binder) and so changes the problem; here that raises. A
    line break between `:=` and `by` (`:=\\n  by sorry`) is accepted and canonicalised, which
    Goedel's literal split cannot handle at all."""
    s = raw.strip("\n")
    hits = list(_BY.finditer(s))
    if len(hits) != 1:
        raise ContractError(f"statement must contain exactly one ':= by', found {len(hits)}")
    m = hits[0]
    tail = s[m.end() :]
    if tail.strip() not in ("", "sorry"):
        raise ContractError(f"text after ':= by' is not a bare sorry: {tail.strip()[:60]!r}")
    return s[: m.start()] + ":= by sorry"


def _strict():
    from align.verify.lean import strict  # lazy: keeps `import p2_common` cheap

    return strict


def statement_key(statement: str) -> tuple[str, str, tuple[str, ...]]:
    """(theorem name, header text, signature tokens) of a formal statement, via the SAME parser
    the verifier uses. Raises ContractError when the verifier could not use the statement."""
    p = _strict().parse_statement(statement)
    if not p.ok:
        raise ContractError(f"unsupported statement: {p.problem}")
    return p.name, p.header_text, tuple(p.sig_toks)


def statement_sha(statement: str) -> str:
    """sha256 of the whitespace- and comment-normalised signature (lexer tokens joined by one
    space). The theorem NAME is deliberately excluded: a renamed copy is the same problem."""
    _, _, sig = statement_key(statement)
    return sha256_text(" ".join(sig))


_KEEP_LOWER = frozenset(
    "fun forall exists if then else let have by sorry at with from in do match return where "
    "sin cos tan exp log sqrt abs max min id fst snd succ pred gcd lcm choose fact floor ceil "
    "card sum prod range filter map deriv logb arctan sinh cosh cexp nnabs dist norm "
    "theorem lemma".split()
)


def alpha_sha(statement: str) -> str:
    """Like `statement_sha` but invariant to consistent renaming of local (lowercase, undotted)
    identifiers. Deliberately conservative: it also equates `sin x = 1` with `cos y = 1`; a
    false positive here costs one pool statement, a false negative costs a contaminated
    evaluation."""
    _, _, sig = statement_key(statement)
    mp: dict[str, str] = {}
    out = []
    for t in sig:
        if re.fullmatch(r"[a-z_][A-Za-z0-9_'₀-₉]*", t) and t not in _KEEP_LOWER:
            out.append(mp.setdefault(t, f"v{len(mp)}"))
        else:
            out.append(t)
    return sha256_text(" ".join(out))


_LEAN_PROMPTS = None


def _lean_prompts():
    """`extras/vllm_rollout/lean_prompts.py` (the byte-faithful port of Goedel's extraction), loaded
    BY PATH: it is stdlib-only at import and is the single source of truth for what the model's
    code block means, so the SFT selection must not grow a second, slightly different regex."""
    global _LEAN_PROMPTS
    if _LEAN_PROMPTS is None:
        spec = importlib.util.spec_from_file_location("p2_lean_prompts", REPO / "extras" / "vllm_rollout" / "lean_prompts.py")
        mod = importlib.util.module_from_spec(spec)
        spec.loader.exec_module(mod)
        _LEAN_PROMPTS = mod
    return _LEAN_PROMPTS


def displayed_signature(completion_text: str) -> Optional[tuple[str, ...]]:
    """Signature tokens of the theorem the MODEL WROTE in the code block Goedel's extractor takes
    (last ```lean4 fence, first `theorem ... := by`), i.e. BEFORE `problem_check` swaps the original
    statement in. None when there is no such theorem or the verifier's parser cannot read it.

    Why this exists: `problem_check` replaces whatever statement the model wrote with the original,
    so a sample whose displayed theorem is `v = 66` but whose tactics happen to prove the original
    `v = 65` is ACCEPTED, and its completion then teaches the model to state the wrong theorem.
    Goedel's evaluation protocol credits it (and so does ours, to stay comparable); a TRAINING row
    must not."""
    lp = _lean_prompts()
    extracted = lp.DeepSeekCoTHandler().extrac_code(completion_text)
    if extracted == "None":
        return None
    proof = lp.remove_comments(extracted)
    span = lp.return_theorem_to_replace(proof)
    if span is None:
        return None
    shown = re.sub(r":=\s*by$", ":= by sorry", proof[span[0] : span[1]].strip())
    parsed = _strict().parse_statement(shown)
    return tuple(parsed.sig_toks) if parsed.ok else None


def displayed_statement_matches(formal_statement: str, completion_text: str) -> bool:
    """True iff the theorem the model displayed has the same signature tokens as `formal_statement`
    (theorem NAME ignored: `problem_check` renames it anyway and a renamed copy is the same problem)."""
    shown = displayed_signature(completion_text)
    return shown is not None and shown == statement_key(formal_statement)[2]


def norm_name(name: str) -> str:
    return re.sub(r"[^a-z0-9]", "", name.lower())


def informal_tokens(text: str) -> list[str]:
    return re.findall(r"[a-z0-9]+", text.lower())


def ngrams(tokens: list[str], n: int) -> set[tuple[str, ...]]:
    return {tuple(tokens[i : i + n]) for i in range(len(tokens) - n + 1)} if len(tokens) >= n else set()


def proof_hash(lean_code: str) -> str:
    """Hash of a proof invariant to whitespace and comments (lexer tokens), used to prefer
    DISTINCT proofs. Falls back to whitespace-collapsed text when the lexer gives up."""
    lx = _strict().lex(lean_code)
    if lx.broken:
        return sha256_text(" ".join(lean_code.split()))
    return sha256_text(" ".join(t.text for t in lx.toks))


def compile_view(full_code: str) -> str:
    """Goedel's `compile.py:handle`: drop every line whose stripped text starts with import /
    set_option / open. Goedel compiles THIS text in a REPL environment that already ran the
    header, so the line numbers of its error dicts index this text, and `get_error_str` must be
    given this text (not `full_code`). Identical to `extras/vllm_rollout/lean_prompts.compile_view`
    (a test pins the two together)."""
    lines = full_code.split("\n")
    return "\n".join(l for l in lines if not (l.strip().startswith("import") or l.strip().startswith("set_option") or l.strip().startswith("open")))