#!/usr/bin/env python3
"""Aggregate GHCN-Daily station files into weather-heavy JSONL documents."""

from __future__ import annotations

import argparse
import sys
from collections import defaultdict
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.noaa import (
    aggregate_monthly,
    iter_ghcn_daily_csv,
    monthly_docs_for_state,
    resolve_noaa_states_and_max_per_state,
    station_id_to_state_map,
    state_rollup_doc,
)


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)
    noaa = cfg["noaa"]
    paths = cfg["paths"]
    raw_dir = root / paths["raw_noaa"]
    out_path = root / paths["processed"] / "weather.jsonl"
    files = sorted(raw_dir.glob("US*.csv.gz")) + sorted(raw_dir.glob("US*.csv"))
    if not files:
        raise SystemExit(f"No station CSV(.gz) files under {raw_dir}. Run scripts/download_noaa.py first.")

    states_cfg, _ = resolve_noaa_states_and_max_per_state(noaa)
    fallback_state = states_cfg[0]
    inv_path = raw_dir / "ghcnd-stations.txt"
    if inv_path.is_file():
        sid_to_state = station_id_to_state_map(inv_path.read_text(encoding="utf-8", errors="replace"))
    else:
        sid_to_state = {}

    all_docs: list[dict] = []
    seed = int(cfg["merge"]["seed"])
    y0, y1 = int(noaa["year_start"]), int(noaa["year_end"])
    monthly_global: dict = {}
    stations_by_state: dict[str, list[str]] = defaultdict(list)

    for fp in tqdm(files, desc="stations"):
        sid = fp.name.split(".")[0]
        st = sid_to_state.get(sid, fallback_state)
        stations_by_state[st].append(sid)
        daily = iter_ghcn_daily_csv(fp)
        monthly = aggregate_monthly(daily, y0, y1)
        for k, v in monthly.items():
            monthly_global[k] = v
        docs = monthly_docs_for_state(st, [sid], monthly, y0, y1, seed)
        all_docs.extend(docs)

    # One July rollup per state represented in the download (correct geography in text)
    for st in sorted(stations_by_state.keys()):
        sids = stations_by_state[st]
        sub_seed = seed + sum(ord(c) for c in st) * 1_000_003
        rollup = state_rollup_doc(st, sids, monthly_global, y0, y1, month=7, seed=sub_seed)
        if rollup:
            all_docs.append(rollup)

    write_jsonl(out_path, iter(all_docs))
    print("Wrote", len(all_docs), "documents to", out_path)
    return 0


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