#!/usr/bin/env python3
"""Stream English Wikipedia articles into chunked JSONL (Hugging Face `datasets`)."""

from __future__ import annotations

import argparse
import sys
from pathlib import Path

ROOT = Path(__file__).resolve().parents[1]
sys.path.insert(0, str(ROOT / "src"))

from tqdm import tqdm

from weather_llm.config_loader import load_yaml, repo_root_from
from weather_llm.data.jsonl_utils import write_jsonl
from weather_llm.data.wiki import wikipedia_to_chunks


def main() -> int:
    ap = argparse.ArgumentParser()
    ap.add_argument("--config", type=Path, default=None)
    args = ap.parse_args()
    root = repo_root_from(ROOT)
    cfg_path = args.config or (root / "configs/data_default.yaml")
    cfg = load_yaml(cfg_path)
    wiki = cfg["wikipedia"]
    paths = cfg["paths"]
    out_path = root / paths["processed"] / "wikipedia.jsonl"

    stream = wikipedia_to_chunks(
        dataset_name=str(wiki["dataset_name"]),
        dataset_config=str(wiki["dataset_config"]),
        max_articles=int(wiki["max_articles"]),
        min_article_chars=int(wiki["min_article_chars"]),
        chunk_target_chars=int(wiki["chunk_target_chars"]),
        chunk_overlap_chars=int(wiki["chunk_overlap_chars"]),
    )
    rows = []
    for row in tqdm(stream, desc="wiki chunks"):
        rows.append(row)
    write_jsonl(out_path, iter(rows))
    print("Wrote", len(rows), "chunks to", out_path)
    return 0


if __name__ == "__main__":
    raise SystemExit(main())