#!/usr/bin/env python3
"""Build deterministic, document-isolated OC-LM Strict-Small training epochs."""

from __future__ import annotations

import argparse
import hashlib
import json
from pathlib import Path

import numpy as np
from tokenizers import Tokenizer


EXPECTED_WORDS = {
    "bnc_spoken.train.txt": 762_073,
    "childes.train.txt": 2_841_101,
    "gutenberg.train.txt": 2_557_721,
    "open_subtitles.train.txt": 2_282_877,
    "simple_wiki.train.txt": 1_531_437,
    "switchboard.train.txt": 24_791,
}
LENGTHS = (128,) * 6 + (256,) * 2 + (512,) * 2
SPECIAL = {"unk": 0, "bos": 1, "eos": 2, "pad": 3, "mask": 4}


def sha256(path: Path) -> str:
    digest = hashlib.sha256()
    with path.open("rb") as handle:
        for chunk in iter(lambda: handle.read(8 << 20), b""):
            digest.update(chunk)
    return digest.hexdigest()


def flush_row(handles, tokens, segments, words, objective, length):
    if not tokens:
        return 0
    used = len(tokens)
    tokens.extend([SPECIAL["pad"]] * (length - used))
    segments.extend([0] * (length - used))
    np.asarray(tokens, dtype="<u2").tofile(handles["tokens"])
    np.asarray(segments, dtype="<u2").tofile(handles["segments"])
    np.asarray([used], dtype="<u2").tofile(handles["lengths"])
    np.asarray([words], dtype="<u4").tofile(handles["words"])
    np.asarray([objective], dtype="u1").tofile(handles["objectives"])
    return 1


