#!/usr/bin/env python3
"""
Model Download Script for Certificate Pipeline
Downloads all required models to /data/onlyface/models
"""

import os
import sys
import urllib.request
from pathlib import Path
from tqdm import tqdm


MODEL_DIR = Path("/data/onlyface/models")
MODEL_DIR.mkdir(parents=True, exist_ok=True)


MODELS = {
    "birefnet-general.onnx": "https://huggingface.co/zhengpeng7/BiRefNet/resolve/main/BiRefNet-general-epoch_244.onnx",
    "rmbg-2.0.onnx": "https://huggingface.co/briaai/RMBG-2.0/resolve/main/rmbg-2.0.onnx",
    "big-lama.pt": "https://github.com/Sanster/models/releases/download/add_big_lama/big-lama.pt",
    "inswapper_128.onnx": "https://huggingface.co/deepinsight/inswapper/resolve/main/inswapper_128.onnx",
    "genderage.onnx": "https://huggingface.co/deepinsight/genderage/resolve/main/genderage.onnx",
}


class DownloadProgressBar(tqdm):
    def update_to(self, b=1, bsize=1, tsize=None):
        if tsize is not None:
            self.total = tsize
        self.update(b * bsize - self.n)


def download_model(url: str, filepath: Path):
    if filepath.exists():
        print(f"Already exists: {filepath.name}")
        return
    
    print(f"Downloading {filepath.name}...")
    try:
        with DownloadProgressBar(unit='B', unit_scale=True, miniters=1, desc=filepath.name) as t:
            urllib.request.urlretrieve(url, filepath, reporthook=t.update_to)
        print(f"Done: {filepath.name}")
    except urllib.error.HTTPError as e:
        print(f"Failed to download {filepath.name}: {e}")
        print(f"  URL: {url}")


def download_paddleocr_models():
    """PaddleOCR models are downloaded automatically on first use."""
    print("PaddleOCR models will be downloaded automatically on first run.")


def main():
    print(f"Model directory: {MODEL_DIR}")
    print("=" * 50)
    
    for filename, url in MODELS.items():
        filepath = MODEL_DIR / filename
        download_model(url, filepath)
    
    print("=" * 50)
    print("Downloading InsightFace buffalo_l models...")
    download_paddleocr_models()
    
    print("\nAll models downloaded!")
    print(f"Models saved to: {MODEL_DIR}")


if __name__ == "__main__":
    main()