"""Product-image search with the commerce query encoder, alone or fused with BM25.

1) Build an EmbeddingGemma 2 index from a folder of product images (skip if you already have one):
    python example_search.py build --images ./products --index index.npz
2) Search it with Korean or English queries:
    python example_search.py search --index index.npz --query "무선 블루투스 이어폰" --query "stainless steel water bottle"
3) Optional: fuse with a BM25 keyword index over your own product titles (one JSON line per image:
   {"path": "<same path as in the index>", "text": "<title, brand, attributes>"}):
    python example_search.py search --index index.npz --catalog catalog.jsonl --query "나이키 러닝화 270mm"

Fusion is reciprocal-rank fusion (k=60) of the two top-10 lists. In our tests it raised R@10 on some
query sets but lowered R@1 whenever BM25 alone was already strong; check it on your own data.
"""
import argparse
import json
import math
import re
from collections import Counter
from pathlib import Path

import numpy as np

from commerce_query import CommerceQueryEncoder

EXTS = {".jpg", ".jpeg", ".png", ".webp", ".bmp"}
REPO = "ThakiCloud/eg2-tiny-query-commerce-ko"


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 tokens(s):
    return re.findall(r"\w+", (s or "").lower())


class BM25:
    """Okapi BM25 (k1=1.5, b=0.75), whitespace/word tokens. Small and dependency-free."""

    def __init__(self, docs, k1=1.5, b=0.75):
        self.docs = [Counter(tokens(d)) for d in docs]
        self.len = np.array([sum(d.values()) for d in self.docs], float)
        self.avg = float(self.len.mean()) if len(self.docs) else 0.0
        df = Counter(t for d in self.docs for t in d)
        n = len(self.docs)
        self.idf = {t: math.log((n - f + 0.5) / (f + 0.5) + 1) for t, f in df.items()}
        self.k1, self.b = k1, b

    def scores(self, query):
        s = np.zeros(len(self.docs))
        for t in set(tokens(query)):
            if t not in self.idf:
                continue
            tf = np.array([d.get(t, 0) for d in self.docs], float)
            s += self.idf[t] * tf * (self.k1 + 1) / (tf + self.k1 * (1 - self.b + self.b * self.len / max(self.avg, 1e-9)))
        return s


def rrf(*ranked, k=60, top=10):
    sc = {}
    for r in ranked:
        for pos, d in enumerate(r):
            sc[int(d)] = sc.get(int(d), 0.0) + 1.0 / (k + pos + 1)
    return [d for d, _ in sorted(sc.items(), key=lambda x: -x[1])[:top]]


def search(index, queries, k=5, model=REPO, catalog=None, route="hangul"):
    d = np.load(index)
    paths = [str(p) for p in d["paths"]]
    enc = CommerceQueryEncoder.from_pretrained(model, route=route)
    q = enc.encode(queries)
    dense = np.argsort(-(q @ d["vectors"].T), axis=1)[:, :10]
    bm = None
    if catalog:
        text = {}
        for line in open(catalog, encoding="utf-8"):
            r = json.loads(line)
            text[str(r["path"])] = r.get("text", "")
        bm = BM25([text.get(p, "") for p in paths])
    results = []
    for i, t in enumerate(queries):
        top = list(dense[i])
        if bm is not None:
            top = rrf(np.argsort(-bm.scores(t))[:10], top)
        results.append(top)
        print(f"\n{t}" + ("  [model + BM25]" if bm is not None else "  [model]"))
        for j in top[:k]:
            print(f"  {paths[j]}")
    return results


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=REPO)
    s.add_argument("--catalog"); s.add_argument("--route", default="hangul", choices=["hangul", "ko", "en"])
    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, a.catalog, a.route)