"""Prepare frozen TinyStories streams for both training arms."""

import hashlib
from itertools import chain
import json
from pathlib import Path
import tempfile
from urllib.request import Request, urlopen

import numpy as np
import sentencepiece as spm


ROOT = Path(__file__).resolve().parent
CONFIG_PATH = ROOT / "protocol.json"
CONFIG = json.loads(CONFIG_PATH.read_bytes())
BASELINE = ROOT.with_name("stories260k-baseline")
REFERENCE = json.loads((BASELINE / "protocol.json").read_bytes())
TOKENIZER = BASELINE / "assets/tok512.model"
VALIDATION = BASELINE / "assets/TinyStories-valid.txt"
DELIMITER = b"<|endoftext|>"
DATA = ROOT / "data"
TOKENS = DATA / "train.u16"
MANIFEST = DATA / "manifest.json"
BLOCK = 1024 * 1024


def sha256(path):
    with path.open("rb") as stream:
        return hashlib.file_digest(stream, "sha256").hexdigest()


def baseline_asset(path):
    expected = REFERENCE["assets"][path.relative_to(BASELINE).as_posix()]["sha256"]
    if sha256(path) != expected:
        raise ValueError(f"Baseline input hash differs for {path.name}")
    return expected


def atomic_json(path, value):
    temp = path.with_suffix(".tmp")
    temp.write_text(json.dumps(value, indent=2, allow_nan=False) + "\n", encoding="utf-8")
    temp.replace(path)


def stories(stream, digest=None):
    pending = b""
    while chunk := stream.read(BLOCK):
        if digest is not None:
            digest.update(chunk)
        records = (pending + chunk).split(DELIMITER)
        pending = records.pop()
        for record in records:
            yield record.strip()


def corpus_tokens(config):
    return config.get("data_tokens", config["training_tokens"])


def verify(config=CONFIG, directory=DATA):
    manifest = json.loads((directory / MANIFEST.name).read_bytes())
    tokens = directory / TOKENS.name
    if (manifest["training_source"] != config["training_source"]
            or manifest["requested_bytes"] != config["training_source_limit_bytes"]
            or manifest.get("source_sha256") != config.get("training_source_sha256")
            or manifest["tokenizer_sha256"] != baseline_asset(TOKENIZER)
            or manifest["validation_sha256"] != baseline_asset(VALIDATION)
            or manifest["token_count"] != corpus_tokens(config) + 1
            or tokens.stat().st_size != 2 * manifest["token_count"]
            or sha256(tokens) != manifest["tokens_sha256"]):
        raise ValueError("Prepared training data does not match its protocol and manifest")
    return manifest


def prepare(config=CONFIG, directory=DATA):
    manifest_path = directory / MANIFEST.name
    if manifest_path.exists():
        return verify(config, directory)
    tokenizer_sha = baseline_asset(TOKENIZER)
    validation_sha = baseline_asset(VALIDATION)
    tokenizer = spm.SentencePieceProcessor(model_file=str(TOKENIZER))
    if tokenizer.vocab_size() != REFERENCE["architecture"]["vocab_size"]:
        raise ValueError("Unexpected vocabulary size")
    excluded = {hashlib.sha256(row.strip()).digest() for row in VALIDATION.read_bytes().split(DELIMITER)}
    source_sha = config.get("training_source_sha256")
    source_digest = hashlib.sha256() if source_sha else None
    needed = corpus_tokens(config) + 1
    limit = config["training_source_limit_bytes"]
    headers = {} if source_sha else {"Range": f"bytes=0-{limit - 1}"}
    request = Request(config["training_source"], headers=headers)
    directory.mkdir(exist_ok=True)
    with tempfile.NamedTemporaryFile(dir=directory, delete=False) as output:
        temp = Path(output.name)
    try:
        written = kept = skipped = 0
        with urlopen(request, timeout=60) as stream, temp.open("wb") as output:
            if source_sha:
                if stream.status != 200 or stream.headers.get("Content-Length") != str(limit):
                    raise ValueError("Expected the entire pinned training source")
            elif stream.status != 206 or not stream.headers.get("Content-Range", "").startswith(f"bytes 0-{limit - 1}/"):
                raise ValueError("Training source did not honor the bounded range request")
            batch = []
            for story in chain(stories(stream, source_digest), (None,)):
                if story:
                    if hashlib.sha256(story).digest() in excluded:
                        skipped += 1
                    else:
                        batch.append(story.decode("utf-8"))
                if len(batch) < 256 and story is not None:
                    continue
                for ids in tokenizer.encode(batch, add_bos=True, add_eos=True):
                    if not source_sha:
                        ids = ids[:needed - written]
                    output.write(np.asarray(ids, dtype="<u2").tobytes())
                    written += len(ids)
                    kept += 1
                    if not source_sha and written == needed:
                        break
                if not source_sha and written == needed:
                    break
                batch.clear()
            if source_sha:
                if source_digest.hexdigest() != source_sha:
                    raise ValueError("Full training source hash differs from the protocol")
                context = config["context"]
                written = (written - 1) // context * context + 1
                output.truncate(written * 2)
        if written != needed:
            raise ValueError(f"Training source yielded {written:,} tokens, expected {needed:,}")
        manifest = {"training_source": config["training_source"], "requested_bytes": limit,
                    "tokenizer_sha256": tokenizer_sha, "validation_sha256": validation_sha,
                    "tokens_sha256": sha256(temp), "token_count": written, "stories_used": kept,
                    "exact_validation_duplicates_skipped": skipped}
        if source_sha:
            manifest["source_sha256"] = source_sha
        temp.replace(directory / TOKENS.name)
        atomic_json(manifest_path, manifest)
        return manifest
    finally:
        temp.unlink(missing_ok=True)