"""
utils/download_checkpoints.py
------------------------------
Downloads all required model checkpoints if not already present.
Run once before using the pipeline:  python -m utils.download_checkpoints
"""

import os
import sys

# allow running standalone
sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__))))

from config import (
    CHECKPOINT_DIR,
    GDINO_CONFIG, GDINO_WEIGHTS, GDINO_CONFIG_URL, GDINO_WEIGHTS_URL,
    SAM_CHECKPOINT, SAM_CHECKPOINT_URL,
    LAMA_CHECKPOINT, LAMA_URL,
    INPAINT_METHOD,
)
from image_utils import download_file, download_text_file, unzip


def download_groundingdino() -> None:
    print("\n[1/3] GroundingDINO weights")
    download_text_file(GDINO_CONFIG_URL, GDINO_CONFIG)
    download_file(GDINO_WEIGHTS_URL, GDINO_WEIGHTS, "GroundingDINO weights")


def download_sam() -> None:
    print("\n[2/3] SAM (Segment Anything) weights")
    download_file(SAM_CHECKPOINT_URL, SAM_CHECKPOINT, "SAM ViT-H")


def download_lama() -> None:
    print("\n[3/3] LaMa inpainting weights")
    if os.path.isdir(LAMA_CHECKPOINT):
        print("  [DONE] LaMa checkpoint directory already exists")
        return
    zip_path = os.path.join(CHECKPOINT_DIR, "big-lama.zip")
    download_file(LAMA_URL, zip_path, "LaMa big-lama")
    unzip(zip_path, CHECKPOINT_DIR)
    os.remove(zip_path)
    print("  [DONE] LaMa extracted")


def main() -> None:
    os.makedirs(CHECKPOINT_DIR, exist_ok=True)
    print("=" * 60)
    print("  Downloading model checkpoints")
    print("=" * 60)

    download_groundingdino()
    download_sam()

    if INPAINT_METHOD == "lama":
        download_lama()
    else:
        print("\n[3/3] Stable Diffusion -- will auto-download on first run via diffusers")

    print("\n[SUCCESS] All checkpoints ready.\n")


if __name__ == "__main__":
    main()