from __future__ import annotations

import argparse
import os

import py7zr
import requests

RECORD = "4783391"
FILES = {
    "development": ("clotho_audio_development.7z", "clotho_captions_development.csv"),
    "validation": ("clotho_audio_validation.7z", "clotho_captions_validation.csv"),
    "evaluation": ("clotho_audio_evaluation.7z", "clotho_captions_evaluation.csv"),
}

def zenodo_files():
    r = requests.get(f"https://zenodo.org/api/records/{RECORD}", timeout=30)
    r.raise_for_status()
    return {f["key"]: (f["links"]["self"], f["size"]) for f in r.json()["files"]}

def download(url, dest, expected_size=None):
    if os.path.exists(dest) and (expected_size is None or os.path.getsize(dest) == expected_size):
        print(f"[fetch] {dest} already present ({os.path.getsize(dest)} bytes), skipping", flush=True)
        return
    with requests.get(url, stream=True, timeout=60) as r:
        r.raise_for_status()
        total = int(r.headers.get("content-length", 0))
        done = 0
        with open(dest + ".part", "wb") as f:
            for chunk in r.iter_content(chunk_size=1 << 20):
                f.write(chunk)
                done += len(chunk)
                if total:
                    print(f"\r[fetch] {os.path.basename(dest)} {done/1e6:.0f}/{total/1e6:.0f} MB",
                          end="", flush=True)
    print(flush=True)
    os.rename(dest + ".part", dest)
    if expected_size is not None and os.path.getsize(dest) != expected_size:
        raise RuntimeError(f"{dest}: size {os.path.getsize(dest)} != expected {expected_size}")

def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--root", default="/root/clotho_raw")
    ap.add_argument("--splits", nargs="+", default=["development", "validation", "evaluation"])
    args = ap.parse_args()
    os.makedirs(args.root, exist_ok=True)

    links = zenodo_files()
    for split in args.splits:
        audio_key, caption_key = FILES[split]
        audio_path = os.path.join(args.root, audio_key)
        caption_path = os.path.join(args.root, caption_key)

        audio_url, audio_size = links[audio_key]
        caption_url, caption_size = links[caption_key]
        print(f"[fetch] {split}: {audio_key} ({audio_size/1e6:.0f} MB)", flush=True)
        download(audio_url, audio_path, audio_size)
        download(caption_url, caption_path, caption_size)

        extract_dir = os.path.join(args.root, split)
        if not os.path.isdir(extract_dir):
            print(f"[fetch] extracting {audio_key}", flush=True)
            with py7zr.SevenZipFile(audio_path, mode="r") as z:
                z.extractall(path=args.root)
        wavs = [f for f in os.listdir(extract_dir) if f.endswith(".wav")]
        print(f"[fetch] {split}: {len(wavs)} wav files in {extract_dir}", flush=True)

    print("FETCHDONE", flush=True)

if __name__ == "__main__":
    main()