"""Turn a directory of plain-text files into the parquet shards nanochat trains on.

Purpose: prepare a corpus of your own for pretraining: split the text into
         documents at paragraph boundaries, drop exact duplicates, shuffle with a
         recorded seed, hold out a validation split as whole documents, and write
         the result as parquet shards with a "text" column in the layout
         nanochat's dataset loader expects.
Platform: all (pure Python plus pyarrow, which nanochat already depends on)
Minimum memory: 8 GB
Assumes: pyarrow is installed (it is a nanochat dependency, so nanochat's own
         virtual environment has it); the input directory contains UTF-8 .txt
         files; the output directory is a base directory for the domain corpus
         and not the one holding the shards downloaded in the lab.

Usage: python make-domain-shards.py --input domain-corpus/text
                                    [--out-dir ~/.cache/nanochat-domain/base_data_climbmix]
                                    [--target-chars 4000] [--val-fraction 0.05]
                                    [--seed 0] [--force]

The loader treats the alphabetically last parquet file as the validation split and
every other file as training data, so this script always writes at least two shards
and names them shard_00000.parquet upwards.
"""
from __future__ import annotations

import argparse
import hashlib
import json
import os
import random
import re
import sys
from pathlib import Path

# A blank line is the most reliable document boundary in plain prose. Chapters
# would be better but their markers differ from book to book, and a paragraph run
# of a few thousand characters is a reasonable unit for a model with a sequence
# length in the hundreds.
PARAGRAPH_BREAK = re.compile(r"\n\s*\n")


def split_into_documents(text: str, target_chars: int) -> list[str]:
    """Group paragraphs into documents of roughly target_chars characters."""
    documents, current, size = [], [], 0
    for paragraph in PARAGRAPH_BREAK.split(text):
        paragraph = paragraph.strip()
        if not paragraph:
            continue
        current.append(paragraph)
        size += len(paragraph)
        if size >= target_chars:
            documents.append("\n\n".join(current))
            current, size = [], 0
    if current:
        documents.append("\n\n".join(current))
    return documents


def deduplicate(documents: list[str]) -> tuple[list[str], int]:
    """Drop exact duplicates, keeping the first occurrence. Returns (kept, dropped)."""
    seen, kept = set(), []
    for doc in documents:
        digest = hashlib.sha256(doc.encode("utf-8")).hexdigest()
        if digest in seen:
            continue
        seen.add(digest)
        kept.append(doc)
    return kept, len(documents) - len(kept)


def write_shard(path: Path, documents: list[str], rows_per_group: int) -> None:
    """Write one parquet file with a single "text" column and several row groups.

    nanochat's loader iterates row groups rather than whole files, and uses the
    group index to shard across ranks, so a file with one enormous group works but
    parallelises badly. Several groups per file costs nothing and behaves better.
    """
    import pyarrow as pa
    import pyarrow.parquet as pq

    table = pa.table({"text": pa.array(documents, type=pa.string())})
    pq.write_table(table, path, row_group_size=max(1, rows_per_group))


