"""
Tokenize the configured sources into flat uint16 shards plus an index.json.

Sources are weighted: each source contributes roughly `weight` of the final token
budget, sampling round-robin so no single language front-loads the run.

A fraction of documents get the fill-in-the-middle transform applied. FIM is
close to free at prepare time and it is what makes a code model useful for
completion inside an existing file rather than only appending to the end.

Usage:
    python scripts/prepare_data.py --config config/run1.json --tokens 5_000_000_000
"""

import argparse
import hashlib
import json
import os
import random
import sys

import numpy as np
from tokenizers import Tokenizer

sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__))))

from corpus import describe, source_texts
from data import canonical_json_sha256, file_sha256
from scripts.hf_metadata import write_json_atomic

SHARD_TOKENS = 100_000_000  # ~200MB per shard as uint16


def require_fresh_output_dir(path):
    """Refuse every output target except a missing or empty real directory."""
    if not isinstance(path, str) or not path.strip():
        raise ValueError("data output directory must be a non-empty path")
    absolute = os.path.abspath(path)
    if absolute == os.path.abspath(os.sep):
        raise ValueError("data output directory cannot be the filesystem root")
    if os.path.lexists(absolute):
        if os.path.islink(absolute):
            raise ValueError("data output directory cannot be a symbolic link")
        if not os.path.isdir(absolute):
            raise ValueError("data output path exists and is not a directory")
        entries = os.listdir(absolute)
        if entries:
            raise FileExistsError(
                f"refusing non-empty data output directory {absolute}: "
                f"{len(entries)} existing entries"
            )
    return absolute


def validate_build_request(
    config,
    tokens,
    validation_tokens,
    seed,
    fim_rate,
    fim_chunk,
):
    """Bind an attested build to the exact request frozen in its config."""
    if tokens < 1 or validation_tokens < 1:
        raise ValueError("training and validation token budgets must be positive")
    if not 0.0 <= fim_rate <= 1.0:
        raise ValueError("fim_rate must be between 0 and 1")
    if fim_chunk < 0:
        raise ValueError("fim_chunk cannot be negative")
    contract = config.get("data_build_contract")
    if contract is None:
        return None
    expected = {
        "schema_version": 1,
        "train_tokens": tokens,
        "validation_tokens": validation_tokens,
        "seed": seed,
        "require_fresh_output_dir": True,
    }
    if contract != expected:
        raise ValueError(
            "data build request differs from config contract:\n"
            f"expected {contract!r}\n"
            f"actual {expected!r}"
        )
    if fim_rate != config.get("fim_rate"):
        raise ValueError("attested build cannot override config fim_rate")
    if fim_chunk != config.get("fim_chunk"):
        raise ValueError("attested build cannot override config fim_chunk")
    return contract


def apply_fim(ids, sentinels, rng, spm_rate=0.5):
    """Character-free FIM: split the token stream into prefix / middle / suffix."""
    if len(ids) < 16:
        return ids
    a, b = sorted(rng.sample(range(1, len(ids) - 1), 2))
    prefix, middle, suffix = ids[:a], ids[a:b], ids[b:]
    p, m, s = sentinels["prefix"], sentinels["middle"], sentinels["suffix"]
    if rng.random() < spm_rate:
        return [p, s] + suffix + [m] + prefix + middle
    return [p] + prefix + [s] + suffix + [m] + middle


def chunk_document(ids, chunk):
    """
    Split a document into chunks before the FIM transform is applied.

    Whole-document FIM plus random window sampling does not give a FIM-first
    model. The transform frames an entire document, then training draws arbitrary
    `seq_len` windows out of the concatenated stream, so a window landing in the
    middle of a long document sees a fragment: a suffix with no prefix sentinel,
    or a middle with no frame around it at all.

    Measured on the first build, at seq_len 2048: only 43 to 70 percent of
    windows contained all three sentinels despite a correct 70 percent
    document-level transform rate, averaging around 55 percent. The headline
    capability was being diluted by roughly a fifth.

    Chunking first fixes it. With chunks of `chunk` tokens and windows of about
    2051, a window spans roughly two chunks, so it almost always contains at
    least one complete frame. 1024 is chosen so that `2 * chunk` fits inside the
    window; larger chunks reintroduce the problem geometrically.

    The cost is that no single frame spans more than `chunk` tokens of context.
    For a cursor-completion model at 2048 context that is an acceptable trade,
    and it is the capability the model is actually for.
    """
    if not chunk or len(ids) <= chunk:
        return [ids]
    return [ids[i:i + chunk] for i in range(0, len(ids), chunk)]


