"""Search an EmbeddingGemma 2 image index with A1 queries.

Build an index from a folder of images with EmbeddingGemma 2 (skip this if you already have one):
    python example_search.py build --images ./images --index index.npz
Search it with Korean or English queries encoded by A1:
    python example_search.py search --index index.npz --query "월별 매출 추이" --query "login screen"
"""
import argparse
from pathlib import Path

import numpy as np

from a1_query import A1QueryEncoder

EXTS = {".jpg", ".jpeg", ".png", ".webp", ".bmp"}


def build(images, index, eg2="google/embeddinggemma-2", batch=16):
    from PIL import Image
    from sentence_transformers import SentenceTransformer
    paths = sorted(p for p in Path(images).rglob("*") if p.suffix.lower() in EXTS)
    st = SentenceTransformer(eg2, device="cpu")
    vecs = []
    for i in range(0, len(paths), batch):
        ims = [{"image": Image.open(p).convert("RGB")} for p in paths[i:i + batch]]
        vecs.append(st.encode(ims, batch_size=batch, convert_to_numpy=True, normalize_embeddings=True))
    np.savez(index, vectors=np.concatenate(vecs).astype(np.float32), paths=np.array([str(p) for p in paths]))
    print(f"indexed {len(paths)} images -> {index}")


def search(index, queries, k=5, model="ThakiCloud/eg2-tiny-query-a1"):
    d = np.load(index)
    enc = A1QueryEncoder.from_pretrained(model)
    q = enc.encode(queries)
    scores = q @ d["vectors"].T
    for text, s in zip(queries, scores):
        print(f"\n{text}")
        for j in np.argsort(-s)[:k]:
            print(f"  {s[j]:.4f}  {d['paths'][j]}")
    return scores


if __name__ == "__main__":
    ap = argparse.ArgumentParser()
    sub = ap.add_subparsers(dest="cmd", required=True)
    b = sub.add_parser("build"); b.add_argument("--images", required=True); b.add_argument("--index", required=True)
    b.add_argument("--eg2", default="google/embeddinggemma-2")
    s = sub.add_parser("search"); s.add_argument("--index", required=True); s.add_argument("--query", action="append", required=True)
    s.add_argument("--k", type=int, default=5); s.add_argument("--model", default="ThakiCloud/eg2-tiny-query-a1")
    a = ap.parse_args()
    build(a.images, a.index, a.eg2) if a.cmd == "build" else search(a.index, a.query, a.k, a.model)