"""Pre-download XTTS v2.0 model files during Docker build"""
import os, urllib.request, time, sys

MODEL_DIR = os.path.expanduser("~/.local/share/tts/tts_models--multilingual--multi-dataset--xtts_v2")
os.makedirs(MODEL_DIR, exist_ok=True)

FILES = [
    ("https://coqui.gateway.scarf.sh/hf-coqui/XTTS-v2/main/dvae.pth", "dvae.pth"),
    ("https://coqui.gateway.scarf.sh/hf-coqui/XTTS-v2/main/mel_stats.pth", "mel_stats.pth"),
    ("https://coqui.gateway.scarf.sh/hf-coqui/XTTS-v2/main/model.pth", "model.pth"),
    ("https://coqui.gateway.scarf.sh/hf-coqui/XTTS-v2/main/config.json", "config.json"),
    ("https://coqui.gateway.scarf.sh/hf-coqui/XTTS-v2/main/vocab.json", "vocab.json"),
]

for url, name in FILES:
    dest = os.path.join(MODEL_DIR, name)
    if os.path.exists(dest) and os.path.getsize(dest) > 1000:
        print(f"{name} already cached ({os.path.getsize(dest)/1e6:.0f} MB)")
        continue
    for attempt in range(5):
        try:
            print(f"Downloading {name}... attempt {attempt+1}/5", flush=True)
            urllib.request.urlretrieve(url, dest)
            size = os.path.getsize(dest) / 1e6
            print(f"  OK ({size:.0f} MB)")
            break
        except Exception as e:
            print(f"  Error: {str(e)[:100]}")
            if attempt < 4:
                time.sleep(5)
            else:
                print(f"  All retries exhausted for {name}")

print("Pre-download complete!")