#!/usr/bin/env python3
"""Reproduce the frozen bibliography and claim samples; standard library only.
Usage: python reproduce_math_sample.py sampled-math-papers.json > sample-manifest.json
This script performs no network requests and makes no verification judgments.
"""
import hashlib
import json
import re
import sys
from pathlib import Path

SEED = 20261008
EXPECTED_SHA256 = "ec9c84e53baaff4cb83c9a33d85380bf2d95bc3b8ad5b467e89b19212127b8ba"

def balanced_end(text, start):
    depth = 0
    for i in range(start, len(text)):
        if text[i] == "{" and (i == 0 or text[i-1] != "\\"):
            depth += 1
        elif text[i] == "}" and (i == 0 or text[i-1] != "\\"):
            depth -= 1
            if depth == 0:
                return i + 1
    raise ValueError("Unbalanced bibliography entry")

def rank(entry, tag="bibliography"):
    value = f"{SEED}|{tag}|{entry['family_id']}|{entry['key']}"
    return hashlib.sha256(value.encode()).hexdigest()

def build_sample(filename):
    raw = Path(filename).read_bytes()
    digest = hashlib.sha256(raw).hexdigest()
    if digest != EXPECTED_SHA256:
        raise ValueError(f"Unexpected input checksum: {digest}")
    data = json.loads(raw)
    assert len(data["sample"]) == 20
    assert all(hashlib.sha256(r["content"].encode()).hexdigest() == r["sha256"]
               for r in data["records"])
    entries = []
    excluded = 0
    for paper in data["sample"]:
        root = str(Path(paper["representative_pdf"]).parent)
        records = [r for r in data["records"] if r["path"].startswith(root + "/")]
        contexts = {}
        for record in records:
            if not record["path"].endswith(".tex"):
                continue
            text = re.sub(r"(?m)(?<!\\)%[^\n]*", "", record["content"])
            pattern = r"\\(?:cite[a-zA-Z]*|nocite)\*?(?:\s*\[[^\]]*\]){0,2}\s*\{([^}]+)\}"
            for match in re.finditer(pattern, text):
                for key in match.group(1).split(","):
                    contexts.setdefault(key.strip(), []).append({
                        "path": record["path"],
                        "offset": match.start(),
                        "command": match.group(0),
                        "context": text[max(0, match.start()-650):match.end()+650],
                    })
        paper_entries = []
        for record in records:
            text = record["content"]
            if record["path"].endswith(".bib"):
                for match in re.finditer(r"(?m)^@(\w+)\s*\{\s*([^,\s]+)\s*,", text):
                    end = balanced_end(text, text.find("{", match.start()))
                    key = match.group(2)
                    if key not in contexts and "*" not in contexts:
                        excluded += 1
                        continue
                    paper_entries.append({"key": key,
                                          "raw": text[match.start():end],
                                          "path": record["path"]})
            elif "\\bibitem" in text:
                matches = list(re.finditer(r"\\bibitem(?:\s*\[[\s\S]*?\])?\s*\{([^}]+)\}", text))
                for j, match in enumerate(matches):
                    if j+1 < len(matches):
                        end = matches[j+1].start()
                    elif "\\end{thebibliography}" in text[match.end():]:
                        end = text.find("\\end{thebibliography}", match.end())
                    else:
                        end = len(text)
                    paper_entries.append({"key": match.group(1),
                                          "raw": text[match.start():end].strip(),
                                          "path": record["path"]})
        assert len({e["key"] for e in paper_entries}) == len(paper_entries)
        for entry in paper_entries:
            entry.update(family_id=paper["family_id"],
                         contexts=contexts.get(entry["key"], []),
                         paper=root.split("/")[-1])
        entries.extend(paper_entries)
    families = sorted({e["family_id"] for e in entries})
    selected = []
    counts = {}
    for family in families:
        frame = [e for e in entries if e["family_id"] == family]
        assert len(frame) >= 5
        counts[family] = len(frame)
        selected.extend(sorted(frame, key=rank)[:5])
    for i, entry in enumerate(selected, 1):
        entry["sample_id"] = f"R{i:03d}"
        entry["selection_hash"] = rank(entry)
    claim_families = sorted(
        families,
        key=lambda family: hashlib.sha256(
            f"{SEED}|claim-family|{family}".encode()).hexdigest()
    )[:10]
    claims = [
        min([e for e in selected if e["family_id"] == family and e["contexts"]],
            key=lambda e: rank(e, "claim"))
        for family in claim_families
    ]
    claim_ids = [e["sample_id"] for e in claims]
    for entry in selected:
        entry["claim_selected"] = entry["sample_id"] in claim_ids
    assert len(entries) == 409 and excluded == 8
    assert len(selected) == 100 and len(claim_ids) == 10
    return {
        "input_sha256": digest,
        "repository": data["repository"],
        "commit": data["commit"],
        "seed": SEED,
        "frame_count": len(entries),
        "excluded_uncited_bib_entries": excluded,
        "frame_counts_by_family": counts,
        "sampling": "Five lowest SHA-256 ranks per family; no replacement within a family.",
        "claim_selection": "Ten lowest claim-family hashes, then lowest claim hash in each family. First context in supplied source-record order.",
        "context_offset_basis": "Unicode character offsets after LaTeX comment removal, not PDF page numbers.",
        "scope_limit": "Source-reconstructed bibliography; does not compile PDFs or independently validate repository provenance.",
        "claim_sample_ids": claim_ids,
        "sample": selected,
    }

if __name__ == "__main__":
    if len(sys.argv) != 2:
        raise SystemExit("Usage: python reproduce_math_sample.py sampled-math-papers.json")
    print(json.dumps(build_sample(sys.argv[1]), ensure_ascii=False, indent=2))
