"""Vector store abstraction with an in-memory tenant-scoped default.

Stores are addressed by name; the in-memory store keeps tenant + collection
isolation so cross-tenant leaks are impossible by construction.
"""

from __future__ import annotations

from collections.abc import Iterable
from dataclasses import dataclass, field
from typing import Any, Protocol

from simorgh.apps.ai.embeddings import cosine_similarity


class VectorStoreError(RuntimeError):
    """Raised on registration or storage errors."""


@dataclass
class VectorRecord:
    id: str
    embedding: list[float]
    metadata: dict[str, Any] = field(default_factory=dict)


@dataclass
class VectorMatch:
    id: str
    score: float
    metadata: dict[str, Any]


class VectorStore(Protocol):
    name: str

    def upsert(
        self,
        *,
        tenant_id: int,
        collection: str,
        records: Iterable[VectorRecord],
    ) -> int: ...

    def query(
        self,
        *,
        tenant_id: int,
        collection: str,
        embedding: list[float],
        top_k: int = 5,
    ) -> list[VectorMatch]: ...

    def delete(self, *, tenant_id: int, collection: str, ids: Iterable[str]) -> int: ...


class InMemoryVectorStore:
    """Process-local store — sufficient for tests and small workloads."""

    name = "memory"

    def __init__(self) -> None:
        # (tenant_id, collection) -> {id: VectorRecord}
        self._data: dict[tuple[int, str], dict[str, VectorRecord]] = {}

    def _bucket(self, tenant_id: int, collection: str) -> dict[str, VectorRecord]:
        return self._data.setdefault((tenant_id, collection), {})

    def upsert(
        self,
        *,
        tenant_id: int,
        collection: str,
        records: Iterable[VectorRecord],
    ) -> int:
        bucket = self._bucket(tenant_id, collection)
        count = 0
        for rec in records:
            bucket[rec.id] = rec
            count += 1
        return count

    def query(
        self,
        *,
        tenant_id: int,
        collection: str,
        embedding: list[float],
        top_k: int = 5,
    ) -> list[VectorMatch]:
        if top_k <= 0:
            raise VectorStoreError("top_k must be >= 1")
        bucket = self._data.get((tenant_id, collection), {})
        scored = [
            VectorMatch(
                id=rec.id,
                score=cosine_similarity(embedding, rec.embedding),
                metadata=dict(rec.metadata),
            )
            for rec in bucket.values()
        ]
        scored.sort(key=lambda m: m.score, reverse=True)
        return scored[:top_k]

    def delete(self, *, tenant_id: int, collection: str, ids: Iterable[str]) -> int:
        bucket = self._data.get((tenant_id, collection))
        if bucket is None:
            return 0
        count = 0
        for vid in ids:
            if bucket.pop(vid, None) is not None:
                count += 1
        return count

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


_STORES: dict[str, VectorStore] = {}


def register_store(store: VectorStore) -> VectorStore:
    name = getattr(store, "name", "")
    if not name:
        raise VectorStoreError("store must expose a non-empty .name")
    _STORES[name] = store
    return store


def get_store(name: str = "memory") -> VectorStore:
    try:
        return _STORES[name]
    except KeyError as exc:
        raise VectorStoreError(f"unknown vector store {name!r}") from exc


def list_stores() -> list[str]:
    return sorted(_STORES)


def reset_for_tests() -> None:
    _STORES.clear()
    register_default_stores()


def register_default_stores() -> None:
    if "memory" not in _STORES:
        register_store(InMemoryVectorStore())
    else:
        # Wipe accumulated state across tests.
        existing = _STORES["memory"]
        if isinstance(existing, InMemoryVectorStore):
            existing.reset()