from fastapi import FastAPI
from pydantic import BaseModel
from sentence_transformers import SentenceTransformer, util

app = FastAPI(title="MovieNotes Ranker (HF Space)")

# Small, fast embedding model
model = SentenceTransformer("sentence-transformers/all-MiniLM-L6-v2")

class RankIn(BaseModel):
    q: str
    catalogue: list[str]
    limit: int = 5

class RankOut(BaseModel):
    indices: list[int]
    scores: list[float]

@app.get("/")
def root():
    return {"status": "ok", "endpoints": ["/rank"]}

@app.post("/rank", response_model=RankOut)
def rank(req: RankIn):
    q = req.q or ""
    k = max(1, min(req.limit, 10))
    q_emb = model.encode(q, normalize_embeddings=True)
    cat_emb = model.encode(req.catalogue, normalize_embeddings=True)
    scores = util.cos_sim(q_emb, cat_emb).tolist()[0]
    order = sorted(range(len(scores)), key=lambda i: -scores[i])[:k]
    return {"indices": order, "scores": [scores[i] for i in order]}