#!/usr/bin/env python
"""Build every corpus this project trains on, in the only legal order.

    python scripts/build_data.py --out <build root> [--tokenizer <dir>] [--curation-target N]

Stages: eval sets -> decontamination hash set -> SFT -> Arm A -> Arm B.
Each refuses to run without its upstream receipt; the final manifest
(`<out>/manifest.json`) pins every content hash, every measured count against
its pinned value, reject and contamination counts, the length-bias baseline,
the delimiter ids, the tokenizer hash, and the dataset revisions. Sections
11/12 quote that manifest in their launch banner and refuse to start without
it.

Host safety (this box, 186GB, NO swap, shared): `apply_resource_guards(
heavy=True)` runs the preflight and takes the exclusive job flock BEFORE
anything heavy, so this build cannot run concurrently with a training job --
that exact overlap killed the predecessor's first SFT run. Every corpus is
streamed; nothing is held resident.
"""

from __future__ import annotations

import argparse
import json
import logging
import sys
from pathlib import Path

sys.path.insert(0, str(Path(__file__).resolve().parent.parent))

from align.utils.guard import apply_resource_guards  # noqa: E402

logging.basicConfig(level=logging.INFO, format="%(asctime)s %(levelname)s %(message)s")
logger = logging.getLogger("build_data")

DEFAULT_OUT = "/apps/development/rlhf/nano-mistral-o1/data"
TOKENIZER_SNAPSHOT = (
    "/apps/voxtral/.cache/huggingface/hub/models--mistralai--Mistral-7B-v0.3/"
    "snapshots/caa1feb0e54d415e2df31207e5f4e273e33509b1"
)
#: Byte-identical tokenizer files vendored in-repo (assets/tokenizer/.../PROVENANCE.md).
#: The HF-cache snapshot is NOT on every box; the build must not die on a missing
#: cache when the same four files are in the repo.
VENDORED_TOKENIZER = str(
    Path(__file__).resolve().parent.parent / "assets" / "tokenizer" / "mistral-7b-v0.3"
)


def default_tokenizer_dir() -> str:
    """The HF snapshot if it holds a tokenizer.json, else the vendored copy."""
    for c in (TOKENIZER_SNAPSHOT, VENDORED_TOKENIZER):
        if (Path(c) / "tokenizer.json").is_file():
            return c
    return TOKENIZER_SNAPSHOT  # fails loudly in load_tokenizer, naming the path


def main() -> int:
    ap = argparse.ArgumentParser()
    ap.add_argument(
        "--out", default=DEFAULT_OUT, help="build root; every path under it is asserted onto /apps"
    )
    ap.add_argument("--tokenizer", default=default_tokenizer_dir())
    ap.add_argument(
        "--curation-target",
        type=int,
        default=None,
        help="Arm B curated pair count (default: the pinned 12,000)",
    )
    ap.add_argument(
        "--max-len",
        type=int,
        default=2048,
        help="build cutoff, deliberately above the 1024 training "
        "seq_len: see PinnedYields.max_len for the measured "
        "length distribution behind this choice",
    )
    ap.add_argument("--held-out-frac", type=float, default=0.05)
    ap.add_argument("--seed", type=int, default=0)
    ap.add_argument(
        "--strict-pins",
        default=True,
        action=argparse.BooleanOptionalAction,
        help="assert every measured count against its pinned "
        "value. --no-strict-pins is RE-MEASUREMENT MODE: it "
        "stamps pins_verified=false into the manifest and "
        "sections 11/12 refuse to launch against it.",
    )
    ap.add_argument(
        "--strict-host-check",
        default=True,
        action=argparse.BooleanOptionalAction,
        help="disable ONLY with a deliberate reason; the check is "
        "the guard that did not exist on 2026-08-12",
    )
    args = ap.parse_args()

    apply_resource_guards(
        cpu_threads=8, heavy=True, label="build_data", strict_host_check=args.strict_host_check
    )

    from align.data.build import PinnedYields, build_all
    from align.data.sources import DataSources
    from align.tokenizer import load_tokenizer

    tokenizer = load_tokenizer(args.tokenizer)
    sources = DataSources.from_local_hub()

    kwargs = {
        "max_len": args.max_len,
        "held_out_frac": args.held_out_frac,
        "seed": args.seed,
        "strict": args.strict_pins,
    }
    if args.curation_target is not None:
        kwargs["arm_b_curated"] = args.curation_target
    pins = PinnedYields(**kwargs)

    manifest = build_all(sources, args.out, tokenizer, pins=pins)
    print(json.dumps(manifest, indent=1))
    logger.info("build complete: %s", Path(args.out) / "manifest.json")
    return 0


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