def main() -> None:
    parser = argparse.ArgumentParser(description=__doc__.splitlines()[0])
    parser.add_argument("--input", required=True, help="directory of UTF-8 .txt files")
    parser.add_argument(
        "--out-dir",
        default=os.path.join(
            os.environ.get("NANOCHAT_BASE_DIR", str(Path.home() / ".cache" / "nanochat-domain")),
            "base_data_climbmix",
        ),
        help="where the shards are written; must be the directory nanochat reads",
    )
    parser.add_argument("--target-chars", type=int, default=4000, help="approximate characters per document")
    parser.add_argument("--val-fraction", type=float, default=0.05, help="share of documents held out")
    parser.add_argument("--docs-per-shard", type=int, default=2000)
    parser.add_argument("--rows-per-group", type=int, default=200)
    parser.add_argument("--min-chars", type=int, default=200, help="drop documents shorter than this")
    parser.add_argument("--seed", type=int, default=0)
    parser.add_argument("--force", action="store_true", help="overwrite an output directory that already has shards")
    args = parser.parse_args()

    source_dir = Path(args.input).expanduser()
    files = sorted(source_dir.glob("*.txt"))
    if not files:
        sys.exit(f"No .txt files under {source_dir}.")

    out_dir = Path(args.out_dir).expanduser()
    existing = sorted(out_dir.glob("*.parquet")) if out_dir.exists() else []
    if existing and not args.force:
        sys.exit(
            f"{out_dir} already holds {len(existing)} parquet file(s).\n"
            "Refusing to mix two corpora in one directory: choose another --out-dir, or pass\n"
            "--force if you are certain these shards should be replaced."
        )
    out_dir.mkdir(parents=True, exist_ok=True)
    for stale in existing:
        stale.unlink()

    # ---- read and split ---------------------------------------------------
    documents, per_file = [], {}
    for path in files:
        text = path.read_text(encoding="utf-8", errors="replace")
        docs = [d for d in split_into_documents(text, args.target_chars) if len(d) >= args.min_chars]
        per_file[path.name] = len(docs)
        documents.extend(docs)
    print(f"read {len(files)} file(s) -> {len(documents):,} document(s)")

    documents, dropped = deduplicate(documents)
    print(f"deduplicated: dropped {dropped:,}, kept {len(documents):,}")
    if not documents:
        sys.exit("Nothing left after filtering; lower --min-chars or check the input.")

    # Shuffle before splitting, so the validation set is not the end of one book.
    random.Random(args.seed).shuffle(documents)

    n_val = max(1, round(len(documents) * args.val_fraction))
    if n_val >= len(documents):
        sys.exit("--val-fraction leaves no training data.")
    val_docs = documents[:n_val]
    train_docs = documents[n_val:]

    # ---- write ------------------------------------------------------------
    # The loader takes the alphabetically last file as validation, so the training
    # shards are written first and the validation shard last.
    shards, index = [], 0
    for start in range(0, len(train_docs), args.docs_per_shard):
        chunk = train_docs[start : start + args.docs_per_shard]
        path = out_dir / f"shard_{index:05d}.parquet"
        write_shard(path, chunk, args.rows_per_group)
        shards.append((path, len(chunk), sum(len(d) for d in chunk)))
        index += 1
    val_path = out_dir / f"shard_{index:05d}.parquet"
    write_shard(val_path, val_docs, args.rows_per_group)
    shards.append((val_path, len(val_docs), sum(len(d) for d in val_docs)))

    train_chars = sum(len(d) for d in train_docs)
    val_chars = sum(len(d) for d in val_docs)

    manifest = {
        "input": str(source_dir),
        "out_dir": str(out_dir),
        "seed": args.seed,
        "target_chars": args.target_chars,
        "min_chars": args.min_chars,
        "val_fraction": args.val_fraction,
        "documents_per_file": per_file,
        "duplicates_dropped": dropped,
        "train_documents": len(train_docs),
        "val_documents": len(val_docs),
        "train_characters": train_chars,
        "val_characters": val_chars,
        "shards": [{"file": p.name, "documents": n, "characters": c} for p, n, c in shards],
    }
    manifest_path = out_dir / "corpus-manifest.json"
    manifest_path.write_text(json.dumps(manifest, indent=2), encoding="utf-8")

    print(f"\nwrote {len(shards)} shard(s) to {out_dir}")
    for path, n_docs, chars in shards:
        role = "validation" if path == val_path else "train"
        print(f"  {path.name:<24} {n_docs:>7,} docs  {chars:>12,} chars  ({role})")
    print(f"\ntrain characters: {train_chars:,}")
    print(f"val characters:   {val_chars:,}")
    print(f"manifest:         {manifest_path}")
    print(
        "\nA character count is not a token count. Train the tokeniser on these shards,\n"
        "run tok_eval, and divide the characters by the bytes-per-token ratio it reports\n"
        "to get the token budget this corpus actually provides."
    )


if __name__ == "__main__":
    main()