app.py

"""hub-search-v3 backend — FastAPI + in-memory DuckDB brute-force vector search.

Drop-in compatible with the old davanstrien/huggingface-datasets-search-v2 API
(same routes / params / response shapes) so the librarian-bots front-end can be
repointed with no code change. Replaces the ChromaDB boot-time HNSW rebuild
(which timed out) with an in-memory DuckDB FLOAT[dim] table loaded once at boot.

See bench: duckdb-vss-bench-2026-07-18. Query encoding replicates the old
backend so query vectors stay aligned with the seed document embeddings.
"""
import logging
import os
import time
from contextlib import asynccontextmanager
from typing import List, Optional

import duckdb
import httpx
from cashews import cache
from fastapi import FastAPI, Header, HTTPException, Query
from fastapi.middleware.cors import CORSMiddleware
from huggingface_hub import HfApi, hf_hub_download, list_repo_files, login
from pydantic import BaseModel

os.environ.setdefault("HF_HUB_ENABLE_HF_TRANSFER", "1")
logging.basicConfig(level=logging.INFO, format="%(asctime)s - %(levelname)s - %(message)s")
logger = logging.getLogger("hub-search-v3")

# --- config -----------------------------------------------------------------
EMBEDDING_MODEL = "Qwen/Qwen3-Embedding-0.6B"
EMBEDDINGS_REPO = os.getenv("EMBEDDINGS_REPO", "davanstrien/search-v2-embeddings")
EMB_DIM = int(os.getenv("EMB_DIM", "256"))          # MRL truncation dim; 1024 = full
FULL_DIM = 1024
THREADS = int(os.getenv("DUCKDB_THREADS", "2"))
CACHE_TTL = "24h"
TRENDING_CACHE_TTL = "1h"

DS_TASK = "Given a search query, retrieve relevant model and dataset summaries that match the query. "

# ISO-639-3/-2B codes that appear in source cards alongside their two-letter
# equivalents (e.g. datasets tagged `fra` vs `fr`); a two-letter filter matches both.
LANG_EQUIV = {
    "en": ["eng"], "fr": ["fra", "fre"], "de": ["deu", "ger"], "es": ["spa"],
    "zh": ["zho", "chi"], "pt": ["por"], "ru": ["rus"], "ja": ["jpn"],
    "ko": ["kor"], "ar": ["ara"], "it": ["ita"], "nl": ["nld", "dut"],
    "pl": ["pol"], "tr": ["tur"], "vi": ["vie"], "hi": ["hin"], "sv": ["swe"],
}

# config dir names inside the embeddings repo (datasets vs models slices)
DATASET_CFG = os.getenv("DATASET_CONFIG", "dataset_cards")
MODEL_CFG = os.getenv("MODEL_CONFIG", "model_cards")

# prebuilt full-card-text FTS database (hybrid search channel); optional at boot
FTS_FILE = os.getenv("FTS_FILE", "_fts/card-text.duckdb")


def config_files(cfg: str) -> List[str]:
    """Discover the parquet shards for a config dir (shard count varies across seed/v3)."""
    files = [f for f in list_repo_files(EMBEDDINGS_REPO, repo_type="dataset")
             if f.startswith(f"{cfg}/") and f.endswith(".parquet")]
    if not files:
        raise RuntimeError(f"no parquet files found under {cfg}/ in {EMBEDDINGS_REPO}")
    return sorted(files)

HF_TOKEN = os.getenv("HF_TOKEN")
if HF_TOKEN:
    login(token=HF_TOKEN)

cache.setup("mem://", size_limit="1gb")

# in-memory DuckDB (single connection, read-only after boot)
con = duckdb.connect(":memory:")
con.execute(f"SET threads={THREADS}")

_query_model = None
STATE = {"ready": False, "boot_seconds": None, "counts": {}, "revision": None,
         "dim": EMB_DIM, "model": EMBEDDING_MODEL}


def get_query_model():
    global _query_model
    if _query_model is None:
        from sentence_transformers import SentenceTransformer
        logger.info(f"Loading query model {EMBEDDING_MODEL} on CPU")
        _query_model = SentenceTransformer(EMBEDDING_MODEL, device="cpu")
    return _query_model


def embed_query(text: str) -> List[float]:
    """Replicate the old backend's query encoding (prompt_name='query')."""
    model = get_query_model()
    vec = model.encode(text, prompt_name="query", normalize_embeddings=False)
    vec = vec[:EMB_DIM]
    # L2 normalise (matches the stored-vector normalisation; makes cosine stable)
    import numpy as np
    n = np.linalg.norm(vec)
    if n > 0:
        vec = vec / n
    return vec.astype(float).tolist()


