"""Banco vetorial persistente (ChromaDB) e implementação em memória para testes."""

from __future__ import annotations

import logging
import math
from dataclasses import dataclass
from pathlib import Path
from typing import Iterable, Sequence

from .corpus import VerseChunk
from .embedding import Embedder, HashEmbedder

logger = logging.getLogger(__name__)


@dataclass(frozen=True)
class SearchResult:
    """Trecho recuperado da busca.

    `text` é a janela de contexto que a busca encontrou (e que vai para o LLM);
    `verse_text` é só o versículo citado, que é o que o usuário vê.
    """

    book: str
    chapter: int
    verse: int
    file: str
    text: str
    verse_text: str = ""
    distance: float = 0.0

    @property
    def reference(self) -> str:
        return f"{self.book} {self.chapter}:{self.verse}"

    @property
    def cited_text(self) -> str:
        """Texto a exibir como citação: o versículo, ou a janela se não houver."""
        return self.verse_text or self.text

    @property
    def score(self) -> float:
        """Similaridade em escala 0..1 (1 = idêntico)."""
        return max(0.0, min(1.0, 1.0 - self.distance))


def cosine_similarity(left: Sequence[float], right: Sequence[float]) -> float:
    """Similaridade de cosseno entre dois vetores."""
    if not left or not right:
        return 0.0
    size = min(len(left), len(right))
    dot = sum(left[i] * right[i] for i in range(size))
    left_norm = math.sqrt(sum(value * value for value in left[:size]))
    right_norm = math.sqrt(sum(value * value for value in right[:size]))
    if left_norm == 0 or right_norm == 0:
        return 0.0
    return dot / (left_norm * right_norm)


class VectorStore:
    """Fachada sobre o ChromaDB persistente."""

    def __init__(
        self,
        persist_dir: Path | str,
        collection_name: str = "biblia",
        embedder: Embedder | None = None,
    ):
        import chromadb
        from chromadb.config import Settings as ChromaSettings

        self.persist_dir = Path(persist_dir)
        self.persist_dir.mkdir(parents=True, exist_ok=True)
        self.collection_name = collection_name
        self.embedder = embedder or HashEmbedder()
        self._client = chromadb.PersistentClient(
            path=str(self.persist_dir),
            settings=ChromaSettings(anonymized_telemetry=False, allow_reset=True),
        )
        self._collection = self._client.get_or_create_collection(
            name=collection_name,
            metadata={"hnsw:space": "cosine"},
        )

    def upsert(self, chunks: Iterable[VerseChunk], vectors: Sequence[Sequence[float]]) -> None:
        """Grava chunks e vetores, substituindo os de mesmo id."""
        chunk_list = list(chunks)
        if not chunk_list:
            return
        if len(vectors) != len(chunk_list):
            raise ValueError("vectors e chunks devem ter o mesmo tamanho")

        self._collection.upsert(
            ids=[chunk.chunk_id for chunk in chunk_list],
            embeddings=[list(vector) for vector in vectors],
            # O documento gravado é a janela de contexto: é ela que foi embeddada
            # e é ela que serve de evidência ao LLM. O versículo citado vai no
            # metadado, para a UI mostrar a referência exata.
            documents=[chunk.embedding_text for chunk in chunk_list],
            metadatas=[
                {
                    "book": chunk.book,
                    "chapter": chunk.chapter,
                    "verse": chunk.verse,
                    "file": chunk.file,
                    "verse_text": chunk.text,
                    "reference": f"{chunk.book} {chunk.chapter}:{chunk.verse}",
                }
                for chunk in chunk_list
            ],
        )

    def delete_file(self, file_name: str) -> None:
        """Remove todos os chunks de um arquivo."""
        existing = self._collection.get(where={"file": file_name})
        ids = existing.get("ids") or []
        if ids:
            self._collection.delete(ids=ids)

    def query(
        self, vector: Sequence[float], n_results: int = 6, min_similarity: float = 0.35
    ) -> list[SearchResult]:
        """Busca os chunks mais próximos, descartando os abaixo do limiar."""
        total = self._collection.count()
        if total == 0:
            return []

        raw = self._collection.query(
            query_embeddings=[list(vector)],
            n_results=min(n_results, total),
            include=["documents", "metadatas", "distances"],
        )

        documents = (raw.get("documents") or [[]])[0]
        metadatas = (raw.get("metadatas") or [[]])[0]
        distances = (raw.get("distances") or [[]])[0]

        results: list[SearchResult] = []
        for document, metadata, distance in zip(documents, metadatas, distances):
            similarity = 1.0 - float(distance)
            if similarity < min_similarity:
                continue
            results.append(
                SearchResult(
                    book=str(metadata.get("book", "")),
                    chapter=int(metadata.get("chapter", 0)),
                    verse=int(metadata.get("verse", 0)),
                    file=str(metadata.get("file", "")),
                    text=document or "",
                    verse_text=str(metadata.get("verse_text", "") or ""),
                    distance=float(distance),
                )
            )
        return results

    def get_by_reference(
        self, book: str, chapter: int, verse: int
    ) -> list[SearchResult]:
        """Busca determinística por livro, capítulo e versículo."""
        raw = self._collection.get(
            where={
                "$and": [
                    {"book": book},
                    {"chapter": int(chapter)},
                    {"verse": int(verse)},
                ]
            },
            include=["documents", "metadatas"],
        )
        documents = raw.get("documents") or []
        metadatas = raw.get("metadatas") or []
        return [
            SearchResult(
                book=str(metadata.get("book", book)),
                chapter=int(metadata.get("chapter", chapter)),
                verse=int(metadata.get("verse", verse)),
                file=str(metadata.get("file", "")),
                text=document or "",
                verse_text=str(metadata.get("verse_text", "") or ""),
                distance=0.0,
            )
            for document, metadata in zip(documents, metadatas)
        ]

    def snapshot(self) -> list[SearchResult]:
        """Inventário de todos os trechos indexados, sem pontuação.

        Serve para (re)construir o índice de palavras-chave depois de uma
        reinicialização, quando o Chroma persistiu mas o BM25 em memória não.
        `distance=1.0` deixa `score` em 0: snapshot não é ranking e não pode
        fingir similaridade.
        """
        raw = self._collection.get(include=["documents", "metadatas"])
        documents = raw.get("documents") or []
        metadatas = raw.get("metadatas") or []
        return [
            SearchResult(
                book=str(metadata.get("book", "")),
                chapter=int(metadata.get("chapter", 0)),
                verse=int(metadata.get("verse", 0)),
                file=str(metadata.get("file", "")),
                text=document or "",
                verse_text=str(metadata.get("verse_text", "") or ""),
                distance=1.0,
            )
            for document, metadata in zip(documents, metadatas)
        ]

    def count(self) -> int:
        """Total de chunks indexados."""
        return int(self._collection.count())


