"""知识检索中的关键词召回与多路结果融合。

第二阶段同时使用两类检索：

1. Milvus 负责语义检索，适合找到“说法不同但含义相近”的内容；
2. 本模块的 BM25 负责关键词检索，适合保险名称、责任条款等精确文字；
3. 两路结果通过 RRF 融合，避免只依赖单一检索方式。

本模块只负责“找候选内容和排序”，不负责让大模型生成答案。
"""

import math
import re
from collections import Counter
from collections.abc import Iterable
from dataclasses import replace
from typing import Protocol

from zbt.domains.knowledge.models import KnowledgeChunk, KnowledgeSearchHit

ASCII_TERM_PATTERN = re.compile(r"[a-z0-9]+")
CHINESE_RUN_PATTERN = re.compile(r"[\u4e00-\u9fff]+")


class KnowledgeReranker(Protocol):
    """语义重排器的最小接口。

    具体实现可以使用本地 BGE Reranker，也可以替换成其他重排服务。
    """

    def rerank(
        self,
        query: str,
        hits: list[KnowledgeSearchHit],
        *,
        limit: int,
    ) -> list[KnowledgeSearchHit]: ...


def rank_bm25(
    query: str,
    chunks: Iterable[KnowledgeChunk],
    *,
    limit: int,
    k1: float = 1.5,
    b: float = 0.75,
) -> list[KnowledgeSearchHit]:
    """按照关键词相关度对知识分块排序。

    ``chunks`` 是已经切分好的知识片段；返回值统一转换成
    ``KnowledgeSearchHit``，方便后续与 Milvus 的语义结果合并。
    """

    candidates = list(chunks)
    if not candidates or not query.strip():
        return []
    tokenized = [_tokenize(chunk.content) for chunk in candidates]
    query_terms = set(_tokenize(query))
    if not query_terms:
        return []
    document_frequency = Counter(
        term
        for terms in tokenized
        for term in query_terms.intersection(terms)
    )
    average_length = sum(len(terms) for terms in tokenized) / len(tokenized)
    ranked: list[KnowledgeSearchHit] = []
    for chunk, terms in zip(candidates, tokenized, strict=True):
        frequencies = Counter(terms)
        score = 0.0
        for term in query_terms:
            frequency = frequencies[term]
            if frequency == 0:
                continue
            idf = math.log(
                1
                + (
                    len(candidates)
                    - document_frequency[term]
                    + 0.5
                )
                / (document_frequency[term] + 0.5)
            )
            length_normalizer = k1 * (
                1 - b + b * len(terms) / max(average_length, 1)
            )
            score += idf * frequency * (k1 + 1) / (frequency + length_normalizer)
        if score > 0:
            ranked.append(_chunk_to_hit(chunk, score))
    return sorted(ranked, key=lambda hit: hit.score, reverse=True)[:limit]


def reciprocal_rank_fusion(
    rankings: Iterable[list[KnowledgeSearchHit]],
    *,
    limit: int,
    rank_constant: int = 60,
) -> list[KnowledgeSearchHit]:
    """使用 RRF 合并多份排行榜。

    RRF 关注同一片段在各排行榜中的名次，不直接比较不同检索器的原始分数，
    因而适合融合 BM25 分数和向量相似度这两种量纲不同的结果。
    """

    hits_by_id: dict[str, KnowledgeSearchHit] = {}
    scores: dict[str, float] = {}
    for ranking in rankings:
        for rank, hit in enumerate(ranking, start=1):
            hits_by_id.setdefault(hit.chunk_id, hit)
            scores[hit.chunk_id] = scores.get(hit.chunk_id, 0.0) + 1 / (
                rank_constant + rank
            )
    ordered_ids = sorted(scores, key=scores.__getitem__, reverse=True)[:limit]
    return [
        replace(hits_by_id[chunk_id], score=scores[chunk_id])
        for chunk_id in ordered_ids
    ]


def _tokenize(text: str) -> list[str]:
    """将中英文混合内容转换成可供 BM25 使用的词项。

    英文和数字按连续字符串提取；中文使用相邻两个字组成的二元词，
    这样不依赖额外中文分词服务也能完成本地检索。
    """

    lowered = text.lower()
    terms = ASCII_TERM_PATTERN.findall(lowered)
    for run in CHINESE_RUN_PATTERN.findall(lowered):
        if len(run) == 1:
            terms.append(run)
            continue
        terms.extend(run[index : index + 2] for index in range(len(run) - 1))
    return terms


def _chunk_to_hit(chunk: KnowledgeChunk, score: float) -> KnowledgeSearchHit:
    return KnowledgeSearchHit(
        chunk_id=chunk.id,
        document_id=chunk.document_id,
        document_version=chunk.document_version,
        ordinal=chunk.ordinal,
        title=chunk.title,
        content=chunk.content,
        document_type=chunk.document_type,
        source_name=chunk.source_name,
        product_code=chunk.product_code,
        score=score,
    )