class ShardWriter:
    def __init__(self, out_dir, split, vocab_size):
        self.out_dir = out_dir
        self.split = split
        self.vocab_size = vocab_size
        self.buf = []
        self.buf_len = 0
        self.shards = []
        self.total = 0
        os.makedirs(out_dir, exist_ok=True)

    def add(self, ids):
        self.buf.append(np.asarray(ids, dtype=np.uint16))
        self.buf_len += len(ids)
        self.total += len(ids)
        if self.buf_len >= SHARD_TOKENS:
            self.flush()

    def flush(self):
        if not self.buf:
            return
        arr = np.concatenate(self.buf)
        name = f"{self.split}_{len(self.shards):04d}.bin"
        path = os.path.join(self.out_dir, name)
        if os.path.lexists(path):
            raise FileExistsError(f"refusing to replace shard: {path}")
        digest = hashlib.sha256()
        digest.update(memoryview(arr).cast("B"))
        with open(path, "xb") as f:
            arr.tofile(f)
            f.flush()
            os.fsync(f.fileno())
        self.shards.append(
            {
                "path": name,
                "tokens": int(arr.shape[0]),
                "bytes": int(arr.nbytes),
                "sha256": digest.hexdigest(),
            }
        )
        print(f"  wrote {name}  ({arr.shape[0]:,} tokens, {self.total:,} total)")
        self.buf, self.buf_len = [], 0

    def manifest(self):
        for entry in self.shards:
            entry["total_tokens"] = self.total
        return self.shards


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--config", required=True)
    ap.add_argument("--tokens", type=int, default=5_000_000_000)
    ap.add_argument("--val-tokens", type=int, default=20_000_000)
    ap.add_argument("--fim-rate", type=float, default=None,
                    help="override the config's fim_rate (default 0.5 if neither is set)")
    ap.add_argument("--fim-chunk", type=int, default=None,
                    help="split documents into this many tokens before applying FIM, "
                         "so a training window contains a complete frame. 0 disables.")
    ap.add_argument("--seed", type=int, default=1337)
    cli = ap.parse_args()

    with open(cli.config) as f:
        cfg = json.load(f)

    fim_rate = cli.fim_rate if cli.fim_rate is not None else cfg.get("fim_rate", 0.5)
    fim_chunk = cli.fim_chunk if cli.fim_chunk is not None else cfg.get("fim_chunk", 0)
    build_contract = validate_build_request(
        cfg,
        cli.tokens,
        cli.val_tokens,
        cli.seed,
        fim_rate,
        fim_chunk,
    )
    print(f"fim_rate {fim_rate}, fim_chunk {fim_chunk or 'off (whole documents)'}")

    out_dir = require_fresh_output_dir(cfg["data_dir"])
    tok = Tokenizer.from_file(cfg["tokenizer_path"])
    vocab_size = tok.get_vocab_size()
    if vocab_size > 65535:
        raise ValueError("uint16 shards require vocab_size <= 65535")

    eos = tok.token_to_id("<|endoftext|>")
    sentinels = {
        "prefix": tok.token_to_id("<|fim_prefix|>"),
        "middle": tok.token_to_id("<|fim_middle|>"),
        "suffix": tok.token_to_id("<|fim_suffix|>"),
    }
    if eos is None or any(v is None for v in sentinels.values()):
        raise ValueError("tokenizer is missing required special tokens")

    rng = random.Random(cli.seed)
    sources = cfg["sources"]
    weights = np.array([s.get("weight", 1.0) for s in sources], dtype=np.float64)
    weights = weights / weights.sum()
    budgets = (weights * cli.tokens).astype(np.int64)

    val_writer = ShardWriter(out_dir, "val", vocab_size)
    train_writer = ShardWriter(out_dir, "train", vocab_size)

    # Validation is allocated per source, in the same proportions as training.
    # Filling one shared counter from the first source instead would hand the
    # entire validation set to whichever language happens to be listed first,
    # and the val loss would then be blind to every other language in the
    # mixture. With python first at 24 percent that is exactly what happened.
    val_budgets = (weights * cli.val_tokens).astype(np.int64)

    for src, budget, val_budget in zip(sources, budgets, val_budgets):
        print(f"source {describe(src)}: target {budget:,} train, "
              f"{val_budget:,} val tokens")
        produced = 0
        val_remaining = int(val_budget)
        batch_texts = []
        stream = source_texts(src)

        def drain(texts):
            nonlocal produced, val_remaining
            if not texts:
                return
            for enc in tok.encode_batch(texts):
                for ids in chunk_document(enc.ids, fim_chunk):
                    if fim_rate > 0 and rng.random() < fim_rate:
                        ids = apply_fim(ids, sentinels, rng)
                    ids = ids + [eos]
                    emit(ids)

        def emit(ids):
            nonlocal produced, val_remaining
            if val_remaining > 0:
                val_writer.add(ids)
                val_remaining -= len(ids)
            else:
                train_writer.add(ids)
                produced += len(ids)

        for text in stream:
            batch_texts.append(text)
            if len(batch_texts) >= 1000:
                drain(batch_texts)
                batch_texts = []
                if produced >= budget:
                    break
        drain(batch_texts)
        print(f"  produced {produced:,} tokens")

    val_writer.flush()
    train_writer.flush()

    index = {
        "schema_version": 2 if build_contract is not None else 1,
        "vocab_size": vocab_size,
        "fim_rate": fim_rate,
        "fim_chunk": fim_chunk,
        "splits": {"train": train_writer.manifest(), "val": val_writer.manifest()},
    }
    if build_contract is not None:
        index["build"] = {
            "completed": True,
            "train_tokens_requested": cli.tokens,
            "validation_tokens_requested": cli.val_tokens,
            "seed": cli.seed,
            "fresh_output_directory": True,
            "config_path": cli.config,
            "config_canonical_sha256": canonical_json_sha256(cfg),
            "tokenizer_path": cfg["tokenizer_path"],
            "tokenizer_sha256": file_sha256(cfg["tokenizer_path"]),
            "sources": sources,
            "sources_canonical_sha256": canonical_json_sha256(sources),
        }
    index_path = os.path.join(out_dir, "index.json")
    write_json_atomic(index_path, index)
    print(f"wrote {index_path}: "
          f"{train_writer.total:,} train / {val_writer.total:,} val tokens")


if __name__ == "__main__":
    main()