"""Phase 3 — préparation du corpus texte tokenisé pour l'entraînement.

Produit deux tableaux binaires `uint16` (le vocabulaire 32 000 tient sur 16 bits) :
`train.bin` et `val.bin`, plus un `meta.json` (provenance, tailles, dtypes).

Chaque document est terminé par `<|eos|>` : le modèle apprend où s'arrêtent les
documents. Aucun modèle pré-entraîné n'est utilisé — seulement du texte brut.

Usage :

    python -m src.data.prepare_text --preset wikipedia_fr --max-mb 200 \\
        --out-dir data/processed/wikipedia_fr
"""

from __future__ import annotations

import argparse
import json
from collections.abc import Iterable, Iterator
from pathlib import Path

import numpy as np

from src.tokenizer import CortexTokenizer, load_tokenizer
from src.tokenizer.special_tokens import EOS_ID
from src.tokenizer.train_tokenizer import CORPUS_PRESETS, iter_hf_corpus, iter_local_corpus

ROOT = Path(__file__).resolve().parents[2]
DTYPE = np.uint16


def tokenize_corpus(
    texts: Iterable[str],
    tokenizer: CortexTokenizer,
    add_eos: bool = True,
    batch_size: int = 1000,
) -> np.ndarray:
    """Concatène les documents tokenisés (+ `<|eos|>` entre eux) en un tableau `uint16`.

    Les documents sont encodés par **lots** (``encode_batch``, implémenté en Rust),
    ce qui est nettement plus rapide qu'un appel Python par document.
    """
    out: list[int] = []
    batch: list[str] = []

    def flush() -> None:
        if not batch:
            return
        for ids in tokenizer.raw.encode_batch(batch):
            out.extend(ids.ids)
            if add_eos:
                out.append(EOS_ID)
        batch.clear()

    for text in texts:
        batch.append(text)
        if len(batch) >= batch_size:
            flush()
    flush()

    arr = np.asarray(out, dtype=np.int64)
    if arr.size and arr.max() > np.iinfo(DTYPE).max:
        raise ValueError("id de token > 65535 : dtype uint16 insuffisant")
    return arr.astype(DTYPE)


def split_train_val(arr: np.ndarray, val_frac: float = 0.01) -> tuple[np.ndarray, np.ndarray]:
    """Découpe séquentielle (pas de fuite : la val est la fin du flux)."""
    if not 0.0 < val_frac < 1.0:
        raise ValueError("val_frac doit être dans ]0, 1[")
    n_val = max(1, int(len(arr) * val_frac))
    return arr[:-n_val], arr[-n_val:]


def save_bin(arr: np.ndarray, path: Path) -> None:
    path = Path(path)
    path.parent.mkdir(parents=True, exist_ok=True)
    arr.tofile(path)


def load_bin(path: Path, mode: str = "r") -> np.memmap:
    """Charge un tableau tokenisé en mémoire virtuelle (aucune copie RAM)."""
    return np.memmap(path, dtype=DTYPE, mode=mode)


def prepare(
    texts: Iterable[str],
    out_dir: Path,
    tokenizer: CortexTokenizer,
    val_frac: float = 0.01,
    extra_meta: dict | None = None,
) -> dict:
    """Tokenise, découpe, écrit `train.bin`/`val.bin`/`meta.json` et retourne les métadonnées."""
    out_dir = Path(out_dir)
    out_dir.mkdir(parents=True, exist_ok=True)

    arr = tokenize_corpus(texts, tokenizer, add_eos=True)
    train, val = split_train_val(arr, val_frac)
    save_bin(train, out_dir / "train.bin")
    save_bin(val, out_dir / "val.bin")

    meta = {
        "tokenizer_vocab_size": tokenizer.vocab_size,
        "dtype": "uint16",
        "add_eos": True,
        "val_frac": val_frac,
        "n_train_tokens": int(train.size),
        "n_val_tokens": int(val.size),
        "n_total_tokens": int(arr.size),
    }
    if extra_meta:
        meta.update(extra_meta)
    (out_dir / "meta.json").write_text(json.dumps(meta, ensure_ascii=False, indent=2), encoding="utf-8")
    return meta


def _iter_source(args) -> Iterator[str]:
    max_bytes = int(args.max_mb * 1024 * 1024)
    if args.local_corpus:
        return iter_local_corpus(args.local_corpus, max_bytes=max_bytes)
    return iter_hf_corpus(args.preset, max_bytes=max_bytes)


def main() -> None:
    parser = argparse.ArgumentParser(description="Préparer un corpus texte tokenisé.")
    src = parser.add_mutually_exclusive_group(required=True)
    src.add_argument("--local-corpus", type=Path)
    src.add_argument("--preset", choices=sorted(CORPUS_PRESETS))
    parser.add_argument("--max-mb", type=float, default=200.0)
    parser.add_argument("--val-frac", type=float, default=0.01)
    parser.add_argument("--out-dir", type=Path, required=True)
    parser.add_argument("--tokenizer-dir", type=Path, default=None)
    args = parser.parse_args()

    tokenizer = load_tokenizer(args.tokenizer_dir)
    meta = prepare(
        _iter_source(args),
        args.out_dir,
        tokenizer,
        val_frac=args.val_frac,
        extra_meta={"preset": args.preset, "max_mb": args.max_mb},
    )
    print(json.dumps(meta, ensure_ascii=False, indent=2))


if __name__ == "__main__":
    main()