"""
BM25 倒排索引搜索引擎
用于图片 OCR 内容的全文搜索
参考 Meilisearch / Whoosh 设计
"""

from __future__ import annotations

import re
import math
import logging
from collections import Counter, defaultdict
from typing import List, Dict, Optional, Tuple

logger = logging.getLogger(__name__)

# 尝试导入 jieba，失败则用简单的字符分割
try:
    import jieba
    jieba.setLogLevel(logging.WARNING)
    def tokenize(text: str) -> List[str]:
        """中文分词"""
        return [w for w in jieba.cut_for_search(text) if len(w.strip()) > 1]
except ImportError:
    logger.warning("jieba 未安装，使用简单分词（pip install jieba 获得更好效果）")
    def tokenize(text: str) -> List[str]:
        """简单分词：按空格和标点分割"""
        return [w for w in re.split(r'[\s,.\-;:!?"\'/\\|()\[\]{}<>，。；：！？""''、（）【】]+', text) if len(w.strip()) > 1]


class BM25Index:
    """
    BM25 倒排索引
    - 支持增量添加文档
    - 支持删除文档
    - 支持中文分词（jieba）
    - BM25 评分排序
    """

    def __init__(self, k1: float = 1.5, b: float = 0.75):
        """
        Args:
            k1: 词频饱和参数（默认 1.5）
            b: 文档长度归一化参数（默认 0.75）
        """
        self.k1 = k1
        self.b = b

        # 倒排索引：term -> {doc_id: tf}
        self.inverted_index: Dict[str, Dict[int, int]] = defaultdict(lambda: defaultdict(int))

        # 文档信息
        self.doc_lengths: Dict[int, int] = {}  # doc_id -> 文档长度（词数）
        self.doc_count: int = 0
        self.avg_doc_length: float = 0.0

        # 文档存储（存完整数据）
        self.documents: Dict[int, dict] = {}

        # IDF 缓存
        self._idf_cache: Dict[str, float] = {}

    def add_document(self, doc_id: int, text: str, metadata: Optional[dict] = None):
        """
        添加文档到索引

        Args:
            doc_id: 文档 ID
            text: 待索引的文本内容
            metadata: 附加元数据（如 file_path, tags 等）
        """
        # 如果已存在，先删除
        if doc_id in self.documents:
            self.remove_document(doc_id)

        tokens = tokenize(text)
        self.doc_lengths[doc_id] = len(tokens)
        self.documents[doc_id] = {
            "text": text,
            "tokens": tokens,
            "metadata": metadata or {},
        }

        # 构建倒排索引
        tf = Counter(tokens)
        for term, count in tf.items():
            self.inverted_index[term][doc_id] = count

        self.doc_count += 1
        self._update_avg_doc_length()
        self._idf_cache.clear()  # 文档数变了，IDF 需要重算

    def remove_document(self, doc_id: int):
        """从索引中删除文档"""
        if doc_id not in self.documents:
            return

        # 从倒排索引中移除
        tokens = self.documents[doc_id].get("tokens", [])
        for term in set(tokens):
            if term in self.inverted_index and doc_id in self.inverted_index[term]:
                del self.inverted_index[term][doc_id]
                if not self.inverted_index[term]:
                    del self.inverted_index[term]

        del self.documents[doc_id]
        del self.doc_lengths[doc_id]
        self.doc_count -= 1
        self._update_avg_doc_length()
        self._idf_cache.clear()

    def _update_avg_doc_length(self):
        """更新平均文档长度"""
        if self.doc_count > 0:
            self.avg_doc_length = sum(self.doc_lengths.values()) / self.doc_count
        else:
            self.avg_doc_length = 0.0

    def _idf(self, term: str) -> float:
        """
        计算逆文档频率 IDF
        IDF(t) = ln((N - df(t) + 0.5) / (df(t) + 0.5) + 1)
        其中 N 是文档总数，df(t) 是包含 term 的文档数
        """
        if term in self._idf_cache:
            return self._idf_cache[term]

        df = len(self.inverted_index.get(term, {}))
        idf = math.log((self.doc_count - df + 0.5) / (df + 0.5) + 1)

        self._idf_cache[term] = idf
        return idf

    def search(self, query: str, top_k: int = 10) -> List[Tuple[int, float]]:
        """
        BM25 搜索

        Args:
            query: 搜索查询（会自动分词）
            top_k: 返回前 K 个结果

        Returns:
            [(doc_id, score), ...] 按 BM25 分数降序排列
        """
        if self.doc_count == 0 or self.avg_doc_length == 0:
            return []

        query_tokens = tokenize(query)
        if not query_tokens:
            return []

        scores: Dict[int, float] = defaultdict(float)

        for term in query_tokens:
            idf = self._idf(term)
            postings = self.inverted_index.get(term, {})

            for doc_id, tf in postings.items():
                dl = self.doc_lengths[doc_id]
                # BM25 公式
                numerator = tf * (self.k1 + 1)
                denominator = tf + self.k1 * (1 - self.b + self.b * dl / self.avg_doc_length)
                scores[doc_id] += idf * numerator / denominator

        # 按分数降序排序
        ranked = sorted(scores.items(), key=lambda x: x[1], reverse=True)
        return ranked[:top_k]

    def get_document(self, doc_id: int) -> Optional[dict]:
        """获取文档完整数据"""
        return self.documents.get(doc_id)

    def clear(self):
        """清空索引"""
        self.inverted_index.clear()
        self.doc_lengths.clear()
        self.documents.clear()
        self._idf_cache.clear()
        self.doc_count = 0
        self.avg_doc_length = 0.0

    @property
    def size(self) -> int:
        """索引中的文档数量"""
        return self.doc_count