"""Shared singletons: embedding model, re-ranker and the vector DB (all thread-safe)."""
import logging
import threading
from functools import lru_cache
from pathlib import Path

from qdrant_client import QdrantClient
from qdrant_client.models import (
    Distance,
    FieldCondition,
    Filter,
    FilterSelector,
    MatchAny,
    PointStruct,
    VectorParams,
)

from .config import settings

log = logging.getLogger("rag.store")

__all__ = ["PointStruct"]

_qlock = threading.RLock()  # qdrant local mode is not thread-safe, so every call goes through this
_model_lock = threading.Lock()
_embedder = None
_reranker = None
_reranker_failed = False


# ---------------------------------------------------------------- embeddings
def set_embedder(obj) -> None:
    """Swap the embedder (used by tests with a fake one)."""
    global _embedder
    _embedder = obj


def get_embedder():
    global _embedder
    if _embedder is None:
        with _model_lock:
            if _embedder is None:
                from sentence_transformers import SentenceTransformer  # lazy: heavy import

                log.info("Loading embedding model %s", settings.embed_model)
                _embedder = SentenceTransformer(settings.embed_model)
    return _embedder


def embed(texts: list[str]) -> list[list[float]]:
    vecs = get_embedder().encode(texts, normalize_embeddings=True, batch_size=32, show_progress_bar=False)
    return vecs.tolist()


def embedding_dim() -> int:
    return int(get_embedder().get_sentence_embedding_dimension())


# ------------------------------------------------------------------ reranker
def get_reranker():
    """Cross-encoder used to re-rank candidates. Returns None if disabled or unavailable."""
    global _reranker, _reranker_failed
    if not settings.use_reranker or _reranker_failed:
        return None
    if _reranker is None:
        with _model_lock:
            if _reranker is None and not _reranker_failed:
                try:
                    from sentence_transformers import CrossEncoder  # lazy

                    log.info("Loading re-ranker %s", settings.rerank_model)
                    _reranker = CrossEncoder(settings.rerank_model, max_length=512)
                except Exception as e:  # no internet, blocked DLL, etc. -> degrade gracefully
                    log.warning("Re-ranker unavailable, continuing without it: %s", e)
                    _reranker_failed = True
    return _reranker


# -------------------------------------------------------------------- qdrant
def _ensure_collection(client: QdrantClient) -> None:
    existing = {c.name for c in client.get_collections().collections}
    if settings.collection not in existing:
        client.create_collection(
            collection_name=settings.collection,
            vectors_config=VectorParams(size=embedding_dim(), distance=Distance.COSINE),
        )


@lru_cache
def get_qdrant() -> QdrantClient:
    path = Path(settings.data_dir) / "qdrant"
    path.mkdir(parents=True, exist_ok=True)
    with _qlock:
        client = QdrantClient(path=str(path))
        _ensure_collection(client)
    return client


def make_filter(doc_ids=None, pages=None) -> Filter | None:
    must = []
    if doc_ids:
        must.append(FieldCondition(key="doc_id", match=MatchAny(any=list(doc_ids))))
    if pages:
        must.append(FieldCondition(key="page", match=MatchAny(any=[int(p) for p in pages])))
    return Filter(must=must) if must else None


def _hit(p) -> dict:
    return {"id": str(p.id), "score": float(getattr(p, "score", 0.0) or 0.0), **(p.payload or {})}


def upsert(points: list[PointStruct]) -> None:
    with _qlock:
        get_qdrant().upsert(collection_name=settings.collection, points=points)


def search(vector: list[float], flt: Filter | None, limit: int) -> list[dict]:
    with _qlock:
        res = get_qdrant().query_points(
            collection_name=settings.collection,
            query=vector,
            query_filter=flt,
            limit=limit,
            with_payload=True,
        )
    return [_hit(p) for p in res.points]


def scroll_all(flt: Filter | None = None) -> list[dict]:
    """Every stored chunk (payload only) that matches the filter."""
    out: list[dict] = []
    offset = None
    while True:
        with _qlock:
            points, offset = get_qdrant().scroll(
                collection_name=settings.collection,
                scroll_filter=flt,
                limit=512,
                offset=offset,
                with_payload=True,
                with_vectors=False,
            )
        out.extend(_hit(p) for p in points)
        if offset is None:
            break
    return out


def delete_doc(doc_id: str) -> None:
    with _qlock:
        get_qdrant().delete(
            collection_name=settings.collection,
            points_selector=FilterSelector(filter=make_filter([doc_id])),
        )


def clear_all() -> None:
    with _qlock:
        client = get_qdrant()
        client.delete_collection(settings.collection)
        _ensure_collection(client)