def main() -> None:
    parser = argparse.ArgumentParser()
    parser.add_argument("--raw", type=Path, required=True)
    parser.add_argument("--tokenizer", type=Path, required=True)
    parser.add_argument("--output", type=Path, required=True)
    parser.add_argument("--seed", type=int, default=20260904)
    parser.add_argument("--batch-lines", type=int, default=4096)
    parser.add_argument("--schedule", choices=["cycle", "fixed512"], required=True)
    args = parser.parse_args()
    args.output.mkdir(parents=True, exist_ok=False)
    schedule = (128,256,512,128,256,512,128,256,512,512) if args.schedule == "cycle" else (512,)*10

    tokenizer = Tokenizer.from_file(str(args.tokenizer / "tokenizer.json"))
    if sha256(args.tokenizer / "tokenizer.json") != "98dfab9eabdd78aed27c025ec1bdbd881172c6339c0b491f9e1612e9a36fcbf8":
        raise RuntimeError("shared DST tokenizer hash mismatch")
    observed_special = {
        "unk": tokenizer.token_to_id("<unk>"),
        "bos": tokenizer.token_to_id("<s>"),
        "eos": tokenizer.token_to_id("</s>"),
        "pad": tokenizer.token_to_id("<pad>"),
        "mask": tokenizer.token_to_id("<mask>"),
    }
    if observed_special != SPECIAL or tokenizer.get_vocab_size() != 16_384:
        raise RuntimeError(f"tokenizer contract mismatch: {observed_special}")

    documents: list[list[int]] = []
    word_counts: list[int] = []
    raw_manifest = {}
    for path in sorted(args.raw.glob("*.txt")):
        expected = EXPECTED_WORDS.get(path.name)
        if expected is None:
            raise RuntimeError(f"unexpected raw file: {path.name}")
        file_words = 0
        batch_text: list[str] = []
        batch_words: list[int] = []

        def emit_batch() -> None:
            if not batch_text:
                return
            encoded = [item.ids for item in tokenizer.encode_batch(batch_text, add_special_tokens=False)]
            for ids, count in zip(encoded, batch_words, strict=True):
                documents.append([SPECIAL["bos"], *ids, SPECIAL["eos"]])
                word_counts.append(count)
            batch_text.clear()
            batch_words.clear()

        with path.open("r", encoding="utf-8") as handle:
            for raw_line in handle:
                text = raw_line.rstrip("\r\n")
                count = len(text.split())
                if count == 0:
                    continue
                file_words += count
                batch_text.append(text)
                batch_words.append(count)
                if len(batch_text) >= args.batch_lines:
                    emit_batch()
        emit_batch()
        if file_words != expected:
            raise RuntimeError(f"word mismatch for {path.name}: {file_words} != {expected}")
        raw_manifest[path.name] = {
            "bytes": path.stat().st_size,
            "sha256": sha256(path),
            "words": file_words,
        }

    if sum(word_counts) != 10_000_000:
        raise RuntimeError("Strict-Small corpus must contain exactly 10,000,000 words")

    docs_token_count = sum(map(len, documents))
    epochs = []
    for epoch, length in enumerate(schedule):
        epoch_dir = args.output / f"epoch_{epoch:02d}_len_{length}"
        epoch_dir.mkdir(parents=True, exist_ok=True)
        paths = {
            name: epoch_dir / f"{name}.bin"
            for name in ("tokens", "segments", "lengths", "words", "objectives")
        }
        handles = {name: path.open("wb") for name, path in paths.items()}
        rng = np.random.default_rng(args.seed + epoch)
        order = rng.permutation(len(documents))
        boundary_rng = np.random.default_rng(args.seed + 100000 + epoch)
        expected_stream = hashlib.sha256()
        row_tokens: list[int] = []
        row_segments: list[int] = []
        row_words = 0
        segment = 0
        rows = 0
        for doc_index in order:
            doc = documents[int(doc_index)]
            expected_stream.update(np.asarray(doc, dtype='<u2').tobytes())
            # Long documents get a new first-chunk boundary each epoch.
            # Retain the prefix and tail; never rotate tokens or join future to past.
            first_size = length - int(boundary_rng.integers(0, length)) if len(doc) > length else None
            if first_size is not None and row_tokens:
                rows += flush_row(handles, row_tokens, row_segments, row_words, (rows + epoch) & 1, length)
                row_tokens, row_segments, row_words, segment = [], [], 0, 0
            cursor = 0
            while cursor < len(doc):
                room = length - len(row_tokens)
                if cursor == 0 and first_size is not None:
                    room = min(room, first_size)
                take = min(room, len(doc) - cursor)
                row_tokens.extend(doc[cursor : cursor + take])
                row_segments.extend([segment] * take)
                cursor += take
                if cursor == len(doc):
                    row_words += word_counts[int(doc_index)]
                    segment += 1
                if len(row_tokens) == length or (first_size is not None and cursor == first_size):
                    rows += flush_row(
                        handles,
                        row_tokens,
                        row_segments,
                        row_words,
                        (rows + epoch) & 1,
                        length,
                    )
                    row_tokens, row_segments, row_words, segment = [], [], 0, 0
        if row_tokens:
            rows += flush_row(
                handles,
                row_tokens,
                row_segments,
                row_words,
                (rows + epoch) & 1,
                length,
            )
        for handle in handles.values():
            handle.flush()
            handle.close()

        words = np.memmap(paths["words"], mode="r", dtype="<u4", shape=(rows,))
        lengths = np.memmap(paths["lengths"], mode="r", dtype="<u2", shape=(rows,))
        objectives = np.memmap(paths["objectives"], mode="r", dtype="u1", shape=(rows,))
        if int(words.sum(dtype=np.uint64)) != 10_000_000:
            raise RuntimeError(f"epoch {epoch} word total mismatch")
        if int(lengths.sum(dtype=np.uint64)) != docs_token_count:
            raise RuntimeError(f"epoch {epoch} token total mismatch")
        token_rows = np.memmap(paths['tokens'], mode='r', dtype='<u2', shape=(rows, length))
        observed_stream = hashlib.sha256()
        for start in range(0, rows, 4096):
            chunk = token_rows[start:start+4096]
            mask = np.arange(length)[None, :] < np.asarray(lengths[start:start+4096])[:, None]
            observed_stream.update(chunk[mask].tobytes())
        if observed_stream.hexdigest() != expected_stream.hexdigest():
            raise RuntimeError('Packed token order/content changed')
        if abs(int((objectives == 0).sum()) - int((objectives == 1).sum())) > 1:
            raise RuntimeError(f"epoch {epoch} objective balance mismatch")
        epoch_record = {
            "epoch": epoch,
            "nonpad_stream_sha256": observed_stream.hexdigest(),
            "stream_readback": "PASS_EXACT_ORDER_CONTENT_NO_OMISSIONS",
            "sequence_length": length,
            "rows": rows,
            "nonpad_token_positions": int(lengths.sum(dtype=np.uint64)),
            "raw_words": int(words.sum(dtype=np.uint64)),
            "causal_sequences": int((objectives == 0).sum()),
            "mntp_sequences": int((objectives == 1).sum()),
            "files": {
                name: {"bytes": path.stat().st_size, "sha256": sha256(path)}
                for name, path in paths.items()
            },
        }
        (epoch_dir / "manifest.json").write_text(
            json.dumps(epoch_record, indent=2, sort_keys=True) + "\n", encoding="utf-8"
        )
        epochs.append(epoch_record)
        print(json.dumps(epoch_record, sort_keys=True), flush=True)

    manifest = {
        "schema_version": 1,
        "builder": "DST_CYCLIC_ENRICHED_100M_20260908/scripts/build_schedule_data.py",
        "boundary_policy": "Long-document first chunk length uniform 1..sequence_length each epoch; all prefixes/tails retained; document-isolated causal masks",
        "builder_sha256": sha256(Path(__file__)),
        "schedule": args.schedule,
        "seed": args.seed,
        "raw": raw_manifest,
        "documents": len(documents),
        "words_per_pass": sum(word_counts),
        "token_positions_per_pass": docs_token_count,
        "tokenizer": {
            "path": str(args.tokenizer),
            "tokenizer_json_sha256": sha256(args.tokenizer / "tokenizer.json"),
            "tokenizer_config_sha256": sha256(args.tokenizer / "tokenizer_config.json"),
            "vocab_size": tokenizer.get_vocab_size(),
            "special_ids": observed_special,
        },
        "epochs": epochs,
    }
    manifest_path = args.output / "dataset_manifest.json"
    manifest_path.write_text(json.dumps(manifest, indent=2, sort_keys=True) + "\n", encoding="utf-8")
    print(f"manifest_sha256={sha256(manifest_path)}")


if __name__ == "__main__":
    main()