# backend/app/rag/store.py
import math

from sqlmodel import Session, select

from ..models import RagChunk


def cosine(a: list[float], b: list[float]) -> float:
    if not a or not b:
        return 0.0
    dot = sum(x * y for x, y in zip(a, b, strict=False))
    na = math.sqrt(sum(x * x for x in a)) or 1.0
    nb = math.sqrt(sum(y * y for y in b)) or 1.0
    return dot / (na * nb)


def add_chunks(s: Session, chunks: list[RagChunk]) -> None:
    for c in chunks:
        s.add(c)
    s.commit()


def query_store(s: Session, qvec: list[float], k: int, source_ids: list[str] | None):
    rows = s.exec(select(RagChunk)).all()
    if source_ids:
        rows = [r for r in rows if r.source_id in source_ids]
    scored = [(cosine(qvec, r.embedding), r) for r in rows]
    scored.sort(key=lambda x: x[0], reverse=True)
    return scored[:k]