def load_table(cfg: str, table: str, has_param: bool):
    """Load one config into an in-memory FLOAT[EMB_DIM] table (MRL trunc + renorm).

    Filters NULL embeddings (v3 refusal rows have no embedding). Discovers shards
    dynamically so the seed and the ~1.17M-row search-v3 dataset both work.
    """
    paths = [hf_hub_download(EMBEDDINGS_REPO, f, repo_type="dataset") for f in config_files(cfg)]
    plist = "[" + ",".join(f"'{p}'" for p in paths) + "]"
    param_sel = "COALESCE(param_count,0)::BIGINT AS param_count" if has_param else "0::BIGINT AS param_count"
    # slice to EMB_DIM, L2-renormalise, cast to fixed-size FLOAT[] array (nested, no correlated subquery)
    # task/license/language are copied straight from the source cards; cast to
    # VARCHAR because the models config's all-null `language` reads back as a
    # NULL-typed column that would otherwise reject `=` filters at query time.
    con.execute(f"""
        CREATE OR REPLACE TABLE {table} AS
        SELECT id, summary,
               COALESCE(likes,0)::BIGINT AS likes,
               COALESCE(downloads,0)::BIGINT AS downloads,
               last_modified,
               CAST(task AS VARCHAR) AS task,
               CAST(license AS VARCHAR) AS license,
               CAST(language AS VARCHAR) AS language,
               {param_sel},
               list_transform(s, x -> (x / nrm))::FLOAT[{EMB_DIM}] AS emb
        FROM (
          SELECT *, sqrt(list_dot_product(s, s)) AS nrm
          FROM (
            SELECT *, embedding[1:{EMB_DIM}] AS s
            FROM read_parquet({plist}, union_by_name=True)
            WHERE embedding IS NOT NULL
            QUALIFY row_number() OVER (PARTITION BY id ORDER BY last_modified DESC NULLS LAST) = 1
          )
        )
    """)
    n = con.execute(f"SELECT count(*) FROM {table}").fetchone()[0]
    STATE["counts"][table] = n
    logger.info(f"loaded {table} (from {cfg}): {n:,} rows @ {EMB_DIM}d")


@asynccontextmanager
async def lifespan(app: FastAPI):
    t0 = time.time()
    logger.info(f"boot: loading model + embeddings (dim={EMB_DIM})")
    get_query_model()  # load model up front so first query isn't slow
    try:
        one = hf_hub_download(EMBEDDINGS_REPO, config_files(MODEL_CFG)[0], repo_type="dataset")
        cols = con.execute(f"DESCRIBE SELECT * FROM read_parquet('{one}')").fetchall()
        has_param = any(c[0] == "param_count" for c in cols)
    except Exception as e:
        logger.warning(f"param_count detection failed ({e}); assuming absent")
        has_param = False
    load_table(DATASET_CFG, "dataset_cards", has_param=False)
    load_table(MODEL_CFG, "model_cards", has_param=has_param)
    # Optional: open the prebuilt card-text FTS db on its OWN connection.
    # (Not ATTACHed: the fts match_bm25 macro references its internal schema
    # unqualified, which breaks across catalog boundaries.) Any failure
    # degrades hybrid search to vector-only rather than blocking boot.
    global fts_con
    try:
        fts_path = hf_hub_download(EMBEDDINGS_REPO, FTS_FILE, repo_type="dataset")
        fts_con = duckdb.connect(fts_path, read_only=True)
        fts_con.execute("INSTALL fts; LOAD fts;")
        STATE["fts_cards"] = fts_con.execute("SELECT count(*) FROM cards").fetchone()[0]
        logger.info(f"card-text FTS attached: {STATE['fts_cards']:,} cards")
    except Exception as e:
        fts_con = None
        STATE["fts_cards"] = None
        logger.warning(f"card-text FTS unavailable ({e}); hybrid=vector-only")
    try:
        info = HfApi().dataset_info(EMBEDDINGS_REPO)
        STATE["revision"] = info.sha[:8] if info.sha else None
        STATE["last_modified"] = str(getattr(info, "last_modified", None))
    except Exception:
        pass
    STATE["boot_seconds"] = round(time.time() - t0, 1)
    STATE["ready"] = True
    logger.info(f"boot complete in {STATE['boot_seconds']}s; ready")
    yield
    await cache.close()


app = FastAPI(title="hub-search-v3", lifespan=lifespan)
app.add_middleware(
    CORSMiddleware,
    allow_origin_regex=r"https://.*\.(hf\.space|huggingface\.co)",
    allow_credentials=True,
    allow_methods=["*"],
    allow_headers=["*"],
)


