import wget
import os

# Base URL for downloading model files
BASE_URL = "https://huggingface.co./microsoft/OmniParser/resolve/main"

# Define folder structure and create folders
os.makedirs("weights/icon_detect", exist_ok=True)
os.makedirs("weights/icon_caption_florence", exist_ok=True)
os.makedirs("weights/icon_caption_blip2", exist_ok=True)

# List of files to download
model_files = {
    "weights/icon_detect/model.safetensors": f"{BASE_URL}/icon_detect/model.safetensors",
    "weights/icon_detect/model.yaml": f"{BASE_URL}/icon_detect/model.yaml",
    "weights/icon_caption_florence/model.safetensors": f"{BASE_URL}/icon_caption_florence/model.safetensors",
    "weights/icon_caption_florence/config.json": f"{BASE_URL}/icon_caption_florence/config.json",
    "weights/icon_caption_blip2/pytorch_model-00001-of-00002.bin": f"{BASE_URL}/icon_caption_blip2/pytorch_model-00001-of-00002.bin",
    "weights/icon_caption_blip2/pytorch_model-00002-of-00002.bin": f"{BASE_URL}/icon_caption_blip2/pytorch_model-00002-of-00002.bin",
    "weights/icon_caption_blip2/pytorch_model.bin.index.json": f"{BASE_URL}/icon_caption_blip2/pytorch_model.bin.index.json",
    "weights/icon_caption_blip2/config.json": f"{BASE_URL}/icon_caption_blip2/config.json"
}

# Download each file
for file_path, url in model_files.items():
    try:
        print(f"Downloading {url} to {file_path}...")
        wget.download(url, file_path)
        print(f"\nDownloaded {file_path}")
    except Exception as e:
        print(f"Failed to download {url}: {e}")
        exit(1)

print("All required model and configuration files downloaded and organised.")