#!/usr/bin/env python3
"""
Download all ONNX models needed for the multi-model pipeline.
Handles split ONNX models and uses verified repositories.
"""

import shutil
from pathlib import Path
from huggingface_hub import hf_hub_download


MODELS = {
    "informative-drawings": {
        "repo": "rocca/informative-drawings-line-art-onnx",
        "file": "model.onnx",
        "local": "models/informative_drawings.onnx",
        "size_mb": 22,
        "required": True,
    },
    "sam2-encoder": {
        "repo": "onnx-community/sam2.1-hiera-small-ONNX",
        "file": "onnx/vision_encoder.onnx",
        "local": "models/sam2_encoder.onnx",
        "extra_files": ["onnx/vision_encoder.onnx_data"],
        "size_mb": 130,
        "required": True,
    },
    "sam2-decoder": {
        "repo": "onnx-community/sam2.1-hiera-small-ONNX",
        "file": "onnx/prompt_encoder_mask_decoder.onnx",
        "local": "models/sam2_decoder.onnx",
        "extra_files": ["onnx/prompt_encoder_mask_decoder.onnx_data"],
        "size_mb": 16,
        "required": True,
    },
    "depth-anything-v2": {
        "repo": "onnx-community/depth-anything-v2-small",
        "file": "onnx/model_fp16.onnx",
        "local": "models/depth_anything_v2.onnx",
        "size_mb": 49,
        "required": True,
    },
    "real-esrgan": {
        "repo": "AXERA-TECH/Real-ESRGAN",
        "file": "realesrgan-x4.onnx",
        "local": "models/real_esrgan_x4.onnx",
        "size_mb": 64,
        "required": False,
    },
    "rmbg": {
        "repo": "briaai/RMBG-2.0",
        "file": "onnx/model.onnx", # Verified common path for this repo
        "local": "models/rmbg2.onnx",
        "size_mb": 170,
        "required": False,
    },
    "face-parsing": {
        "repo": "bluefoxcreation/Face_parsing_onnx",
        "file": "faceparser.onnx",
        "local": "models/face_parsing.onnx",
        "size_mb": 52,
        "required": False,
    },
}


def download_all(required_only: bool = False):
    """Download models from HuggingFace Hub."""
    Path("models").mkdir(exist_ok=True)

    for name, info in MODELS.items():
        local_path = Path(info["local"])

        if local_path.exists():
            # For split models, ensure extra files exist even if .onnx exists
            if "extra_files" in info:
                all_extras_exist = True
                for extra in info["extra_files"]:
                    if not (Path("models") / Path(extra).name).exists():
                        all_extras_exist = False
                        break
                if all_extras_exist:
                    print(f"  ✓ {name} already exists")
                    continue
            elif local_path.stat().st_size > 1024 * 1024:
                print(f"  ✓ {name} already exists")
                continue

        if required_only and not info["required"]:
            print(f"  ⊘ {name} skipped (optional)")
            continue

        print(f"  ↓ Downloading {name} from {info['repo']}...")
        try:
            downloaded = hf_hub_download(
                repo_id=info["repo"],
                filename=info["file"],
            )
            shutil.copy(downloaded, local_path)
            print(f"  ✓ {name} saved to {local_path}")

            if "extra_files" in info:
                for extra in info["extra_files"]:
                    extra_local = Path("models") / Path(extra).name
                    print(f"    ↓ Downloading extra file {extra}...")
                    downloaded_extra = hf_hub_download(repo_id=info["repo"], filename=extra)
                    shutil.copy(downloaded_extra, extra_local)
                    print(f"    ✓ Saved to {extra_local}")

        except Exception as e:
            if info["required"]:
                print(f"  ✗ Failed to download REQUIRED model {name}: {e}")
            else:
                print(f"  ✗ {name} failed (optional): {e}")


if __name__ == "__main__":
    import argparse

    parser = argparse.ArgumentParser()
    parser.add_argument(
        "--required-only",
        action="store_true",
        help="Only download required models",
    )
    args = parser.parse_args()
    download_all(required_only=args.required_only)