# download_models.py
import os
import gdown
import zipfile

def download_all():
    # Make necessary directories
    os.makedirs("output", exist_ok=True)
    os.makedirs("P&ID-Symbols-3/train", exist_ok=True)

    # File ID mapping
    file_map = {
        "output/checkpoint0009.pth": "1WfsV8ZuDwlgvBsompA8jdpG_XqNGsWGT",
        "output/checkpoint_best_total.pth": "1UpVLNeKDrocU4UgBe361IwvHnpFkKues",
        "rf-detr-base.pth": "1L7mU1jyQNLxJcex3jTd5wccdtUrhgQ2c",
        "P&ID-Symbols-3/train/_annotations.coco.json": "159ArLMxS1PZ4zD6CAf7bbJfEvLToFA-C",
        "P&ID-Symbols-3.zip": "1HfjxUN7j92XCciKWB4_ZkC_W6UUKWtDd"
    }

    # Download each file
    for path, file_id in file_map.items():
        if not os.path.exists(path):
            print(f"\nDownloading {path} ...")
            gdown.download(id=file_id, output=path, quiet=False)

    # Unzip image folder if not already extracted
    zip_path = "P&ID-Symbols-3.zip"
    if os.path.exists(zip_path):
        print(f"\nExtracting {zip_path} ...")
        with zipfile.ZipFile(zip_path, 'r') as zip_ref:
            zip_ref.extractall(".")
        print("? Extraction complete.")

    print("? All files downloaded and set up.")