"""Download only the selected GLaDOS reference WAVs from the HF dataset API."""
import json
import pathlib
import urllib.parse
import urllib.request

DATASET = "ray0rf1re/GLaDOS-audio-v2"
INDICES = [0, 7, 25, 46, 53, 55, 72, 109, 133, 142, 143, 177, 208, 224, 234, 258, 285, 304, 349, 553]
OUT = pathlib.Path(__file__).parent / "references"

def row_for(index: int) -> dict:
    query = urllib.parse.urlencode({"dataset": DATASET, "config": "default", "split": "train", "offset": index, "length": 1})
    with urllib.request.urlopen("https://datasets-server.huggingface.co/rows?" + query, timeout=60) as response:
        return json.load(response)["rows"][0]

def main() -> None:
    OUT.mkdir(exist_ok=True)
    for index in INDICES:
        item = row_for(index)
        row = item["row"]
        audio_url = row["audio"][0]["src"]
        target = OUT / f"{index:04d}.wav"
        if not target.exists():
            urllib.request.urlretrieve(audio_url, target)
        print(index, row["filename"], target, sep="\t")

if __name__ == "__main__":
    main()