# --- response models (identical to old backend) -----------------------------
class QueryResult(BaseModel):
    dataset_id: str
    similarity: float
    summary: str
    likes: int
    downloads: int
    task: Optional[str] = None
    license: Optional[str] = None
    language: Optional[str] = None
    last_modified: Optional[str] = None


class QueryResponse(BaseModel):
    results: List[QueryResult]


class ModelQueryResult(BaseModel):
    model_id: str
    similarity: float
    summary: str
    likes: int
    downloads: int
    param_count: Optional[int] = None
    task: Optional[str] = None
    license: Optional[str] = None
    language: Optional[str] = None
    last_modified: Optional[str] = None


class ModelQueryResponse(BaseModel):
    results: List[ModelQueryResult]


# --- core search -------------------------------------------------------------
def _where(min_likes, min_downloads, min_param=0, max_param=None, use_param=False,
           task=None, license=None, language=None, modified_after=None):
    """Build a WHERE clause + its bind params.

    Numeric bounds are inlined (ints are injection-safe); the free-text metadata
    filters (task/license/language/modified_after) are bound as `?` params.
    task/language are comma-joined multi-value strings in the source cards, so
    they match by token membership, not string equality; two-letter language
    codes also match their ISO-639-3 equivalents (LANG_EQUIV).
    Returns ``(where_sql, params)`` for ``_knn``.
    """
    conds, params = [], []
    if min_likes > 0:
        conds.append(f"likes >= {int(min_likes)}")
    if min_downloads > 0:
        conds.append(f"downloads >= {int(min_downloads)}")
    if use_param:
        conds.append("param_count > 0")
        if min_param > 0:
            conds.append(f"param_count >= {int(min_param)}")
        if max_param is not None:
            conds.append(f"param_count <= {int(max_param)}")
    token_match = "list_contains(string_split(replace(lower({col}), ' ', ''), ','), ?)"
    if task:
        conds.append(token_match.format(col="task"))
        params.append(task.lower())
    if license:
        conds.append("license = ?")
        params.append(license)
    if language:
        variants = [language.lower()] + LANG_EQUIV.get(language.lower(), [])
        ors = " OR ".join(token_match.format(col="language") for _ in variants)
        conds.append(f"({ors})")
        params.extend(variants)
    if modified_after:
        # last_modified is an ISO-8601 string; lexicographic >= is chronological.
        conds.append("last_modified >= ?")
        params.append(modified_after)
    return (("WHERE " + " AND ".join(conds)) if conds else ""), params


def _vec_literal(vec):
    return "[" + ",".join(repr(float(x)) for x in vec) + f"]::FLOAT[{EMB_DIM}]"


DUP_SIM = 0.985  # near-identical cards (mirrors/re-uploads) collapse to the top-ranked copy


def _knn(table, qvec, k, where_sql, cols, params=None, return_emb=False):
    """Top-k by cosine, with near-duplicate collapse.

    Overfetches 2x, then greedily drops rows whose embedding is ~identical
    (cos >= DUP_SIM) to a higher-ranked kept row — mirror repos otherwise eat
    adjacent result slots. Returns (cols..., sim) tuples; with return_emb=True,
    (cols..., emb, sim).
    """
    import numpy as np
    q = f"""SELECT {cols}, emb, array_cosine_similarity(emb, {_vec_literal(qvec)}) AS sim
            FROM {table} {where_sql} ORDER BY sim DESC LIMIT {int(k) * 2}"""
    rows = con.execute(q, params or []).fetchall()
    kept, kept_vecs = [], []
    for r in rows:
        v = np.asarray(r[-2], dtype=np.float32)
        if kept_vecs and float(np.max(np.stack(kept_vecs) @ v)) >= DUP_SIM:
            continue
        kept.append(r if return_emb else r[:-2] + (r[-1],))
        kept_vecs.append(v)
        if len(kept) >= int(k):
            break
    return kept


def _emb_of(table, repo_id):
    row = con.execute(f"SELECT emb FROM {table} WHERE id = ?", [repo_id]).fetchone()
    return row[0] if row else None


fts_con = None


def _bm25_ids(kind: str, text: str, n: int) -> list:
    """Top-n repo ids by BM25 over full card text (empty if FTS unavailable)."""
    if fts_con is None:
        return []
    rows = fts_con.execute(
        """SELECT id FROM (
               SELECT id, fts_main_cards.match_bm25(doc_id, ?) AS s
               FROM cards WHERE repo_type = ?
           ) WHERE s IS NOT NULL ORDER BY s DESC LIMIT ?""",
        [text, kind, int(n)],
    ).fetchall()
    return [r[0] for r in rows]


