"""基于本地 embedding 模型的语义分块器。

把文本按句子/子句切分成语义单元后，计算相邻单元向量相似度，
在相似度显著下降处建立分块边界，从而实现按话题主题分块。
"""
import logging
import math
import os
import re
from typing import List, Tuple

logger = logging.getLogger("semantic-chunker")

# 语义单元边界：段落/句子级别。与 text_splitter 分隔符保持一致，但粒度更粗。
_SEMANTIC_UNIT_SEPARATORS = ["\n\n", "。", "！", "？", "；", "\n"]
_UNIT_SPLIT_PATTERN = re.compile(
    "(" + "|".join(re.escape(s) for s in _SEMANTIC_UNIT_SEPARATORS) + ")"
)


def _split_to_units(text: str) -> List[str]:
    """
    把文本切分为语义单元，每个单元包含末尾分隔符。
    例如：['第一句话。', '第二句话！', '第三句话；']

    注意：regex split 在相邻分隔符之间会产生空字符串（如 "。
" 之间），
    且孤立的换行符不能独立成单元。这里把分隔符附加到当前单元，直到
    单元包含非空白内容后再结束，从而保证拼接结果等于原文。
    """
    parts = _UNIT_SPLIT_PATTERN.split(text)
    units: List[str] = []
    current = ""
    for part in parts:
        if part in _SEMANTIC_UNIT_SEPARATORS:
            current += part
            # 当前单元已包含非空白内容时结束该单元；否则继续等待后续文本
            if current.strip():
                units.append(current)
                current = ""
        else:
            current += part
    if current:
        units.append(current)
    return units


def _merge_tiny_units(units: List[str], min_length: int = 12) -> List[str]:
    """
    把过短的单元（如标题、单个语气词）向前合并，减少 embedding 噪声。
    合并后若仍过短则继续合并，直到末尾。
    """
    if not units:
        return []
    merged: List[str] = [units[0]]
    for unit in units[1:]:
        if len(merged[-1].strip()) < min_length:
            merged[-1] = merged[-1] + unit
        elif len(unit.strip()) < min_length:
            merged[-1] = merged[-1] + unit
        else:
            merged.append(unit)
    return merged


