"""
build_index.py — embed the knowledge base for retrieval.

Reads  data/kb.json  (produced by build_kb.py)
Writes data/index.npz — normalized embedding matrix + aligned ids

Needs GEMINI_API_KEY (in .env). Run AFTER build_kb.py:
    python3 scripts/build_kb.py
    python3 scripts/build_index.py
"""
import json
import os
import numpy as np
from dotenv import load_dotenv
from google import genai
from google.genai import types

ROOT = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
load_dotenv(os.path.join(ROOT, ".env"))

KB = os.path.join(ROOT, "data", "kb.json")
OUT = os.path.join(ROOT, "data", "index.npz")
EMBED_MODEL = os.getenv("GEMINI_EMBED_MODEL", "gemini-embedding-001")
DIM = 1536  # truncated Matryoshka dim — plenty for ~500 entries, cheaper to store

client = genai.Client(api_key=os.environ["GEMINI_API_KEY"])


def embed_batch(texts, task_type):
    resp = client.models.embed_content(
        model=EMBED_MODEL,
        contents=texts,
        config=types.EmbedContentConfig(task_type=task_type, output_dimensionality=DIM),
    )
    return [e.values for e in resp.embeddings]


def normalize(mat):
    mat = np.asarray(mat, dtype=np.float32)
    norms = np.linalg.norm(mat, axis=1, keepdims=True)
    norms[norms == 0] = 1.0
    return mat / norms


if __name__ == "__main__":
    kb = json.load(open(KB))
    # exclude rows the AI must not send verbatim (coach-only scripts, placeholder/merge-field rows)
    skipped = [e for e in kb if e.get("coach_only") or e.get("has_placeholder")]
    kb = [e for e in kb if not e.get("coach_only") and not e.get("has_placeholder")]
    print(f"excluded {len(skipped)} coach-only/placeholder rows from the answer index")
    # retrieval key = question + tags + topic (what a user query should match against)
    texts = [
        " ".join(filter(None, [e["question"], e.get("tags", ""), e.get("topic", "")]))
        for e in kb
    ]

    vectors = []
    for i in range(0, len(texts), 100):
        chunk = texts[i:i + 100]
        vectors.extend(embed_batch(chunk, "RETRIEVAL_DOCUMENT"))
        print(f"embedded {min(i + 100, len(texts))}/{len(texts)}")

    mat = normalize(vectors)
    ids = np.array([e["id"] for e in kb], dtype=np.int32)
    np.savez(OUT, vectors=mat, ids=ids)
    print(f"Wrote {mat.shape[0]} vectors (dim {mat.shape[1]}) to {OUT}")