def _hybrid_fuse(table, kind, query, qvec, raw, k, where_sql, wparams):
    """RRF-fuse vector KNN rows with a BM25 channel over full card text.

    BM25-only hits are fetched from the in-memory table (with their true cosine
    similarity) and must pass the same metadata filters as the vector channel.
    Returns rows in fused order, same tuple shape as _knn output.
    """
    bm_ids = _bm25_ids(kind, query, k * 2)
    if not bm_ids:
        return raw
    rrf: dict = {}
    for rank, rid in enumerate(r[0] for r in raw):
        rrf[rid] = rrf.get(rid, 0.0) + 1.0 / (60 + rank)
    for rank, rid in enumerate(bm_ids):
        rrf[rid] = rrf.get(rid, 0.0) + 1.0 / (60 + rank)
    order = sorted(rrf, key=rrf.get, reverse=True)
    known = {r[0]: r for r in raw}
    missing = [rid for rid in order if rid not in known][: k * 2]
    if missing:
        ph = ",".join("?" for _ in missing)
        cond = f"AND {where_sql[6:]}" if where_sql else ""
        extra = con.execute(
            f"""SELECT {COLS}, array_cosine_similarity(emb, {_vec_literal(qvec)}) AS sim
                FROM {table} WHERE id IN ({ph}) {cond}""",
            missing + list(wparams or []),
        ).fetchall()
        known.update({r[0]: r for r in extra})
    return [known[rid] for rid in order if rid in known]


async def _sort_results(rows, id_field, sort_by, k):
    """rows: list of dicts with keys id, similarity, summary, likes, downloads, [param_count]."""
    if sort_by == "trending":
        scores = {}
        tasks = [get_trending_score(r["_id"], id_field) for r in rows]
        import asyncio
        vals = await asyncio.gather(*tasks)
        scores = {r["_id"]: v for r, v in zip(rows, vals)}
        rows.sort(key=lambda r: scores.get(r["_id"], 0), reverse=True)
        rows = rows[:k]
    elif sort_by in ("likes", "downloads"):
        rows.sort(key=lambda r: r[sort_by], reverse=True)
        rows = rows[:k]
    elif sort_by == "updated":
        # last_modified is an ISO-8601 string; lexicographic sort is chronological
        rows.sort(key=lambda r: r.get("last_modified") or "", reverse=True)
        rows = rows[:k]
    else:
        rows = rows[:k]
    return rows


def _rows_to_dataset(raw):
    out = []
    for id_, summary, likes, downloads, param, task, license_, language, last_modified, sim in raw:
        out.append({"_id": id_, "dataset_id": id_, "similarity": float(sim),
                    "summary": summary, "likes": int(likes), "downloads": int(downloads),
                    "task": task, "license": license_, "language": language,
                    "last_modified": last_modified})
    return out


def _rows_to_model(raw):
    out = []
    for id_, summary, likes, downloads, param, task, license_, language, last_modified, sim in raw:
        out.append({"_id": id_, "model_id": id_, "similarity": float(sim),
                    "summary": summary, "likes": int(likes), "downloads": int(downloads),
                    "param_count": int(param),
                    "task": task, "license": license_, "language": language,
                    "last_modified": last_modified})
    return out


META_COLS = "task, license, language, CAST(last_modified AS VARCHAR) AS last_modified"
COLS = f"id, summary, likes, downloads, param_count, {META_COLS}"


# --- routes ------------------------------------------------------------------
@app.get("/")
async def root():
    from fastapi.responses import RedirectResponse
    return RedirectResponse(url="/docs")


@app.get("/health")
async def health():
    return {
        "ready": STATE["ready"],
        "boot_seconds": STATE["boot_seconds"],
        "counts": STATE["counts"],
        "index_revision": STATE.get("revision"),
        "index_last_modified": STATE.get("last_modified"),
        "fts_cards": STATE.get("fts_cards"),
        "embedding_model": STATE["model"],
        "dim": STATE["dim"],
        "embeddings_repo": EMBEDDINGS_REPO,
    }


@app.get("/lookup/{repo_id:path}")
async def lookup(repo_id: str):
    """Resolve a repo id to its type in one call: {"id":..., "type": "dataset"|"model"}.

    Saves clients the 404-fallthrough (try datasets, then models). 404 with a JSON
    detail if the id is in neither index. `:path` so `org/name` ids keep their slash.
    """
    for table, typ in (("dataset_cards", "dataset"), ("model_cards", "model")):
        if con.execute(f"SELECT 1 FROM {table} WHERE id = ? LIMIT 1", [repo_id]).fetchone():
            return {"id": repo_id, "type": typ}
    raise HTTPException(status_code=404, detail=f"'{repo_id}' not found as a dataset or model in the index")


