"""In-memory vector store for local development and tests.

Cosine-similarity search over vectors held in process memory. Same interface as
the Qdrant store so the rest of the system is identical either way.
"""

from __future__ import annotations

import math

from legate.rag.vectorstore.base import SearchHit, VectorRecord


def _cosine(a: list[float], b: list[float]) -> float:
    if not a or not b or len(a) != len(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))
    nb = math.sqrt(sum(y * y for y in b))
    if na == 0.0 or nb == 0.0:
        return 0.0
    return dot / (na * nb)


class InMemoryVectorStore:
    def __init__(self) -> None:
        # collection -> {id: (vector, payload)}
        self._data: dict[str, dict[str, tuple[list[float], dict]]] = {}

    def upsert(self, collection: str, records: list[VectorRecord]) -> None:
        store = self._data.setdefault(collection, {})
        for rec in records:
            store[rec.id] = (rec.vector, rec.payload)

    def search(self, collection: str, vector: list[float], top_k: int = 5) -> list[SearchHit]:
        store = self._data.get(collection, {})
        hits = [
            SearchHit(id=pid, score=_cosine(vector, vec), payload=payload)
            for pid, (vec, payload) in store.items()
        ]
        hits.sort(key=lambda h: h.score, reverse=True)
        return hits[:top_k]

    def delete_by_document(self, collection: str, document_id: str) -> int:
        store = self._data.get(collection)
        if not store:
            return 0
        to_remove = [
            pid for pid, (_, payload) in store.items() if payload.get("document_id") == document_id
        ]
        for pid in to_remove:
            del store[pid]
        return len(to_remove)

    def delete_collection(self, collection: str) -> None:
        self._data.pop(collection, None)

    def reset(self) -> None:
        self._data.clear()
