"""Chunking - 文档分块策略

提供两种分块策略：
- TextChunker: 按字符数分块，带重叠
- StructuralChunker: 按文档结构（标题/段落）分块，适用于 Markdown/HTML
"""

from __future__ import annotations

import logging
import re
import uuid
from dataclasses import dataclass, field
from typing import Any

logger = logging.getLogger(__name__)


@dataclass
class Chunk:
    """文档分块"""

    chunk_id: str
    content: str
    metadata: dict[str, Any] = field(default_factory=dict)

    def __post_init__(self) -> None:
        if "source_doc" not in self.metadata:
            self.metadata["source_doc"] = ""
        if "position" not in self.metadata:
            self.metadata["position"] = 0

    def to_dict(self) -> dict[str, Any]:
        return {
            "chunk_id": self.chunk_id,
            "content": self.content,
            "metadata": self.metadata,
        }


class TextChunker:
    """按字符数分块，带重叠

    适用于纯文本文档，按固定字符数切分，相邻块之间有重叠区域。
    """

    def __init__(
        self,
        chunk_size: int = 1000,
        chunk_overlap: int = 200,
        separator: str = "\n\n",
    ):
        """
        Args:
            chunk_size: 每个块的最大字符数
            chunk_overlap: 相邻块之间的重叠字符数
            separator: 优先分割符
        """
        if chunk_overlap >= chunk_size:
            raise ValueError(f"chunk_overlap ({chunk_overlap}) must be less than chunk_size ({chunk_size})")
        self._chunk_size = chunk_size
        self._chunk_overlap = chunk_overlap
        self._separator = separator

    def chunk(
        self,
        text: str,
        source_doc_id: str = "",
        metadata: dict[str, Any] | None = None,
    ) -> list[Chunk]:
        """将文本分块

        Args:
            text: 待分块文本
            source_doc_id: 源文档 ID
            metadata: 附加元数据

        Returns:
            Chunk 列表
        """
        if not text.strip():
            return []

        # 先尝试按分隔符分割
        segments = self._split_by_separator(text)

        # 合并小段，切分大段
        chunks_text = self._merge_and_split(segments)

        base_meta = dict(metadata or {})
        base_meta["source_doc"] = source_doc_id
        base_meta["chunking_strategy"] = "text"

        chunks = []
        for i, chunk_text in enumerate(chunks_text):
            chunk_meta = dict(base_meta)
            chunk_meta["position"] = i
            chunk_meta["char_count"] = len(chunk_text)
            chunks.append(Chunk(
                chunk_id=str(uuid.uuid4()),
                content=chunk_text,
                metadata=chunk_meta,
            ))

        return chunks

    def _split_by_separator(self, text: str) -> list[str]:
        """按分隔符分割文本"""
        segments = text.split(self._separator)
        # 过滤空段
        return [s.strip() for s in segments if s.strip()]

    def _merge_and_split(self, segments: list[str]) -> list[str]:
        """合并小段，切分大段"""
        result: list[str] = []
        current: list[str] = []
        current_len = 0

        for segment in segments:
            seg_len = len(segment)

            # 如果单个段超过 chunk_size，需要进一步切分
            if seg_len > self._chunk_size:
                # 先把当前累积的段输出
                if current:
                    result.append(self._separator.join(current))
                    current = []
                    current_len = 0

                # 切分大段
                for sub in self._split_large_segment(segment):
                    result.append(sub)
                continue

            # 如果加入当前段会超过 chunk_size，先输出当前累积
            if current_len + seg_len + len(self._separator) > self._chunk_size and current:
                result.append(self._separator.join(current))
                # 保留重叠部分
                overlap_text = self._separator.join(current)
                overlap_start = max(0, len(overlap_text) - self._chunk_overlap)
                overlap_segments = self._get_overlap_segments(
                    overlap_text[overlap_start:], segments
                )
                current = overlap_segments
                current_len = sum(len(s) for s in current) + len(self._separator) * max(0, len(current) - 1)

            current.append(segment)
            current_len += seg_len + len(self._separator)

        if current:
            result.append(self._separator.join(current))

        return result

    def _split_large_segment(self, segment: str) -> list[str]:
        """切分超大段"""
        result = []
        start = 0
        while start < len(segment):
            end = start + self._chunk_size
            # 尝试在句子边界切分
            if end < len(segment):
                # 查找最近的句子结束符
                for sep in ["。", ".", "！", "!", "？", "?", "\n"]:
                    last_sep = segment.rfind(sep, start + self._chunk_size // 2, end)
                    if last_sep > start:
                        end = last_sep + len(sep)
                        break
            result.append(segment[start:end].strip())
            start = end - self._chunk_overlap
            if start <= 0 and end >= len(segment):
                break
            if start < 0:
                start = 0
        return [r for r in result if r]

    def _get_overlap_segments(self, overlap_text: str, segments: list[str]) -> list[str]:
        """从重叠文本中提取完整段"""
        # 简化实现：将重叠文本作为一个段
        if overlap_text.strip():
            return [overlap_text.strip()]
        return []


class StructuralChunker:
    """按文档结构分块

    适用于 Markdown 和 HTML 等有标题结构的文档。
    按标题层级分块，每个标题下的内容作为一个块。
    如果某个块超过 chunk_size，则回退到 TextChunker 继续切分。
    """

    def __init__(
        self,
        chunk_size: int = 1000,
        chunk_overlap: int = 200,
        heading_levels: int = 3,
    ):
        """
        Args:
            chunk_size: 每个块的最大字符数
            chunk_overlap: 回退 TextChunker 时的重叠字符数
            heading_levels: 识别的标题层级数（1-6 对应 # 到 ######）
        """
        self._chunk_size = chunk_size
        self._chunk_overlap = chunk_overlap
        self._heading_levels = min(max(heading_levels, 1), 6)
        self._text_chunker = TextChunker(
            chunk_size=chunk_size,
            chunk_overlap=chunk_overlap,
        )

    def chunk(
        self,
        text: str,
        source_doc_id: str = "",
        metadata: dict[str, Any] | None = None,
    ) -> list[Chunk]:
        """将文本按结构分块

        Args:
            text: 待分块文本（Markdown 格式）
            source_doc_id: 源文档 ID
            metadata: 附加元数据

        Returns:
            Chunk 列表
        """
        if not text.strip():
            return []

        sections = self._split_by_headings(text)

        base_meta = dict(metadata or {})
        base_meta["source_doc"] = source_doc_id
        base_meta["chunking_strategy"] = "structural"

        chunks = []
        position = 0

        for section in sections:
            heading = section["heading"]
            content = section["content"]
            level = section["level"]

            if not content.strip():
                continue

            # 如果内容超过 chunk_size，使用 TextChunker 继续切分
            if len(content) > self._chunk_size:
                sub_chunks = self._text_chunker.chunk(
                    content,
                    source_doc_id=source_doc_id,
                    metadata=metadata,
                )
                for sub in sub_chunks:
                    sub.metadata["position"] = position
                    sub.metadata["heading"] = heading
                    sub.metadata["heading_level"] = level
                    sub.metadata["chunking_strategy"] = "structural"
                    position += 1
                    chunks.append(sub)
            else:
                chunk_meta = dict(base_meta)
                chunk_meta["position"] = position
                chunk_meta["heading"] = heading
                chunk_meta["heading_level"] = level
                chunk_meta["char_count"] = len(content)
                chunks.append(Chunk(
                    chunk_id=str(uuid.uuid4()),
                    content=content,
                    metadata=chunk_meta,
                ))
                position += 1

        return chunks

    def _split_by_headings(self, text: str) -> list[dict[str, Any]]:
        """按标题分割 Markdown 文本

        Returns:
            列表，每项包含 heading, content, level
        """
        lines = text.split("\n")
        sections: list[dict[str, Any]] = []
        current_heading = ""
        current_level = 0
        current_lines: list[str] = []

        heading_pattern = re.compile(r"^(#{1," + str(self._heading_levels) + r"})\s+(.+)$")

        for line in lines:
            match = heading_pattern.match(line)
            if match:
                # 保存当前节
                if current_lines:
                    content = "\n".join(current_lines).strip()
                    if content:
                        sections.append({
                            "heading": current_heading,
                            "content": content,
                            "level": current_level,
                        })

                # 开始新节
                current_heading = match.group(2).strip()
                current_level = len(match.group(1))
                current_lines = [line]
            else:
                current_lines.append(line)

        # 保存最后一节
        if current_lines:
            content = "\n".join(current_lines).strip()
            if content:
                sections.append({
                    "heading": current_heading,
                    "content": content,
                    "level": current_level,
                })

        # 如果没有标题结构，整体作为一个块
        if not sections:
            sections.append({
                "heading": "",
                "content": text.strip(),
                "level": 0,
            })

        return sections
