import os
import sys
import torch
import gdown
from pathlib import Path

MODELS_DIR = Path("models")

def download_models():
    """Download required models if they don't exist."""
    os.makedirs(MODELS_DIR, exist_ok=True)
    
    # Model URLs (you should replace these with actual URLs to your models)
    models = {
        "yolov8n.pt": "https://github.com/ultralytics/assets/releases/download/v0.0.0/yolov8n.pt",
        "osnet_x0_25_msmt17.pt": "https://github.com/mikel-brostrom/assets/releases/download/v0.0.0/osnet_x0_25_msmt17.pt"
    }
    
    # Download missing models
    for model_name, url in models.items():
        model_path = MODELS_DIR / model_name
        if not model_path.exists():
            print(f"Downloading {model_name}...")
            try:
                torch.hub.download_url_to_file(url, model_path)
                print(f"✅ Downloaded {model_name}")
            except Exception as e:
                print(f"⚠️ Failed to download {model_name}: {e}")

if __name__ == "__main__":
    download_models()