# dl_models.py
import open_clip
import torch
import os
import gc
from huggingface_hub import snapshot_download
from transformers import AutoTokenizer

if "HF_HUB_OFFLINE" in os.environ: del os.environ["HF_HUB_OFFLINE"]

MODEL_ZOO = {
    # Large
    "CLIP_ViT-L-14_Laion2B":   {'type': 'open_clip', 'arch': 'ViT-L-14', 'data': 'laion2b_s32b_b82k'},
    "CLIP_ViT-L-14_OpenAI":    {'type': 'open_clip', 'arch': 'ViT-L-14-quickgelu', 'data': 'openai'},
    # Medium
    "EVA02_B-16":              {'type': 'open_clip', 'arch': 'EVA02-B-16', 'data': 'merged2b_s8b_b131k'},
    "ConvNeXt_Base":           {'type': 'open_clip', 'arch': 'convnext_base', 'data': 'laion400m_s13b_b51k'},
    # ALIGN & SigLIP (BASE VARIANTS)
    "ALIGN_RoBERTa_ViT-B":     {'type': 'open_clip', 'arch': 'xlm-roberta-base-ViT-B-32', 'data': 'laion5b_s13b_b90k'},
    "SigLIP_ViT-B-16":         {'type': 'open_clip', 'arch': 'ViT-B-16-SigLIP', 'data': 'webli'},
    # Standard Base
    "CLIP_ViT-B-32_Laion2B":   {'type': 'open_clip', 'arch': 'ViT-B-32', 'data': 'laion2b_s34b_b79k'},
    "DFN_ViT-B-32":            {'type': 'open_clip', 'arch': 'ViT-B-32', 'data': 'datacomp_xl_s13b_b90k'},
    "CLIP_ViT-B-32_OpenAI":    {'type': 'open_clip', 'arch': 'ViT-B-32-quickgelu', 'data': 'openai'},
    "CoCa_ViT-B-32":           {'type': 'open_clip', 'arch': 'coca_ViT-B-32', 'data': 'laion2b_s13b_b90k'},
}

def cleanup(): gc.collect(); torch.cuda.empty_cache()

def download_full_zoo():
    print(f"===========================================================")
    print(f"STARTING DOWNLOAD (ALIGN BASE + SIGLIP BASE)")
    print(f"===========================================================\n")

    for key, config in MODEL_ZOO.items():
        arch = config['arch']
        data = config['data']
        print(f"Processing: {key} ({arch})")
        
        # 1. Weights
        try:
            print("  Weights...", end=" ", flush=True)
            model, _, _ = open_clip.create_model_and_transforms(arch, pretrained=data, device='cpu')
            print("DONE")
            del model
        except Exception as e: print(f"FAILED ({e})")

        # 2. Manual Dependencies
        hf_extra = None
        
        # Check for ALIGN Base (xlm-roberta-base)
        if "roberta-base" in arch:
            hf_extra = "xlm-roberta-base"
        # Check for SigLIP Base (timm/ViT-B-16-SigLIP)
        elif "siglip" in arch.lower():
            hf_extra = "timm/ViT-B-16-SigLIP"

        if hf_extra:
            try:
                print(f"  Snapshotting '{hf_extra}'...", end=" ", flush=True)
                snapshot_download(
                    repo_id=hf_extra, 
                    allow_patterns=["*.json", "*.txt", "*.model"], 
                    local_dir=f"./{hf_extra}",  # Download locally to handle ALIGN logic
                    local_dir_use_symlinks=False
                )
                print("DONE")
            except Exception as e: print(f"FAILED ({e})")

        cleanup()
        print("-" * 60)

if __name__ == "__main__":
    download_full_zoo()