@app.get("/search/datasets", response_model=QueryResponse)
@cache(ttl=CACHE_TTL, key="search:d:{query}:{k}:{sort_by}:{min_likes}:{min_downloads}:{task}:{license}:{language}:{modified_after}:{hybrid}")
async def search_datasets(
    query: str,
    k: int = Query(default=5, ge=1, le=100),
    sort_by: str = Query(default="similarity", enum=["similarity", "likes", "downloads", "trending", "updated"]),
    min_likes: int = Query(default=0, ge=0),
    min_downloads: int = Query(default=0, ge=0),
    task: Optional[str] = Query(default=None),
    license: Optional[str] = Query(default=None),
    language: Optional[str] = Query(default=None),
    modified_after: Optional[str] = Query(default=None),
    hybrid: bool = Query(default=False),
):
    try:
        qvec = embed_query(f"Instruct: {DS_TASK}\nQuery:{query}")
        n = k * 4 if sort_by != "similarity" else (k * 2 if hybrid else k)
        where_sql, wparams = _where(min_likes, min_downloads,
                                    task=task, license=license, language=language,
                                    modified_after=modified_after)
        raw = _knn("dataset_cards", qvec, n, where_sql, COLS, wparams)
        if hybrid:
            raw = _hybrid_fuse("dataset_cards", "datasets", query, qvec, raw, k, where_sql, wparams)
        rows = await _sort_results(_rows_to_dataset(raw), "dataset", sort_by, k)
        return QueryResponse(results=[QueryResult(**{k2: v for k2, v in r.items() if k2 != "_id"}) for r in rows])
    except Exception as e:
        logger.error(f"search datasets: {e}")
        raise HTTPException(status_code=500, detail="Search failed")


@app.get("/similarity/datasets", response_model=QueryResponse)
@cache(ttl=CACHE_TTL, key="sim:d:{dataset_id}:{k}:{sort_by}:{min_likes}:{min_downloads}:{task}:{license}:{language}:{modified_after}")
async def find_similar_datasets(
    dataset_id: str,
    k: int = Query(default=5, ge=1, le=100),
    sort_by: str = Query(default="similarity", enum=["similarity", "likes", "downloads", "trending", "updated"]),
    min_likes: int = Query(default=0, ge=0),
    min_downloads: int = Query(default=0, ge=0),
    task: Optional[str] = Query(default=None),
    license: Optional[str] = Query(default=None),
    language: Optional[str] = Query(default=None),
    modified_after: Optional[str] = Query(default=None),
):
    emb = _emb_of("dataset_cards", dataset_id)
    if emb is None:
        raise HTTPException(status_code=404, detail=f"Dataset ID '{dataset_id}' not found")
    try:
        n = k * 4 if sort_by != "similarity" else k + 1
        where_sql, wparams = _where(min_likes, min_downloads,
                                    task=task, license=license, language=language,
                                    modified_after=modified_after)
        raw = _knn("dataset_cards", emb, n, where_sql, COLS, wparams)
        rows = [r for r in _rows_to_dataset(raw) if r["_id"] != dataset_id]
        rows = await _sort_results(rows, "dataset", sort_by, k)
        return QueryResponse(results=[QueryResult(**{k2: v for k2, v in r.items() if k2 != "_id"}) for r in rows])
    except HTTPException:
        raise
    except Exception as e:
        logger.error(f"similarity datasets: {e}")
        raise HTTPException(status_code=500, detail="Similarity search failed")


