# -*- coding: utf-8 -*-
"""
BM25 检索引擎（借鉴 zhikuncode-main MemorySearchEngine.java）

纯 Python 实现，无 Lucene/Jieba 依赖。
- K1=1.2, B=0.75, TITLE_BOOST=2.0
- CJK 分词：Unigram（每字符）+ Bigram（相邻字符对）
- 英文分词：regex `[a-zA-Z0-9_]+`
- 中英文停用词过滤
"""
import math
import re
from collections import Counter
from dataclasses import dataclass
from typing import Any, Dict, List, Optional, Sequence

# BM25 参数（与 zhikuncode MemorySearchEngine 一致）
K1 = 1.2
B = 0.75
TITLE_BOOST = 2.0

# 英文停用词（高频无意义词）
_ENGLISH_STOPWORDS = frozenset({
    "the", "a", "an", "and", "or", "but", "is", "are", "was", "were",
    "be", "been", "being", "have", "has", "had", "do", "does", "did",
    "of", "in", "on", "at", "to", "for", "with", "by", "from", "as",
    "this", "that", "these", "those", "it", "its", "i", "you", "he",
    "she", "we", "they", "my", "your", "his", "her", "our", "their",
})

# 中文停用词（高频虚词）
_CJK_STOPWORDS = frozenset({
    "的", "了", "在", "是", "我", "你", "他", "她", "它", "们",
    "和", "与", "或", "但", "而", "也", "都", "就", "还", "又",
    "把", "被", "让", "使", "对", "为", "以", "于", "之", "其",
    "这", "那", "些", "个", "中", "上", "下", "里", "外", "前",
    "后", "不", "没", "无", "有", "一", "二", "三",
})

# 英文 token 正则
_ENGLISH_TOKEN_RE = re.compile(r"[a-zA-Z0-9_]+")


@dataclass
class ScoredResult:
    """BM25 检索结果：索引 + 分数"""
    index: int
    score: float


def _is_cjk(ch: str) -> bool:
    """判断字符是否为 CJK（中日韩）字符"""
    if not ch:
        return False
    cp = ord(ch)
    # CJK 统一表意文字 + 扩展 A + 日文假名 + 韩文
    return (
        0x4E00 <= cp <= 0x9FFF        # CJK Unified Ideographs
        or 0x3400 <= cp <= 0x4DBF     # CJK Extension A
        or 0x3040 <= cp <= 0x30FF    # Hiragana + Katakana
        or 0xAC00 <= cp <= 0xD7AF    # Hangul Syllables
    )


def tokenize(text: str) -> List[str]:
    """分词：CJK 用 Unigram + Bigram；英文用 regex。

    与 zhikuncode MemorySearchEngine.tokenize 一致：
    - CJK 字符：每个字符作 Unigram，相邻两字符作 Bigram
    - 英文：连续 [a-zA-Z0-9_] 作为一个 token
    - 过滤停用词（仅对英文 token 和 CJK Unigram 过滤，Bigram 保留以保留语序信息）
    """
    if not text:
        return []
    tokens: List[str] = []
    cjk_chars: List[str] = []

    def flush_cjk():
        """把累积的 CJK 字符生成为 Unigram + Bigram"""
        nonlocal cjk_chars
        if not cjk_chars:
            return
        # Unigram（过滤停用词）
        for ch in cjk_chars:
            if ch not in _CJK_STOPWORDS:
                tokens.append(ch)
        # Bigram（保留语序信息，不过滤停用词以保留语义）
        for i in range(len(cjk_chars) - 1):
            tokens.append(cjk_chars[i] + cjk_chars[i + 1])
        cjk_chars = []

    for m in _ENGLISH_TOKEN_RE.finditer(text):
        # 先 flush 之前累积的 CJK
        start = m.start()
        if cjk_chars:
            flush_cjk()
        word = m.group(0).lower()
        if word not in _ENGLISH_STOPWORDS and len(word) > 1:
            tokens.append(word)
        elif word not in _ENGLISH_STOPWORDS:
            # 单字符英文 token 也保留（如变量名 i, j）
            tokens.append(word)

    # 处理文本中未被英文 token 正则匹配的部分（CJK 字符）
    # 重新扫描：遍历原 text，找出非英文 token 的 CJK 字符
    last_pos = 0
    for m in _ENGLISH_TOKEN_RE.finditer(text):
        # 处理 last_pos 到 m.start() 之间的 CJK
        gap = text[last_pos:m.start()]
        for ch in gap:
            if _is_cjk(ch):
                cjk_chars.append(ch)
            else:
                flush_cjk()
        last_pos = m.end()
    # 处理尾部
    tail = text[last_pos:]
    for ch in tail:
        if _is_cjk(ch):
            cjk_chars.append(ch)
        else:
            flush_cjk()
    flush_cjk()

    return tokens


