"""BM25 轻量近似稀疏向量编码器

零第三方依赖的稀疏编码实现：无全局 IDF 的 BM25 近似（仅用词频 tf 加权），
输出 Qdrant SparseVector 所需的 indices/values 格式。
后续可替换为 SPLADE / BM42 等更强的稀疏编码器。
"""

import hashlib
import math
import re

# 稀疏向量维度（哈希空间大小）
SPARSE_DIM = 2**18

# 连续 CJK 字符段（一-鿿）
_CJK_RE = re.compile(r"[一-鿿]+")
# CJK 段或英文/数字连续段
_TOKEN_RE = re.compile(r"[一-鿿]+|[a-zA-Z0-9]+")


def _tokenize(text: str) -> list[str]:
    """分词（纯正则实现）

    - 连续 CJK 字符段：长度 1 时保留 unigram，长度 >= 2 时生成字符 bigram
    - 英文/数字连续段：小写化后作为整词
    """
    tokens: list[str] = []
    for match in _TOKEN_RE.finditer(text):
        seg = match.group()
        if _CJK_RE.fullmatch(seg):
            if len(seg) == 1:
                tokens.append(seg)
            else:
                tokens.extend(seg[i : i + 2] for i in range(len(seg) - 1))
        else:
            tokens.append(seg.lower())
    return tokens


def _hash(token: str) -> int:
    """将词哈希到 [0, SPARSE_DIM) 的索引空间"""
    digest = hashlib.blake2b(token.encode("utf-8"), digest_size=8).digest()
    return int.from_bytes(digest, "big") % SPARSE_DIM


class SparseEncoder:
    """稀疏向量编码器，输出 Qdrant SparseVector 的 indices/values"""

    def encode(self, text: str) -> tuple[list[int], list[float]]:
        """编码单条文本

        权重 = 1 + log(tf)，词经哈希到 [0, SPARSE_DIM)，哈希冲突时权重累加。

        Returns:
            (indices, values)：indices 升序且无重复，与 values 等长
        """
        # 按哈希索引累计词频，天然处理哈希冲突
        tf: dict[int, float] = {}
        for token in _tokenize(text):
            idx = _hash(token)
            tf[idx] = tf.get(idx, 0.0) + 1.0
        indices = sorted(tf)
        values = [1.0 + math.log(tf[idx]) for idx in indices]
        return indices, values

    def encode_batch(self, texts: list[str]) -> list[tuple[list[int], list[float]]]:
        """批量编码，与逐条 encode 结果一致"""
        return [self.encode(text) for text in texts]