# retoor <retoor@molodetz.nl>

from __future__ import annotations

import asyncio
import logging
import math
import re
from collections import Counter
from dataclasses import dataclass, field
from urllib.parse import urlparse

from devplacepy.config import DEEPSEARCH_CHROMA_DIR

logger = logging.getLogger(__name__)

TOKEN_PATTERN = re.compile(r"[a-z0-9]+")
BM25_K1 = 1.5
BM25_B = 0.75
RRF_K = 60.0
DEFAULT_TOP_K = 8
CANDIDATE_MULTIPLIER = 4
CHROMADB_TIMEOUT = 30.0


@dataclass
class Chunk:
    uid: str
    text: str
    url: str
    title: str
    depth: int = 0
    source: str = ""
    position: int = 0
    score: float = 0.0
    metadata: dict = field(default_factory=dict)
    embedding: list[float] | None = None


def _tokenize(text: str) -> list[str]:
    return TOKEN_PATTERN.findall((text or "").lower())


class VectorStore:
    def __init__(self, collection_name: str) -> None:
        self.collection_name = collection_name
        self._client = None
        self._collection = None
        self._dims: int | None = None

    async def _run_sync(self, func, *args, timeout: float = CHROMADB_TIMEOUT):
        try:
            return await asyncio.wait_for(
                asyncio.to_thread(func, *args), timeout=timeout
            )
        except asyncio.TimeoutError:
            logger.error(
                "deepsearch ChromaDB operation timed out after %.1fs on collection %s",
                timeout,
                self.collection_name,
            )
            raise

    def _ensure(self):
        if self._collection is not None:
            return self._collection
        import chromadb

        DEEPSEARCH_CHROMA_DIR.mkdir(parents=True, exist_ok=True)
        self._client = chromadb.PersistentClient(path=str(DEEPSEARCH_CHROMA_DIR))
        self._collection = self._client.get_or_create_collection(
            name=self.collection_name, metadata={"hnsw:space": "cosine"}
        )
        return self._collection

    @property
    def dims(self) -> int | None:
        if self._dims is not None:
            return self._dims
        try:
            collection = self._ensure()
            data = collection.get(include=["embeddings"], limit=1)
            rows = data.get("embeddings") or []
            if rows and rows[0]:
                self._dims = len(rows[0])
        except Exception:
            return self._dims
        return self._dims

    def _add_sync(self, chunks: list[Chunk], vectors: list[list[float]]) -> None:
        if not chunks:
            return
        keep_chunks: list[Chunk] = []
        keep_vectors: list[list[float]] = []
        for chunk, vector in zip(chunks, vectors):
            if not vector:
                continue
            if self._dims is None:
                self._dims = len(vector)
            if len(vector) != self._dims:
                logger.warning(
                    "deepsearch dropping chunk with mismatched embedding dim %d != %d",
                    len(vector),
                    self._dims,
                )
                continue
            keep_chunks.append(chunk)
            keep_vectors.append(vector)
        if not keep_chunks:
            return
        chunks = keep_chunks
        vectors = keep_vectors
        collection = self._ensure()
        collection.add(
            ids=[chunk.uid for chunk in chunks],
            embeddings=vectors,
            documents=[chunk.text for chunk in chunks],
            metadatas=[
                {
                    "url": chunk.url,
                    "title": chunk.title,
                    "depth": chunk.depth,
                    "source": chunk.source,
                    "position": chunk.position,
                }
                for chunk in chunks
            ],
        )

    async def add(self, chunks: list[Chunk], vectors: list[list[float]]) -> None:
        if not chunks:
            return
        await self._run_sync(self._add_sync, chunks, vectors)

    def _all_chunks_sync(self, limit: int = 0) -> list[Chunk]:
        collection = self._ensure()
        kwargs: dict = {"include": ["documents", "metadatas"]}
        if limit > 0:
            kwargs["limit"] = limit
        data = collection.get(**kwargs)
        chunks: list[Chunk] = []
        ids = data.get("ids") or []
        documents = data.get("documents") or []
        metadatas = data.get("metadatas") or []
        for index, uid in enumerate(ids):
            meta = metadatas[index] if index < len(metadatas) else {}
            chunks.append(
                Chunk(
                    uid=uid,
                    text=documents[index] if index < len(documents) else "",
                    url=str(meta.get("url", "")),
                    title=str(meta.get("title", "")),
                    depth=int(meta.get("depth", 0) or 0),
                    source=str(meta.get("source", "")),
                    position=int(meta.get("position", 0) or 0),
                    metadata=dict(meta),
                )
            )
        return chunks

    async def all_chunks(self, limit: int = 1000) -> list[Chunk]:
        return await self._run_sync(self._all_chunks_sync, limit)

    async def count(self) -> int:
        try:
            collection = self._ensure()
            return await self._run_sync(collection.count)
        except Exception:
            return 0

    def _vector_search_sync(
        self, query_vector: list[float], top_k: int, where: dict | None = None
    ) -> list[Chunk]:
        collection = self._ensure()
        result = collection.query(
            query_embeddings=[query_vector],
            n_results=top_k,
            where=where or None,
            include=["documents", "metadatas", "distances", "embeddings"],
        )
        ids = (result.get("ids") or [[]])[0]
        documents = (result.get("documents") or [[]])[0]
        metadatas = (result.get("metadatas") or [[]])[0]
        distances = (result.get("distances") or [[]])[0]
        embeddings = (result.get("embeddings") or [[]])[0]
        chunks: list[Chunk] = []
        for index, uid in enumerate(ids):
            meta = metadatas[index] if index < len(metadatas) else {}
            distance = distances[index] if index < len(distances) else 1.0
            vector = embeddings[index] if index < len(embeddings) else None
            chunks.append(
                Chunk(
                    uid=uid,
                    text=documents[index] if index < len(documents) else "",
                    url=str(meta.get("url", "")),
                    title=str(meta.get("title", "")),
                    depth=int(meta.get("depth", 0) or 0),
                    source=str(meta.get("source", "")),
                    position=int(meta.get("position", 0) or 0),
                    score=1.0 - float(distance),
                    metadata=dict(meta),
                    embedding=[float(v) for v in vector] if vector is not None else None,
                )
            )
        return chunks

    async def vector_search(
        self, query_vector: list[float], top_k: int, where: dict | None = None
    ) -> list[Chunk]:
        return await self._run_sync(
            self._vector_search_sync, query_vector, top_k, where
        )

    def keyword_scores(self, query: str, chunks: list[Chunk]) -> dict[str, float]:
        terms = _tokenize(query)
        if not terms or not chunks:
            return {}
        docs = [_tokenize(chunk.text) for chunk in chunks]
        lengths = [len(doc) for doc in docs]
        avg_len = (sum(lengths) / len(lengths)) if lengths else 0.0
        doc_freq: Counter = Counter()
        for doc in docs:
            for term in set(doc):
                if term in terms:
                    doc_freq[term] += 1
        total_docs = len(docs)
        scores: dict[str, float] = {}
        for index, chunk in enumerate(chunks):
            counts = Counter(docs[index])
            length = lengths[index] or 1
            score = 0.0
            for term in terms:
                freq = counts.get(term, 0)
                if freq == 0:
                    continue
                idf = math.log(
                    1 + (total_docs - doc_freq[term] + 0.5) / (doc_freq[term] + 0.5)
                )
                denom = freq + BM25_K1 * (
                    1 - BM25_B + BM25_B * (length / (avg_len or 1))
                )
                score += idf * (freq * (BM25_K1 + 1)) / (denom or 1)
            scores[chunk.uid] = score
        return scores

    async def hybrid_search(
        self,
        query: str,
        query_vector: list[float],
        top_k: int = DEFAULT_TOP_K,
        where: dict | None = None,
    ) -> list[Chunk]:
        candidates = await self.vector_search(
            query_vector, top_k * CANDIDATE_MULTIPLIER, where
        )
        if not candidates:
            return []
        vector_rank = sorted(
            candidates, key=lambda chunk: chunk.score, reverse=True
        )
        keyword = self.keyword_scores(query, candidates)
        keyword_rank = sorted(
            candidates, key=lambda chunk: keyword.get(chunk.uid, 0.0), reverse=True
        )
        fused: dict[str, float] = {}
        for rank, chunk in enumerate(vector_rank, start=1):
            fused[chunk.uid] = fused.get(chunk.uid, 0.0) + 1.0 / (RRF_K + rank)
        for rank, chunk in enumerate(keyword_rank, start=1):
            fused[chunk.uid] = fused.get(chunk.uid, 0.0) + 1.0 / (RRF_K + rank)
        for chunk in candidates:
            chunk.score = fused.get(chunk.uid, 0.0)
        candidates.sort(key=lambda chunk: chunk.score, reverse=True)
        return candidates[:top_k]

    async def coverage_analytics(self, sample_limit: int = 5000) -> dict:
        chunks = await self.all_chunks(limit=sample_limit)
        if not chunks:
            total = await self.count()
            return {"chunks": total, "domains": 0, "sources": 0, "avg_chunk_chars": 0}
        domains = {
            urlparse(chunk.metadata.get("url", "")).netloc for chunk in chunks
        }
        domains.discard("")
        sources = {chunk.source for chunk in chunks}
        sources.discard("")
        avg_chars = int(sum(len(chunk.text) for chunk in chunks) / len(chunks))
        return {
            "chunks": len(chunks),
            "domains": len(domains),
            "sources": len(sources),
            "avg_chunk_chars": avg_chars,
        }

    def drop(self) -> None:
        try:
            import chromadb

            DEEPSEARCH_CHROMA_DIR.mkdir(parents=True, exist_ok=True)
            client = self._client or chromadb.PersistentClient(
                path=str(DEEPSEARCH_CHROMA_DIR)
            )
            client.delete_collection(self.collection_name)
        except Exception as exc:
            logger.info("deepsearch collection drop skipped for %s: %s", self.collection_name, exc)
        finally:
            self._collection = None