class BM25Engine:
    """BM25 检索引擎。

    使用方式：
        engine = BM25Engine()
        engine.index([entry1, entry2, ...])  # entry 为含 text/title 字段的 dict 或 str
        results = engine.search("查询文本", top_k=5)
    """

    def __init__(self, k1: float = K1, b: float = B, title_boost: float = TITLE_BOOST):
        self.k1 = k1
        self.b = b
        self.title_boost = title_boost
        # 文档数据
        self._docs: List[Dict[str, Any]] = []
        # 每个文档的 token 列表（body 部分）
        self._doc_tokens: List[List[str]] = []
        # 每个文档的 title token 列表
        self._title_tokens: List[List[str]] = []
        # 文档频率：token -> 出现该 token 的文档数
        self._df: Counter = Counter()
        # 平均文档长度（body）
        self._avgdl: float = 0.0
        # 平均 title 长度（使用原始值，不除以 5）
        self._avg_title_dl: float = 0.0

    def index(self, entries: Sequence[Any]) -> None:
        """建立索引。entries 中每项可以是：
        - str：纯 body 文本
        - dict：可含 title 与 content/body/text 字段
        """
        self._docs = []
        self._doc_tokens = []
        self._title_tokens = []
        self._df = Counter()
        total_len = 0
        total_title_len = 0

        for entry in entries:
            if isinstance(entry, str):
                body = entry
                title = ""
            elif isinstance(entry, dict):
                body = (
                    entry.get("content")
                    or entry.get("body")
                    or entry.get("text")
                    or ""
                )
                title = entry.get("title") or ""
            else:
                # 尝试 getattr
                body = getattr(entry, "content", "") or getattr(entry, "text", "") or ""
                title = getattr(entry, "title", "") or ""

            body_tokens = tokenize(body) if body else []
            title_tokens_list = tokenize(title) if title else []

            self._docs.append({"title": title, "body": body})
            self._doc_tokens.append(body_tokens)
            self._title_tokens.append(title_tokens_list)
            total_len += len(body_tokens)
            total_title_len += len(title_tokens_list)

            # 更新 df（文档级去重：token 出现在 body 或 title 中即计 1 次，避免 df 超过 N）
            unique_tokens = set(body_tokens) | set(title_tokens_list)
            for t in unique_tokens:
                self._df[t] += 1

        n = len(self._docs)
        self._avgdl = (total_len / n) if n > 0 else 0.0
        # title avgdl 使用原始值（不除以 5，避免过度惩罚长 title）
        self._avg_title_dl = (total_title_len / n) if n > 0 else 0.0

    def _idf(self, df: int) -> float:
        """计算 IDF：log((N - df + 0.5) / (df + 0.5) + 1.0)"""
        n = len(self._docs)
        if n == 0:
            return 0.0
        return math.log((n - df + 0.5) / (df + 0.5) + 1.0)

    def _bm25_score(self, query_tokens: List[str], doc_tokens: List[str],
                    avgdl: float) -> float:
        """计算单个文档的 BM25 分数"""
        if not doc_tokens or not query_tokens:
            return 0.0
        tf_map = Counter(doc_tokens)
        dl = len(doc_tokens)
        denom_avgdl = avgdl if avgdl > 0 else 1.0
        score = 0.0
        for q in query_tokens:
            tf = tf_map.get(q, 0)
            if tf == 0:
                continue
            df = self._df.get(q, 0)
            if df == 0:
                continue
            idf = self._idf(df)
            # BM25 公式
            norm = 1 - self.b + self.b * (dl / denom_avgdl)
            score += idf * (tf * (self.k1 + 1)) / (tf + self.k1 * norm)
        return score

    def search(self, query: str, top_k: int = 5) -> List[ScoredResult]:
        """检索：返回 top_k 个结果，按分数降序"""
        if not self._docs or not query:
            return []
        query_tokens = tokenize(query)
        if not query_tokens:
            return []

        scored: List[ScoredResult] = []
        for i in range(len(self._docs)):
            body_score = self._bm25_score(query_tokens, self._doc_tokens[i], self._avgdl)
            title_score = self._bm25_score(
                query_tokens, self._title_tokens[i], self._avg_title_dl
            )
            # title 加权（与 zhikuncode 一致）
            total = body_score + title_score * self.title_boost
            scored.append(ScoredResult(index=i, score=total))

        scored.sort(key=lambda r: r.score, reverse=True)
        # 过滤掉 0 分结果
        return [r for r in scored[:top_k] if r.score > 0]

    @property
    def doc_count(self) -> int:
        """已索引的文档数"""
        return len(self._docs)
