import os
from huggingface_hub import hf_hub_download

def download_models():
    """Download all models from the dedicated Model Repo using huggingface_hub."""
    repo_id = "icanedit2/bakoai-models"
    models_to_download = [
        "player_detector.pt",
        "ball_detector.pt",
        "court_keypoint_detector.pt",
        "swish_ball_rim.pt",
        "swish_pose.pt",
        "yolov8n-pose.pt"
    ]
    
    os.makedirs("models", exist_ok=True)
    
    token = os.environ.get("HF_TOKEN")
    
    for model_file in models_to_download:
        dest_path = os.path.join("models", model_file)
        if os.path.exists(dest_path):
            print(f"✅ {dest_path} already exists.")
            continue
            
        print(f"⏳ Downloading {model_file} from {repo_id}...")
        try:
            # hf_hub_download automatically handles LFS, Xet, and auth
            path = hf_hub_download(
                repo_id=repo_id,
                filename=model_file,
                local_dir="models",
                token=token
            )
            print(f"✅ Successfully downloaded {model_file} to {path}")
        except Exception as e:
            print(f"❌ Failed to download {model_file}: {e}")
            raise

if __name__ == "__main__":
    download_models()