# ============================================================
# SPID — embeddings & semantic search
#
# Local sentence-transformers (default all-MiniLM-L6-v2) over
# project descriptions; cosine similarity search.  Vectors are
# stored in DuckDB (FLOAT[]) — pgvector/Qdrant are drop-in
# alternatives via the same vector column.
# ============================================================
from __future__ import annotations

import numpy as np

from spid.config import get_logger, get_settings
from spid.store import Store

log = get_logger("embed")

_MODEL = None


def _model(name: str):
    global _MODEL
    if _MODEL is None:
        from sentence_transformers import SentenceTransformer
        log.info("Loading embedding model %s ...", name)
        _MODEL = SentenceTransformer(name)
    return _MODEL


def _doc(row: dict) -> str:
    parts = [row.get("project_name") or "", row.get("project_type") or "",
             row.get("description") or ""]
    techs = row.get("technologies") or []
    if techs:
        parts.append("technologies: " + ", ".join(techs))
    return ". ".join(p for p in parts if p)


def build_embeddings(store: Store, batch: int = 256) -> dict:
    s = get_settings()
    rows = store.con.execute(
        "SELECT project_id, project_name, project_type, description, technologies FROM projects"
    ).to_arrow_table().to_pylist()
    if not rows:
        return {"embedded": 0}
    model = _model(s.embedding_model)
    docs = [_doc(r) for r in rows]
    vecs = model.encode(docs, batch_size=batch, show_progress_bar=False,
                        normalize_embeddings=True)
    dim = int(vecs.shape[1])
    store.con.execute("DELETE FROM project_embeddings")
    payload = [(rows[i]["project_id"], s.embedding_model, dim,
                [float(x) for x in vecs[i]]) for i in range(len(rows))]
    store.con.executemany(
        "INSERT OR REPLACE INTO project_embeddings VALUES (?,?,?,?)", payload)
    log.info("Embedded %d projects (dim=%d, model=%s).", len(rows), dim, s.embedding_model)
    return {"embedded": len(rows), "dim": dim, "model": s.embedding_model}


def semantic_search(store: Store, query: str, k: int = 10,
                    where_sql: str | None = None) -> list[dict]:
    s = get_settings()
    rows = store.con.execute(
        "SELECT project_id, vector FROM project_embeddings"
    ).fetchall()
    if not rows:
        return []
    ids = [r[0] for r in rows]
    mat = np.asarray([r[1] for r in rows], dtype=np.float32)
    qv = _model(s.embedding_model).encode([query], normalize_embeddings=True)[0]
    sims = mat @ qv  # vectors already normalised
    order = np.argsort(-sims)

    # join back to project metadata, applying optional filter
    top_ids = [ids[i] for i in order[: max(k * 5, k)]]
    quoted = ",".join("'" + i + "'" for i in top_ids)
    filt = f" AND {where_sql}" if where_sql else ""
    meta = {r["project_id"]: r for r in store.con.execute(
        f"SELECT * FROM projects WHERE project_id IN ({quoted}){filt}"
    ).to_arrow_table().to_pylist()}
    out = []
    sim_by_id = {ids[i]: float(sims[i]) for i in range(len(ids))}
    for pid in top_ids:
        if pid in meta:
            row = dict(meta[pid]); row["score"] = round(sim_by_id[pid], 4)
            out.append(row)
        if len(out) >= k:
            break
    return out