@app.get("/search/models", response_model=ModelQueryResponse)
@cache(ttl=CACHE_TTL, key="search:m:{query}:{k}:{sort_by}:{min_likes}:{min_downloads}:{min_param_count}:{max_param_count}:{task}:{license}:{language}:{modified_after}:{hybrid}")
async def search_models(
    query: str,
    k: int = Query(default=5, ge=1, le=100),
    sort_by: str = Query(default="similarity", enum=["similarity", "likes", "downloads", "trending", "updated"]),
    min_likes: int = Query(default=0, ge=0),
    min_downloads: int = Query(default=0, ge=0),
    min_param_count: int = Query(default=0, ge=0),
    max_param_count: Optional[int] = Query(default=None, ge=0),
    task: Optional[str] = Query(default=None),
    license: Optional[str] = Query(default=None),
    language: Optional[str] = Query(default=None),
    modified_after: Optional[str] = Query(default=None),
    hybrid: bool = Query(default=False),
):
    try:
        use_param = min_param_count > 0 or max_param_count is not None
        qvec = embed_query(f"search_query: {query}")
        n = k * 4 if sort_by != "similarity" else (k * 2 if hybrid else k)
        where_sql, wparams = _where(min_likes, min_downloads, min_param_count, max_param_count, use_param,
                                    task=task, license=license, language=language,
                                    modified_after=modified_after)
        raw = _knn("model_cards", qvec, n, where_sql, COLS, wparams)
        if hybrid:
            raw = _hybrid_fuse("model_cards", "models", query, qvec, raw, k, where_sql, wparams)
        rows = await _sort_results(_rows_to_model(raw), "model", sort_by, k)
        return ModelQueryResponse(results=[ModelQueryResult(**{k2: v for k2, v in r.items() if k2 != "_id"}) for r in rows])
    except Exception as e:
        logger.error(f"search models: {e}")
        raise HTTPException(status_code=500, detail="Model search failed")


@app.get("/similarity/models", response_model=ModelQueryResponse)
@cache(ttl=CACHE_TTL, key="sim:m:{model_id}:{k}:{sort_by}:{min_likes}:{min_downloads}:{min_param_count}:{max_param_count}:{task}:{license}:{language}:{modified_after}")
async def find_similar_models(
    model_id: str,
    k: int = Query(default=5, ge=1, le=100),
    sort_by: str = Query(default="similarity", enum=["similarity", "likes", "downloads", "trending", "updated"]),
    min_likes: int = Query(default=0, ge=0),
    min_downloads: int = Query(default=0, ge=0),
    min_param_count: int = Query(default=0, ge=0),
    max_param_count: Optional[int] = Query(default=None, ge=0),
    task: Optional[str] = Query(default=None),
    license: Optional[str] = Query(default=None),
    language: Optional[str] = Query(default=None),
    modified_after: Optional[str] = Query(default=None),
):
    emb = _emb_of("model_cards", model_id)
    if emb is None:
        raise HTTPException(status_code=404, detail=f"Model ID '{model_id}' not found")
    try:
        use_param = min_param_count > 0 or max_param_count is not None
        n = k * 4 if sort_by != "similarity" else k + 1
        where_sql, wparams = _where(min_likes, min_downloads, min_param_count, max_param_count, use_param,
                                    task=task, license=license, language=language,
                                    modified_after=modified_after)
        raw = _knn("model_cards", emb, n, where_sql, COLS, wparams)
        rows = [r for r in _rows_to_model(raw) if r["_id"] != model_id]
        rows = await _sort_results(rows, "model", sort_by, k)
        return ModelQueryResponse(results=[ModelQueryResult(**{k2: v for k2, v in r.items() if k2 != "_id"}) for r in rows])
    except HTTPException:
        raise
    except Exception as e:
        logger.error(f"similarity models: {e}")
        raise HTTPException(status_code=500, detail="Model similarity search failed")


# --- trending (proxies HF API, joins summaries from the in-memory table) -----
@cache(ttl=TRENDING_CACHE_TTL, key="tscore:{item_type}:{item_id}")
async def get_trending_score(item_id: str, item_type: str) -> float:
    try:
        endpoint = "models" if item_type == "model" else "datasets"
        async with httpx.AsyncClient(timeout=10) as c:
            r = await c.get(f"https://huggingface.co/api/{endpoint}/{item_id}?expand=trendingScore")
            r.raise_for_status()
            return r.json().get("trendingScore", 0)
    except Exception:
        return 0


@app.get("/trending/datasets", response_model=QueryResponse)
@cache(ttl=TRENDING_CACHE_TTL, key="trend:d:{limit}:{min_likes}:{min_downloads}")
async def trending_datasets(
    limit: int = Query(default=10, ge=1, le=100),
    min_likes: int = Query(default=0, ge=0),
    min_downloads: int = Query(default=0, ge=0),
):
    try:
        async with httpx.AsyncClient(timeout=15) as c:
            r = await c.get("https://huggingface.co/api/datasets?sort=trendingScore&limit=200")
            r.raise_for_status()
            items = r.json()
    except Exception as e:
        logger.error(f"trending datasets fetch: {e}")
        raise HTTPException(status_code=502, detail="Failed to fetch trending datasets")
    items = [d for d in items if d.get("likes", 0) >= min_likes and d.get("downloads", 0) >= min_downloads]
    ids = [d["id"] for d in items[: limit * 3]]
    if not ids:
        return QueryResponse(results=[])
    placeholders = ",".join("?" for _ in ids)
    found = {r[0]: r for r in con.execute(
        f"SELECT id, summary, likes, downloads, {META_COLS} FROM dataset_cards WHERE id IN ({placeholders})", ids
    ).fetchall()}
    out = []
    for d in items:
        row = found.get(d["id"])
        if row:
            out.append(QueryResult(dataset_id=d["id"], similarity=1.0, summary=row[1],
                                   likes=d.get("likes", 0), downloads=d.get("downloads", 0),
                                   task=row[4], license=row[5], language=row[6],
                                   last_modified=row[7]))
        if len(out) >= limit:
            break
    return QueryResponse(results=out)


