"""Download public benchmark parquets and preserve exact normalized task splits.

This script performs no model work. Public SST-2 test labels are withheld (-1);
the validation split is therefore the labelled held-out evaluation split.
"""
from __future__ import annotations

import hashlib
import json
import time
import urllib.parse
import urllib.request
from collections import Counter
from concurrent.futures import ThreadPoolExecutor
from datetime import datetime, timezone
from pathlib import Path

import pyarrow as pa
import pyarrow.parquet as pq

ROOT = Path(__file__).resolve().parent / "data"
SOURCES = [("sst2", "stanfordnlp/sst2", "default"), ("race_middle", "ehovy/race", "middle")]


def digest(path: Path) -> str:
    value = hashlib.sha256()
    with path.open("rb") as stream:
        for block in iter(lambda: stream.read(1024 * 1024), b""):
            value.update(block)
    return value.hexdigest()


def fetch(url: str, path: Path) -> None:
    if path.exists():
        return
    path.parent.mkdir(parents=True, exist_ok=True)
    temporary = path.with_suffix(path.suffix + ".partial")
    for attempt in range(4):
        try:
            request = urllib.request.Request(url, headers={"User-Agent": "WiggleGPT-benchmark-preparation/1.0"})
            with urllib.request.urlopen(request, timeout=90) as response, temporary.open("wb") as stream:
                while block := response.read(1024 * 1024):
                    stream.write(block)
            temporary.replace(path)
            return
        except Exception:
            if attempt == 3:
                raise
            time.sleep(2 ** attempt)


def normalize(task: str, row: dict, index: int, split: str) -> dict:
    if task == "sst2":
        return {"example_id": f"sst2/{split}/{row.get('idx', index)}", "text": row["sentence"], "label": int(row["label"])}
    answer = row["answer"]
    label = "ABCD".index(answer) if isinstance(answer, str) else int(answer)
    options = list(row["options"])
    if len(options) != 4 or not 0 <= label < 4:
        raise ValueError(f"Invalid RACE record {index}")
    return {"example_id": str(row.get("example_id", f"race_middle/{split}/{index}")),
            "article": row["article"], "question": row["question"], "options": options, "label": label}


def prepare_source(spec: tuple[str, str, str]) -> dict:
    task, dataset, config = spec
    metadata_url = "https://datasets-server.huggingface.co/parquet?" + urllib.parse.urlencode({"dataset": dataset})
    metadata_path = ROOT / "source_parquets" / task / "parquet_metadata.json"
    fetch(metadata_url, metadata_path)
    metadata = json.loads(metadata_path.read_text(encoding="utf-8"))
    if metadata.get("partial"):
        raise RuntimeError(f"Partial parquet conversion for {dataset}")
    files = [entry for entry in metadata["parquet_files"] if entry["config"] == config]
    result = {"dataset": dataset, "config": config, "metadata_url": metadata_url, "splits": {}}
    for split in ("train", "validation", "test"):
        split_files = sorted([entry for entry in files if entry["split"] == split], key=lambda entry: entry["filename"])
        if not split_files:
            raise RuntimeError(f"Missing {dataset} {config} {split}")
        destination = ROOT / f"{task}_{split}.jsonl"
        temporary = destination.with_suffix(".jsonl.partial")
        count, labels, source_files = 0, Counter(), []
        with temporary.open("w", encoding="utf-8", newline="\n") as output:
            for ordinal, entry in enumerate(split_files):
                parquet_path = ROOT / "source_parquets" / task / split / f"{ordinal:03d}_{entry['filename']}"
                fetch(entry["url"], parquet_path)
                if parquet_path.stat().st_size != entry["size"]:
                    raise RuntimeError(f"Size mismatch for {parquet_path}")
                source_files.append({"url": entry["url"], "path": str(parquet_path), "bytes": parquet_path.stat().st_size, "sha256": digest(parquet_path)})
                for batch in pq.ParquetFile(parquet_path).iter_batches(batch_size=1024):
                    for row in batch.to_pylist():
                        item = normalize(task, row, count, split)
                        output.write(json.dumps(item, ensure_ascii=False) + "\n")
                        labels[str(item["label"])] += 1
                        count += 1
        temporary.replace(destination)
        result["splits"][split] = {"path": str(destination), "rows": count, "sha256": digest(destination), "label_counts": dict(sorted(labels.items())), "labelled": all(int(label) >= 0 for label in labels), "source_parquets": source_files}
        print(f"{task} {split}: {count:,} rows; labels={dict(labels)}", flush=True)
    return result


def main() -> None:
    ROOT.mkdir(parents=True, exist_ok=True)
    pa.set_cpu_count(2)
    pa.set_io_thread_count(2)
    with ThreadPoolExecutor(max_workers=2) as pool:
        prepared = list(pool.map(prepare_source, SOURCES))
    manifest = {"prepared_at_utc": datetime.now(timezone.utc).isoformat(), "tasks": {spec[0]: result for spec, result in zip(SOURCES, prepared)},
                "split_policy": {"sst2": "Train on official train; validation is the only public labelled holdout. Official test labels are hidden (-1).",
                                 "race_middle": "Train on official middle/train; select using middle/validation; evaluate once on middle/test."},
                "normalization": "Exact input strings retained; sentiment labels are 0=negative,1=positive; RACE A/B/C/D mapped to integer 0/1/2/3; no examples generated or reassigned."}
    manifest_path = ROOT / "dataset_manifest.json"
    manifest_path.write_text(json.dumps(manifest, ensure_ascii=False, indent=2) + "\n", encoding="utf-8")
    print(f"Manifest: {manifest_path}", flush=True)


if __name__ == "__main__":
    main()