"""Materialize every corpus in corpora.yaml into data/corpora/<name>.txt.

    python src/prepare_data.py            # download what's missing
    python src/prepare_data.py --force    # re-download everything

Sources: `url` (single text file), `hf` (HuggingFace datasets), `path` (already
local — just validated).
"""

from __future__ import annotations

import argparse
import os
import sys
import urllib.request

import config

# The machine-wide HuggingFace cache (~/.cache/huggingface) is root-owned and
# not writable for us, so keep datasets' cache inside the project.
os.environ.setdefault("HF_HOME", os.path.join(config.ROOT, ".hf_cache"))


def _from_url(url: str) -> str:
    with urllib.request.urlopen(url, timeout=60) as resp:
        return resp.read().decode("utf-8", errors="replace")


def _from_hf(src: dict) -> str:
    from datasets import load_dataset
    ds = load_dataset(src["dataset"], src.get("config"), split=src.get("split", "train"))
    col = src.get("text_column", "text")
    # WikiText ships lots of blank lines; keep non-empty ones.
    parts = [t for t in ds[col] if t and t.strip()]
    return "".join(parts)


def prepare_one(corpus: dict, force: bool = False) -> None:
    name = corpus["name"]
    src = corpus["source"]
    out_path = config.corpus_text_path(name)

    if src["type"] == "path":
        local = src["path"] if os.path.isabs(src["path"]) else os.path.join(config.ROOT, src["path"])
        if not os.path.exists(local):
            print(f"  [{name}] MISSING local file: {local}", file=sys.stderr)
            return
        if os.path.abspath(local) != os.path.abspath(out_path):
            text = open(local, encoding="utf-8", errors="replace").read()
            _write(out_path, text)
        print(f"  [{name}] using local file ({os.path.getsize(out_path)/1e6:.2f} MB)")
        return

    if os.path.exists(out_path) and not force:
        print(f"  [{name}] already present ({os.path.getsize(out_path)/1e6:.2f} MB) — skip")
        return

    print(f"  [{name}] fetching via {src['type']} ...")
    text = _from_url(src["url"]) if src["type"] == "url" else _from_hf(src)
    _write(out_path, text)
    print(f"  [{name}] wrote {os.path.getsize(out_path)/1e6:.2f} MB to {out_path}")


def _write(path: str, text: str) -> None:
    os.makedirs(os.path.dirname(path), exist_ok=True)
    with open(path, "w", encoding="utf-8") as f:
        f.write(text)


def main() -> None:
    ap = argparse.ArgumentParser()
    ap.add_argument("--force", action="store_true", help="re-download even if present")
    args = ap.parse_args()

    os.makedirs(config.DATA_DIR, exist_ok=True)
    print(f"Preparing corpora -> {config.DATA_DIR}")
    for corpus in config.corpora():
        try:
            prepare_one(corpus, force=args.force)
        except Exception as e:  # noqa: BLE001 — report and continue with others
            print(f"  [{corpus['name']}] FAILED: {e}", file=sys.stderr)
    print("Done.")


if __name__ == "__main__":
    main()