#!/usr/bin/env python3
"""Create a tiny mixed corpus for offline tokenizer/smoke tests (no network)."""

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 weather_llm.config_loader import repo_root_from
from weather_llm.data.jsonl_utils import write_jsonl
from weather_llm.data.merge import merge_weighted_shuffle


def main() -> int:
    ap = argparse.ArgumentParser()
    ap.add_argument("--out_subdir", type=str, default="processed", help="Under data/")
    args = ap.parse_args()
    root = repo_root_from(ROOT)
    proc = root / "data" / args.out_subdir
    proc.mkdir(parents=True, exist_ok=True)
    weather = [
        {
            "text": (
                "In Texas, July is typically very hot, with mean daily highs often near 95–100°F "
                "based on historical GHCN-Daily-style summaries. This reflects past observations, "
                "not current conditions."
            ),
            "source": "synthetic",
        },
        {
            "text": (
                "Historical station aggregates for Minnesota in January often show subzero lows and "
                "snow-prone patterns; specifics depend on year and site. Not a live forecast."
            ),
            "source": "synthetic",
        },
    ]
    wiki_like = [
        {
            "text": (
                "Photosynthesis\n\nPhotosynthesis is a process used by plants and other organisms "
                "to convert light energy into chemical energy that can be later released to fuel "
                "the organism's activities."
            ),
            "source": "synthetic_wiki",
        },
        {
            "text": (
                "Climate of the United States\n\nThe climate of the United States varies due to "
                "changes in latitude and a range of geographic features."
            ),
            "source": "synthetic_wiki",
        },
        {
            "text": (
                "Atlantic hurricane season\n\nThe Atlantic hurricane season is the period in a year "
                "when hurricanes usually form in the Atlantic Ocean."
            ),
            "source": "synthetic_wiki",
        },
        {
            "text": (
                "Drought\n\nA drought is an event of prolonged shortages in the water supply, "
                "whether atmospheric, surface water or ground water."
            ),
            "source": "synthetic_wiki",
        },
        {
            "text": (
                "El Niño\n\nEl Niño is a climate pattern that describes the unusual warming of "
                "surface waters in the eastern tropical Pacific Ocean."
            ),
            "source": "synthetic_wiki",
        },
    ]
    # Duplicate with light paraphrase to give SentencePiece enough surface forms on tiny installs.
    weather_extra = []
    for i in range(12):
        weather_extra.append(
            {
                "text": (
                    f"Sample {i}: Arizona summers are hot and dry in many lower-elevation areas, "
                    "while winters are milder; this is descriptive climatology, not a current observation."
                ),
                "source": "synthetic",
            }
        )
    weather.extend(weather_extra)
    write_jsonl(proc / "weather.jsonl", iter(weather))
    write_jsonl(proc / "wikipedia.jsonl", iter(wiki_like))
    merge_weighted_shuffle(
        weather_path=proc / "weather.jsonl",
        wiki_path=proc / "wikipedia.jsonl",
        out_train=proc / "train.jsonl",
        out_val=proc / "val.jsonl",
        weather_weight=0.85,
        wiki_weight=0.15,
        val_ratio=0.25,
        seed=42,
    )
    print("Wrote minimal corpus under", proc)
    return 0


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