@app.get("/trending/models", response_model=ModelQueryResponse)
@cache(ttl=TRENDING_CACHE_TTL, key="trend:m:{limit}:{min_likes}:{min_downloads}:{min_param_count}:{max_param_count}")
async def trending_models(
    limit: int = Query(default=10, ge=1, le=100),
    min_likes: int = Query(default=0, ge=0),
    min_downloads: int = Query(default=0, ge=0),
    min_param_count: int = Query(default=0, ge=0),
    max_param_count: Optional[int] = Query(default=None, ge=0),
):
    try:
        async with httpx.AsyncClient(timeout=15) as c:
            r = await c.get("https://huggingface.co/api/models?sort=trendingScore&limit=200")
            r.raise_for_status()
            items = r.json()
    except Exception as e:
        logger.error(f"trending models fetch: {e}")
        raise HTTPException(status_code=502, detail="Failed to fetch trending models")
    items = [m for m in items if m.get("likes", 0) >= min_likes and m.get("downloads", 0) >= min_downloads]
    ids = [m["id"] for m in items[: limit * 3]]
    if not ids:
        return ModelQueryResponse(results=[])
    placeholders = ",".join("?" for _ in ids)
    found = {r[0]: r for r in con.execute(
        f"SELECT id, summary, likes, downloads, param_count, {META_COLS} FROM model_cards WHERE id IN ({placeholders})", ids
    ).fetchall()}
    use_param = min_param_count > 0 or max_param_count is not None
    out = []
    for m in items:
        row = found.get(m["id"])
        if not row:
            continue
        param = int(row[4])
        if use_param:
            if param == 0:
                continue
            if min_param_count > 0 and param < min_param_count:
                continue
            if max_param_count is not None and param > max_param_count:
                continue
        out.append(ModelQueryResult(model_id=m["id"], similarity=1.0, summary=row[1],
                                    likes=m.get("likes", 0), downloads=m.get("downloads", 0),
                                    param_count=param,
                                    task=row[5], license=row[6], language=row[7],
                                    last_modified=row[8]))
        if len(out) >= limit:
            break
    return ModelQueryResponse(results=out)


# --- embedding primitives (personal-feeds prototype + compare tooling) -------
class EmbBatchReq(BaseModel):
    repo_type: str
    ids: List[str]


@app.post("/embeddings_batch")
async def embeddings_batch(req: EmbBatchReq):
    """Return stored embeddings for up to 200 repo ids (e.g. a user's likes)."""
    if req.repo_type not in ("datasets", "models"):
        raise HTTPException(status_code=422, detail="repo_type must be datasets|models")
    table = "dataset_cards" if req.repo_type == "datasets" else "model_cards"
    ids = req.ids[:200]
    if not ids:
        return {"embeddings": {}}
    placeholders = ",".join("?" for _ in ids)
    rows = con.execute(
        f"SELECT id, emb FROM {table} WHERE id IN ({placeholders})", ids
    ).fetchall()
    return {"embeddings": {r[0]: [float(x) for x in r[1]] for r in rows}}


@app.get("/embed_query")
@cache(ttl=CACHE_TTL, key="embq:{repo_type}:{text}")
async def embed_query_route(
    text: str = Query(min_length=3, max_length=200),
    repo_type: str = Query(default="datasets"),
):
    """Embed a free-text topic phrase with the serving query encoder.

    repo_type picks the query prefix the corpus was embedded against
    (datasets use the instruct prompt, models the search_query prefix).
    """
    if repo_type not in ("datasets", "models"):
        raise HTTPException(status_code=422, detail="repo_type must be datasets|models")
    prompt = (f"Instruct: {DS_TASK}\nQuery:{text}" if repo_type == "datasets"
              else f"search_query: {text}")
    return {"embedding": embed_query(prompt)}


