from datasets import load_dataset, Audio
import soundfile as sf
import numpy as np
import io
from pathlib import Path

Path("data").mkdir(exist_ok=True)

ds = load_dataset("oddadmix/arabic-audio-collection-mostafa-mahmoud", split="train")
ds = ds.cast_column("audio", Audio(decode=False))

for row in ds:
    dur = row.get("duration") or 0
    if not (6.0 <= dur <= 10.0):
        continue
    audio_bytes = row["audio"].get("bytes")
    if not audio_bytes:
        continue
    try:
        arr, sr = sf.read(io.BytesIO(audio_bytes))
        arr = arr.astype(np.float32)
        snr = arr.std() / (np.abs(arr).mean() + 1e-9)
        if snr > 0.15:
            sf.write("data/ref.wav", arr, sr)
            print(f"ref.wav written — dur={dur:.1f}s  snr={snr:.3f}")
            print(row["transcript_text"][:80])
            break
    except Exception as e:
        print(f"skip: {e}")