"""语义检索模块."""
import logging
from src.core.db import VectorDB
from src.core.embedder import Embedder

logger = logging.getLogger("md-vector-db")


class Searcher:
    """向量检索器."""

    def __init__(self, db: VectorDB, embedder: Embedder, collection_name: str):
        self.db = db
        self.embedder = embedder
        self.collection_name = collection_name

    @property
    def collection(self):
        return self.db.get_or_create_collection(self.collection_name)

    def search(
        self,
        query: str,
        top_k: int = 10,
        source_file: str | None = None,
    ) -> list[dict]:
        """语义检索, 返回格式化结果列表."""
        query_embedding = self.embedder.embed([query])[0]

        where_filter = None
        if source_file:
            where_filter = {"source_file": source_file}

        results = self.collection.query(
            query_embeddings=[query_embedding],
            n_results=top_k,
            where=where_filter,
            include=["documents", "metadatas", "distances"],
        )

        formatted = []
        if results["ids"] and results["ids"][0]:
            for i, doc_id in enumerate(results["ids"][0]):
                metadata = results["metadatas"][0][i] if results["metadatas"] else {}
                distance = results["distances"][0][i] if results["distances"] else 0.0
                # ChromaDB 默认用余弦距离, 转为相似度分数 (0~1)
                score = round(1.0 - distance, 4)

                formatted.append({
                    "id": doc_id,
                    "content": results["documents"][0][i] if results["documents"] else "",
                    "source_file": metadata.get("source_file", ""),
                    "section_title": metadata.get("section_title", ""),
                    "heading_level": metadata.get("heading_level", 0),
                    "chunk_index": metadata.get("chunk_index", 0),
                    "score": max(0.0, score),
                })

        return formatted

    def get_collection_info(self) -> dict:
        """获取 collection 信息."""
        return {
            "name": self.collection_name,
            "count": self.collection.count(),
        }

    def list_sources(self) -> list[str]:
        """列出所有已入库的源文件."""
        if self.collection.count() == 0:
            return []
        result = self.collection.get(include=["metadatas"])
        sources = set()
        if result and result["metadatas"]:
            for m in result["metadatas"]:
                if m and "source_file" in m:
                    sources.add(m["source_file"])
        return sorted(sources)

    def delete_by_source(self, file_name: str) -> bool:
        """按文件名删除文档 (线程安全)."""
        try:
            with self.db.write_lock:
                existing = self.collection.get(
                    where={"source_file": file_name}
                )
                if existing and existing["ids"]:
                    self.collection.delete(ids=existing["ids"])
                    return True
        except ValueError:
            pass  # collection 为空时 ChromaDB 抛 ValueError
        except Exception:
            logger.exception("删除文档失败: %s", file_name)
        return False