"""Dense retrieval over a shard index (BGE-M3 by default; any model in embed_models.py).

    from search import Index
    idx = Index()                              # BGE-M3, data/index
    idx = Index(model="sealion")               # SEA-LION-E5, data/index_sealion (metadata shared from data/index)
    hits = idx.search("ကျောင်းဆရာ လစာ", k=5)      # -> list of (score, meta)

Shards: emb_NNNNN.npy (float16, L2-normalised) per model directory; meta_NNNNN.jsonl lives once, in data/index,
because every model embeds the same passages in the same order (make_embed_input.py fixed the order).
"""
import glob, json, os
import numpy as np
from common import DATA, INDEX_DIR, load_embedder, header_of, citation_of
from embed_models import spec

def load_meta(meta_dir=INDEX_DIR):
    metas = []
    for mp in sorted(glob.glob(os.path.join(meta_dir, "meta_*.jsonl"))):
        metas.extend(json.loads(l) for l in open(mp, encoding="utf-8"))
    return metas

class Index:
    def __init__(self, model="bge", index_dir=None, meta_dir=INDEX_DIR, device=None, embedder=None, meta=None):
        self.key = model
        self.spec = spec(model)
        self.index_dir = index_dir or os.path.join(DATA, self.spec["dir"])
        paths = sorted(glob.glob(os.path.join(self.index_dir, "emb_*.npy")))
        if not paths: raise RuntimeError(f"no shards found in {self.index_dir} — build the '{model}' index first")
        self.n_shards = len(paths)
        self.emb = np.concatenate([np.load(p) for p in paths]).astype(np.float32)   # [N, dim]
        self.meta = meta if meta is not None else load_meta(meta_dir)[: len(self.emb)]
        if len(self.meta) != len(self.emb):
            raise RuntimeError(f"{len(self.emb):,} embeddings but {len(self.meta):,} metadata rows — index incomplete?")
        self.model = embedder or load_embedder(device, model=model)[0]

    def __len__(self): return len(self.meta)

    def encode_query(self, q):
        q = self.spec["query"] + q
        return self.model.encode([q], normalize_embeddings=True, convert_to_numpy=True)[0].astype(np.float32)

    def scores(self, query):
        """Cosine score of every passage for the query (used by hybrid fusion)."""
        return self.emb @ self.encode_query(query)

    def search(self, query, k=5, filter_fn=None, per_speech=None):
        """Top-k passages. per_speech=N caps how many passages of the same speech may appear (pilot: top-k
        was filled by one long speech)."""
        sc = self.scores(query)
        order = np.argsort(-sc)
        out, seen = [], {}
        for j in order:
            m = self.meta[j]
            if filter_fn and not filter_fn(m): continue
            if per_speech:
                c = seen.get(m["speech_id"], 0)
                if c >= per_speech: continue
                seen[m["speech_id"]] = c + 1
            out.append((float(sc[j]), m))
            if len(out) >= k: break
        return out

def format_hit(rank, score, m, n_chars=220):
    return (f"{rank}. [{score:.3f}] {citation_of(m)}\n"
            f"   {m['text'][:n_chars].replace(chr(10), ' ')}{'…' if len(m['text']) > n_chars else ''}")