"""
Generate CLIP image embeddings for a folder of reference images.

These embeddings are then used to train a STANNO as a style autoencoder,
which can be loaded into the ComfyUI STANNODreamCond or STANNODynamicLoRA
nodes for conditioning/weight-patch injection.

Usage:
    python scripts/generate_clip_embeddings.py \
        --dir  my_style_images/ \
        --out  style_embeddings.npy \
        [--model ViT-L-14]  [--pretrained openai]

Requirements:
    pip install open-clip-torch Pillow

Outputs:
    A .npy file of shape (N, 768) — one 768-dim CLIP embedding per image.
    Compatible with SD 1.5 CLIP-L text encoder embedding space.
"""

from __future__ import annotations
import argparse
import sys
from pathlib import Path

import numpy as np


def parse_args() -> argparse.Namespace:
    p = argparse.ArgumentParser(description="Generate CLIP embeddings for a folder of images")
    p.add_argument("--dir", required=True, help="Folder of input images (png, jpg, webp)")
    p.add_argument("--out", required=True, help="Output .npy path")
    p.add_argument("--model", default="ViT-L-14", help="OpenCLIP model name")
    p.add_argument("--pretrained", default="openai", help="OpenCLIP pretrained weights")
    p.add_argument("--batch", type=int, default=16, help="Batch size for encoding")
    p.add_argument("--device", default="cuda", help="Device: cuda or cpu")
    return p.parse_args()


def main() -> None:
    args = parse_args()

    try:
        import torch
        import open_clip
        from PIL import Image
    except ImportError as e:
        print(f"Missing dependency: {e}")
        print("Install with: pip install open-clip-torch Pillow")
        sys.exit(1)

    image_paths = sorted(
        p for ext in ("*.png", "*.jpg", "*.jpeg", "*.webp")
        for p in Path(args.dir).glob(ext)
    )
    if not image_paths:
        print(f"No images found in {args.dir}")
        sys.exit(1)

    print(f"Found {len(image_paths)} images in {args.dir}")

    model, _, preprocess = open_clip.create_model_and_transforms(
        args.model, pretrained=args.pretrained
    )
    model.eval().to(args.device)

    all_embeddings: list[np.ndarray] = []

    for i in range(0, len(image_paths), args.batch):
        batch_paths = image_paths[i : i + args.batch]
        imgs = torch.stack(
            [preprocess(Image.open(str(p)).convert("RGB")) for p in batch_paths]
        ).to(args.device)

        with torch.no_grad():
            feats = model.encode_image(imgs)

        all_embeddings.append(feats.cpu().numpy())
        print(f"  Encoded {min(i + args.batch, len(image_paths))}/{len(image_paths)}")

    embeddings = np.concatenate(all_embeddings, axis=0).astype(np.float32)
    out_path = Path(args.out)
    out_path.parent.mkdir(parents=True, exist_ok=True)
    np.save(str(out_path), embeddings)

    print(f"\nSaved {embeddings.shape} embeddings → {out_path}")
    print(f"Use this file with train_stanno_on_embeddings.py or STANNOTrainImages (ComfyUI).")


if __name__ == "__main__":
    main()