class VectorSearchReq(BaseModel):
    repo_type: str
    vector: List[float]
    k: int = 20
    exclude_ids: List[str] = []
    include_embeddings: bool = False


@app.post("/search_vector")
async def search_vector(req: VectorSearchReq):
    """KNN search with a caller-supplied (already 256-d, normalised) vector.

    Powers relevance-feedback loops: refine a preference vector client-side
    from votes, then pull fresh results for it. exclude_ids drops already-seen
    items server-side so refinement surfaces new material.
    """
    if req.repo_type not in ("datasets", "models"):
        raise HTTPException(status_code=422, detail="repo_type must be datasets|models")
    if len(req.vector) != EMB_DIM:
        raise HTTPException(status_code=422, detail=f"vector must be {EMB_DIM}-d")
    table = "dataset_cards" if req.repo_type == "datasets" else "model_cards"
    k = max(1, min(int(req.k), 100))
    n = k + len(req.exclude_ids[:300])
    raw = _knn(table, req.vector, n, "", COLS, return_emb=req.include_embeddings)
    excl = set(req.exclude_ids[:300])
    out = []
    for row in raw:
        emb_val = row[-2] if req.include_embeddings else None
        (id_, summary, likes, downloads, param, task, license_, language, last_modified) = row[:9]
        if id_ in excl:
            continue
        item = {"id": id_, "summary": summary, "likes": int(likes),
                "downloads": int(downloads), "param_count": int(param),
                "task": task, "license": license_, "language": language,
                "last_modified": last_modified, "similarity": float(row[-1])}
        if req.include_embeddings:
            item["embedding"] = [float(x) for x in emb_val]
        out.append(item)
        if len(out) >= k:
            break
    return {"results": out}


# --- per-user feeds store (OAuth token -> whoami -> one JSON doc per user) ---
FEEDS_STORE_REPO = os.getenv("FEEDS_STORE_REPO", "davanstrien/hub-feeds-store")
FEEDS_DOC_LIMIT = 300_000  # bytes


@cache(ttl="10m", key="whoami:{token}")
async def _resolve_user(token: str) -> str:
    async with httpx.AsyncClient(timeout=10) as c:
        r = await c.get("https://huggingface.co/api/whoami-v2",
                        headers={"Authorization": f"Bearer {token}"})
    if r.status_code != 200:
        raise HTTPException(status_code=401, detail="Invalid or expired HF token")
    name = r.json().get("name")
    if not name:
        raise HTTPException(status_code=401, detail="Could not resolve username")
    return name


async def _auth_user(authorization: str) -> str:
    if not authorization or not authorization.startswith("Bearer "):
        raise HTTPException(status_code=401, detail="Missing Bearer token")
    return await _resolve_user(authorization[7:])


@app.get("/feeds_store")
async def feeds_load(authorization: str = Header(default=None)):
    user = await _auth_user(authorization)
    import json as _json
    try:
        path = hf_hub_download(FEEDS_STORE_REPO, f"users/{user}.json",
                               repo_type="dataset", force_download=True)
        with open(path) as f:
            return {"user": user, "doc": _json.load(f)}
    except Exception:
        return {"user": user, "doc": None}


class FeedsSaveReq(BaseModel):
    doc: dict


@app.put("/feeds_store")
async def feeds_save(req: FeedsSaveReq, authorization: str = Header(default=None)):
    user = await _auth_user(authorization)
    import json as _json
    payload = _json.dumps(req.doc).encode()
    if len(payload) > FEEDS_DOC_LIMIT:
        raise HTTPException(status_code=413, detail="Feeds doc too large")
    HfApi().upload_file(path_or_fileobj=payload, path_in_repo=f"users/{user}.json",
                        repo_id=FEEDS_STORE_REPO, repo_type="dataset",
                        commit_message=f"feeds sync: {user}")
    return {"ok": True, "user": user, "bytes": len(payload)}


# --- suggest (autocomplete; old backend never implemented it — bonus) --------
@app.get("/suggest/{repo_type}")
@cache(ttl="1h", key="suggest:{repo_type}:{q}")
async def suggest(repo_type: str, q: str):
    if len(q) < 2:
        return {"suggestions": []}
    table = "dataset_cards" if repo_type == "datasets" else "model_cards"
    rows = con.execute(
        f"SELECT id FROM {table} WHERE id ILIKE ? ORDER BY downloads DESC LIMIT 10",
        [f"%{q}%"],
    ).fetchall()
    return {"suggestions": [r[0] for r in rows]}


if __name__ == "__main__":
    import uvicorn
    uvicorn.run(app, host="0.0.0.0", port=7860)