import os
import urllib.request

MODELS_DIR = "models"
MODELS = {
    "yolov8n.pt": "https://github.com/ultralytics/assets/releases/download/v8.2.0/yolov8n.pt",
    "FastSAM-s.pt": "https://github.com/ultralytics/assets/releases/download/v8.2.0/FastSAM-s.pt"
}

def download_file(url: str, dest: str):
    print(f"Downloading {os.path.basename(dest)}...")
    try:
        urllib.request.urlretrieve(url, dest)
        print(f"Successfully downloaded to {dest}")
    except Exception as e:
        print(f"Failed to download from {url}: {e}")

def download():
    os.makedirs(MODELS_DIR, exist_ok=True)
    
    # 1. Download PyTorch models if missing
    for filename, url in MODELS.items():
        filepath = os.path.join(MODELS_DIR, filename)
        if not os.path.exists(filepath):
            download_file(url, filepath)
        else:
            print(f"Model {filename} already exists in {MODELS_DIR}/. Skipping download.")

    # 2. Pre-download OWLv2 via transformers (caches to HF_HOME which is models/hf_cache)
    print("Pre-downloading OWLv2 model...")
    try:
        from transformers import Owlv2ForObjectDetection, Owlv2Processor
        MODEL_ID = "google/owlv2-base-patch16-ensemble"
        # The transformers library automatically handles skipping if already cached
        Owlv2Processor.from_pretrained(MODEL_ID)
        Owlv2ForObjectDetection.from_pretrained(MODEL_ID)
        print("OWLv2 downloaded and cached successfully.")
    except Exception as e:
        print(f"Failed to download OWLv2: {e}")

    # 3. Pre-fetch caption fonts (otherwise lazily downloaded at first render,
    #    which hard-fails on cold boots with no network headroom)
    print("Pre-fetching caption fonts...")
    try:
        from captiona.core.styles import StyleManager
        StyleManager().ensure_fonts()
        print("Caption fonts cached successfully.")
    except Exception as e:
        print(f"Failed to prefetch fonts: {e}")

if __name__ == "__main__":
    download()