class SemanticChunker:
    """
    基于向量相似度的语义分块器。

    分块策略：
    1. 先把文本切分为语义单元（句子/段落）。
    2. 用本地 embedding 模型为每个单元生成向量。
    3. 计算相邻单元余弦相似度。
    4. 当相似度低于阈值，或当前段落长度超过 max_chunk_size 时，建立新边界。

    阈值策略：
    - 默认启用自适应阈值：取相邻单元相似度分布的 25 分位数，且不低于
      similarity_threshold 地板值。对中文等相似度整体偏高的语料更稳健。
    - 仅低于阈值还不够，必须同时满足“显著下降”条件：当前相似度比前一对
      单元下降超过 significant_drop，从而避免在整体都低但平稳的语料内部过切。
    - 关闭自适应阈值时，使用固定 similarity_threshold。

    重叠策略：
    - 当 chunk_overlap > 0 时，语义边界处的“承上启下”单元会被追加到下一个
      chunk 的开头，实现相邻 chunk 在句子/子句边界的重叠，而不是字符级滑动。
    - 重叠量受 chunk_overlap 上限限制，避免无限重复。
    """

    def __init__(
        self,
        similarity_threshold: float = 0.55,
        max_chunk_size: int = 800,
        min_unit_length: int = 12,
        batch_size: int = 32,
        use_adaptive_threshold: bool = True,
        significant_drop: float = 0.08,
        chunk_overlap: int = 0,
    ):
        self.similarity_threshold = similarity_threshold
        self.max_chunk_size = max_chunk_size
        self.min_unit_length = min_unit_length
        self.batch_size = batch_size
        self.use_adaptive_threshold = use_adaptive_threshold
        self.significant_drop = significant_drop
        self.chunk_overlap = chunk_overlap

    def split_text(self, text: str) -> List[str]:
        """
        对 text 执行语义分块。

        返回的段落列表按顺序拼接包含原文全部内容；若启用重叠，相邻段落之间会
        共享边界处的语义单元（承上启下句子），因此拼接结果会大于原文。
        """
        if not text:
            return []
        if len(text) <= self.max_chunk_size:
            return [text]

        units = _split_to_units(text)
        units = _merge_tiny_units(units, self.min_unit_length)
        if not units:
            return [text]
        if len(units) == 1:
            return [text]

        embeddings = self._embed_units(units)
        return self._build_chunks(units, embeddings)

    def _embed_units(self, units: List[str]) -> List[List[float]]:
        """批量获取语义单元向量。"""
        # 延迟导入，避免在 embedding-bridge 启动时立即加载模型
        try:
            from backend.indexing.embedding import embedding_service
        except ImportError:
            logger.warning("embedding_service 不可用，回退到规则分块")
            return []

        try:
            texts = [u.strip() for u in units]
            all_embeddings: List[List[float]] = []
            for i in range(0, len(texts), self.batch_size):
                batch = texts[i : i + self.batch_size]
                emb = embedding_service.get_embeddings(batch)
                all_embeddings.extend(emb)
            return all_embeddings
        except Exception as e:
            logger.warning("语义向量计算失败，回退到规则分块: %s", e)
            return []

    def _build_chunks(
        self, units: List[str], embeddings: List[List[float]]
    ) -> List[str]:
        """根据相似度构建语义段落，支持边界单元重叠。"""
        chunks: List[str] = []
        current_chunk = units[0]
        current_len = len(units[0])

        # 预计算所有相邻单元相似度
        similarities: List[float] = []
        if embeddings and len(embeddings) == len(units):
            for i in range(1, len(units)):
                similarities.append(_cosine_similarity(embeddings[i - 1], embeddings[i]))

        effective_threshold = self.similarity_threshold
        if self.use_adaptive_threshold and similarities:
            effective_threshold = max(
                self.similarity_threshold,
                _percentile(similarities, 0.25),
            )
            logger.debug("自适应语义阈值: %.3f (地板值: %.3f)", effective_threshold, self.similarity_threshold)

        for i in range(1, len(units)):
            unit = units[i]
            unit_len = len(unit)

            # 强制边界：超过最大长度必须切开
            if current_len + unit_len > self.max_chunk_size:
                chunks.append(current_chunk)
                current_chunk = unit
                current_len = unit_len
                continue

            # 语义边界：相似度低于有效阈值且出现显著下降时切开
            should_split = False
            if similarities and i - 1 < len(similarities):
                sim = similarities[i - 1]
                if sim < effective_threshold:
                    prev_sim = similarities[i - 2] if i >= 2 else 1.0
                    if prev_sim - sim > self.significant_drop:
                        should_split = True
                        logger.debug(
                            "语义边界: units[%s]~[%s] similarity=%.3f < %.3f, drop=%.3f",
                            i - 1,
                            i,
                            sim,
                            effective_threshold,
                            prev_sim - sim,
                        )

            if should_split:
                chunks.append(current_chunk)
                # 边界处的重叠单元：把当前 chunk 末尾若干完整语义单元共享给下一个 chunk
                overlap_units = self._extract_overlap_units(current_chunk)
                current_chunk = "".join(overlap_units) + unit
                current_len = sum(len(u) for u in overlap_units) + unit_len
            else:
                current_chunk += unit
                current_len += unit_len

        if current_chunk:
            chunks.append(current_chunk)

        return chunks

    def _extract_overlap_units(self, current_chunk: str) -> List[str]:
        """从当前 chunk 末尾提取不超过 chunk_overlap 的完整语义单元。"""
        if self.chunk_overlap <= 0 or not current_chunk:
            return []

        current_units = _split_to_units(current_chunk)
        if not current_units:
            return []

        overlap: List[str] = []
        total_len = 0
        for u in reversed(current_units):
            candidate_len = total_len + len(u)
            if overlap and candidate_len > self.chunk_overlap:
                break
            overlap.insert(0, u)
            total_len = candidate_len
            if total_len >= self.chunk_overlap:
                break

        return overlap


def _percentile(values: List[float], q: float) -> float:
    """计算分位数（q 在 [0, 1] 之间），不依赖 numpy。"""
    if not values:
        return 0.0
    sorted_values = sorted(values)
    n = len(sorted_values)
    if n == 1:
        return sorted_values[0]
    index = q * (n - 1)
    lower = int(index)
    upper = lower + 1
    if upper >= n:
        return sorted_values[-1]
    weight = index - lower
    return sorted_values[lower] * (1 - weight) + sorted_values[upper] * weight


def _cosine_similarity(a: List[float], b: List[float]) -> float:
    """计算两个向量的余弦相似度。"""
    if not a or not b or len(a) != len(b):
        return 0.0
    dot = sum(x * y for x, y in zip(a, b))
    norm_a = math.sqrt(sum(x * x for x in a))
    norm_b = math.sqrt(sum(x * x for x in b))
    if norm_a == 0 or norm_b == 0:
        return 0.0
    return dot / (norm_a * norm_b)