class InMemoryVectorStore:
    """Banco vetorial em memória, usado nos testes.

    Implementa a mesma interface de `VectorStore`, com busca por similaridade
    de cosseno e filtros por metadados.
    """

    def __init__(self, embedder: Embedder | None = None):
        self.embedder = embedder or HashEmbedder()
        self._vectors: dict[str, list[float]] = {}
        self._chunks: dict[str, VerseChunk] = {}

    def upsert(self, chunks: Iterable[VerseChunk], vectors: Sequence[Sequence[float]]) -> None:
        chunk_list = list(chunks)
        if len(vectors) != len(chunk_list):
            raise ValueError("vectors e chunks devem ter o mesmo tamanho")
        for chunk, vector in zip(chunk_list, vectors):
            self._vectors[chunk.chunk_id] = list(vector)
            self._chunks[chunk.chunk_id] = chunk

    def delete_file(self, file_name: str) -> None:
        for chunk_id in [cid for cid, chunk in self._chunks.items() if chunk.file == file_name]:
            self._vectors.pop(chunk_id, None)
            self._chunks.pop(chunk_id, None)

    def query(
        self, vector: Sequence[float], n_results: int = 6, min_similarity: float = 0.35
    ) -> list[SearchResult]:
        scored: list[SearchResult] = []
        for chunk_id, stored in self._vectors.items():
            similarity = cosine_similarity(vector, stored)
            if similarity < min_similarity:
                continue
            chunk = self._chunks[chunk_id]
            scored.append(
                SearchResult(
                    book=chunk.book,
                    chapter=chunk.chapter,
                    verse=chunk.verse,
                    file=chunk.file,
                    text=chunk.embedding_text,
                    verse_text=chunk.text,
                    distance=1.0 - similarity,
                )
            )
        scored.sort(key=lambda result: result.distance)
        return scored[:n_results]

    def get_by_reference(
        self, book: str, chapter: int, verse: int
    ) -> list[SearchResult]:
        return [
            SearchResult(
                book=chunk.book,
                chapter=chunk.chapter,
                verse=chunk.verse,
                file=chunk.file,
                text=chunk.embedding_text,
                verse_text=chunk.text,
                distance=0.0,
            )
            for chunk in self._chunks.values()
            if chunk.book == book and chunk.chapter == chapter and chunk.verse == verse
        ]

    def snapshot(self) -> list[SearchResult]:
        """Inventário de todos os trechos indexados, sem pontuação."""
        return [
            SearchResult(
                book=chunk.book,
                chapter=chunk.chapter,
                verse=chunk.verse,
                file=chunk.file,
                text=chunk.embedding_text,
                verse_text=chunk.text,
                distance=1.0,
            )
            for chunk in self._chunks.values()
        ]

    def count(self) -> int:
        return len(self._chunks)