#!/usr/bin/env python3
"""Fetch the clips the listener's languages are measured on.

  python3 tools/asr/eval/fetch_fleurs.py OUT [CODE=CONFIG ...]

The first 25 test clips (in the archive's order) of each language of FLEURS
(google/fleurs, CC BY 4.0) at a pinned revision, without downloading whole
archives: test.tar.gz is streamed and left once 25 clips listed in test.tsv
are out. Writes OUT/<code>/<file>.wav (16 kHz mono) and OUT/<code>/refs.json
(the clip ids and their raw transcriptions). With no CODE=CONFIG, every
language in LANGS. A language already fetched is left alone.

Standard library only (Python 3.9 or later).
"""
import csv
import io
import json
import os
import sys
import tarfile
import urllib.request

REV = "70bb2e84b976b7e960aa89f1c648e09c59f894dd"
BASE = f"https://huggingface.co/datasets/google/fleurs/resolve/{REV}/data"
CLIPS = 25
# The app's language → FLEURS's configuration (Spanish and Arabic: the only
# ones FLEURS has, Latin American and Egyptian).
LANGS = {
    "ar": "ar_eg", "de": "de_de", "en": "en_us", "es": "es_419", "fr": "fr_fr", "hi": "hi_in", "it": "it_it",
    "ja": "ja_jp", "ko": "ko_kr", "nl": "nl_nl", "pl": "pl_pl", "pt": "pt_br", "ru": "ru_ru", "tr": "tr_tr",
    "uk": "uk_ua", "vi": "vi_vn",
}


def main(argv):
    if len(argv) < 2:
        print(__doc__, file=sys.stderr)
        return 2
    out_root = argv[1]
    langs = dict(a.split("=") for a in argv[2:]) or LANGS
    csv.field_size_limit(10**9)
    for lang, cfg in langs.items():
        d = os.path.join(out_root, lang)
        if os.path.exists(os.path.join(d, "refs.json")):
            continue
        os.makedirs(d, exist_ok=True)
        tsv = urllib.request.urlopen(f"{BASE}/{cfg}/test.tsv").read().decode("utf-8")
        rows = {}
        for r in csv.reader(io.StringIO(tsv), delimiter="\t", quoting=csv.QUOTE_NONE):
            # id, file_name, raw_transcription, transcription, phonemes, num_samples, gender
            rows[r[1]] = {"id": r[0], "raw": r[2], "norm": r[3], "num_samples": int(r[5]), "gender": r[6]}
        got = []
        with urllib.request.urlopen(f"{BASE}/{cfg}/audio/test.tar.gz") as resp, tarfile.open(fileobj=resp, mode="r|gz") as tf:
            for m in tf:
                name = os.path.basename(m.name)
                if not m.isfile() or name not in rows:
                    continue
                with open(os.path.join(d, name), "wb") as f:
                    f.write(tf.extractfile(m).read())
                got.append({"file": name, **rows[name]})
                if len(got) >= CLIPS:
                    break
        with open(os.path.join(d, "refs.json"), "w", encoding="utf-8") as f:
            json.dump({"source": f"google/fleurs@{REV} {cfg} test (CC-BY-4.0)", "clips": got}, f, ensure_ascii=False, indent=1)
        print(lang, cfg, len(got), "clips", flush=True)
    return 0


if __name__ == "__main__":
    sys.exit(main(sys.argv))