#!/usr/bin/env python3
"""Fetch the two public AA-WER v2 datasets and build an Echo suite manifest."""

from __future__ import annotations

import concurrent.futures
import hashlib
import json
import pathlib
import urllib.parse
import urllib.request


ROOT = pathlib.Path(__file__).resolve().parent
DATASETS = {
    "voxpopuli": {
        "repo": "ArtificialAnalysis/VoxPopuli-Cleaned-AA",
        "revision": "07adf4a242e2e5dd955a230909c177573c978ee3",
        "expected_rows": 628,
    },
    "earnings22": {
        "repo": "ArtificialAnalysis/Earnings22-Cleaned-AA",
        "revision": "5999936ee3d8d3cde160e3c65a8cc6ff5b7c5276",
        "expected_rows": 6,
    },
}


def get_json(url: str) -> dict:
    with urllib.request.urlopen(url, timeout=120) as response:
        return json.load(response)


def rows_for(repo: str, expected_rows: int) -> list[dict]:
    rows: list[dict] = []
    for offset in range(0, expected_rows, 100):
        query = urllib.parse.urlencode(
            {
                "dataset": repo,
                "config": "default",
                "split": "test",
                "offset": offset,
                "length": min(100, expected_rows - offset),
            }
        )
        payload = get_json(f"https://datasets-server.huggingface.co/rows?{query}")
        rows.extend(item["row"] for item in payload["rows"])
    if len(rows) != expected_rows:
        raise RuntimeError(f"{repo}: expected {expected_rows} rows, received {len(rows)}")
    return rows


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


def download_one(repo: str, revision: str, remote_path: str, target: pathlib.Path) -> dict:
    target.parent.mkdir(parents=True, exist_ok=True)
    if not target.exists():
        encoded_path = "/".join(urllib.parse.quote(part) for part in remote_path.split("/"))
        url = f"https://huggingface.co/datasets/{repo}/resolve/{revision}/{encoded_path}"
        request = urllib.request.Request(url, headers={"User-Agent": "Echo-AA-public-benchmark/1.0"})
        temporary = target.with_suffix(target.suffix + ".partial")
        with urllib.request.urlopen(request, timeout=300) as response, temporary.open("wb") as output:
            while chunk := response.read(1024 * 1024):
                output.write(chunk)
        temporary.replace(target)
    return {
        "path": str(target.relative_to(ROOT)),
        "bytes": target.stat().st_size,
        "sha256": sha256(target),
    }


def main() -> None:
    fixtures: list[dict] = []
    source_records: list[dict] = []
    downloads: list[tuple[str, str, str, pathlib.Path]] = []

    for dataset_name, metadata in DATASETS.items():
        dataset_rows = rows_for(metadata["repo"], metadata["expected_rows"])
        for row in dataset_rows:
            relative_audio = pathlib.Path("audio") / dataset_name / row["file_name"]
            fixture_id = f"{dataset_name}:{row['id']}"
            fixtures.append(
                {
                    "id": fixture_id,
                    "audio": str(relative_audio),
                    "reference": row["transcript"],
                    "language": row["language"],
                }
            )
            source_records.append(
                {
                    "fixture_id": fixture_id,
                    "dataset": dataset_name,
                    "source_repo": metadata["repo"],
                    "revision": metadata["revision"],
                    "source_path": row["url"],
                    "duration_seconds": row["duration"],
                    "reference": row["transcript"],
                    "local_audio": str(relative_audio),
                }
            )
            downloads.append(
                (
                    metadata["repo"],
                    metadata["revision"],
                    row["url"],
                    ROOT / relative_audio,
                )
            )

    with concurrent.futures.ThreadPoolExecutor(max_workers=12) as executor:
        file_records = list(executor.map(lambda args: download_one(*args), downloads))

    file_by_path = {record["path"]: record for record in file_records}
    for record in source_records:
        record.update(file_by_path[record["local_audio"]])

    manifest = {
        "name": "AA-WER v2 public subset — Echo Core ML reconstruction",
        "provenance": (
            "All 628 VoxPopuli-Cleaned-AA v1 samples and all 6 "
            "Earnings22-Cleaned-AA v1 parent recordings, pinned to the "
            "Artificial Analysis Hugging Face revisions recorded in sources.json. "
            "AA-AgentTalk is proprietary and excluded."
        ),
        "fixtures": fixtures,
    }
    (ROOT / "manifest.json").write_text(
        json.dumps(manifest, indent=2, ensure_ascii=False) + "\n",
        encoding="utf-8",
    )
    (ROOT / "sources.json").write_text(
        json.dumps(
            {
                "protocol": "AA-WER v2 public subset reconstruction",
                "datasets": DATASETS,
                "rows": source_records,
            },
            indent=2,
            ensure_ascii=False,
        )
        + "\n",
        encoding="utf-8",
    )
    print(
        json.dumps(
            {
                "fixtures": len(fixtures),
                "audio_bytes": sum(record["bytes"] for record in file_records),
                "manifest": str(ROOT / "manifest.json"),
            },
            indent=2,
        )
    )


if __name__ == "__main__":
    main()
