"""Encode the raw corpus into flat uint16 token files.

Each story in the source is delimited by a line containing only
``<|endoftext|>``. We stream the file one line at a time, encode each story
with the trained tokenizer, append the end-of-text token id after it, and flush
the ids to disk in batches — so memory stays flat regardless of corpus size.

Output: data/tokens/train.bin and data/tokens/val.bin.

Usage:
    python scripts/encode_corpus.py
"""

import sys
import time
from pathlib import Path

import numpy as np

ROOT = Path(__file__).resolve().parent.parent
sys.path.insert(0, str(ROOT / "src"))

from dataset import TOKEN_DTYPE  # noqa: E402
from tokenizer import BPETokenizer  # noqa: E402

EOT = "<|endoftext|>"
TOK_PATH = ROOT / "artifacts" / "tokenizer.json"
RAW = ROOT / "data" / "raw"
OUT = ROOT / "data" / "tokens"
FLUSH_EVERY = 20_000  # stories between disk writes / progress updates

JOBS = [
    ("TinyStoriesV2-GPT4-train.txt", "train.bin"),
    ("TinyStoriesV2-GPT4-valid.txt", "val.bin"),
]


def encode_file(tok: BPETokenizer, src: Path, dst: Path, progress: bool = True) -> tuple[int, int]:
    """Encode one corpus file to a uint16 .bin. Returns (n_stories, n_tokens)."""
    eot_id = tok.special_tokens[EOT]
    buf: list[int] = []
    story: list[str] = []
    n_stories = n_tokens = 0
    t0 = time.perf_counter()

    def emit(text: str) -> None:
        nonlocal n_stories, n_tokens
        text = text.strip()
        if not text:
            return  # skip blank runs between separators
        ids = tok.encode(text)
        ids.append(eot_id)
        buf.extend(ids)
        n_stories += 1
        n_tokens += len(ids)

    with open(src, "r", encoding="utf-8", errors="ignore") as fin, open(dst, "wb") as fout:
        for line in fin:
            if line.strip() == EOT:
                emit("".join(story))
                story = []
                if n_stories % FLUSH_EVERY == 0 and buf:
                    np.asarray(buf, dtype=TOKEN_DTYPE).tofile(fout)
                    buf.clear()
                    if progress:
                        rate = n_stories / (time.perf_counter() - t0)
                        print(f"\r  {dst.name}: {n_stories:,} stories, "
                              f"{n_tokens:,} tokens ({rate:,.0f}/s)", end="", flush=True)
            else:
                story.append(line)
        emit("".join(story))  # trailing story with no final separator
        if buf:
            np.asarray(buf, dtype=TOKEN_DTYPE).tofile(fout)
    if progress:
        print()
    return n_stories, n_tokens


def main() -> None:
    if not TOK_PATH.exists():
        sys.exit(f"missing {TOK_PATH} — run scripts/train_tokenizer.py first")
    tok = BPETokenizer.load(TOK_PATH)
    OUT.mkdir(parents=True, exist_ok=True)

    for src_name, dst_name in JOBS:
        src, dst = RAW / src_name, OUT / dst_name
        if not src.exists():
            sys.exit(f"missing {src} — run scripts/download_data.py first")
        print(f"encoding {src_name} -> {dst_name}")
        n_stories, n_tokens = encode_file(tok, src, dst)
        size_mb = dst.stat().st_size / 1e6
        print(f"  done: {n_stories:,} stories, {n_tokens:,} tokens, {size_mb:,.0f} MB")


if __name__ == "__main__":
    main()