"""文档 chunk 切分器

按文档原生标题树切分 chunk：
- 结构化文本：每个标题起点切分 section（标题行到下一个标题前），
  超长 section 按空行段落二次切分，单段落仍超长则硬切
- 无结构文本：直接按空行段落累加切分

每个 chunk 记录 section_path（祖先标题链，" / " 连接），与 L2 大纲节点互相定位。
"""

import re

import structlog

from app.config import settings
from app.core.headings import Heading, parse_headings
from app.models.document import ChunkModel

logger = structlog.get_logger()

# 段落分隔：一个或多个空行
_PARAGRAPH_SPLIT_PATTERN = re.compile(r"\n\s*\n")


class Chunker:
    """按标题树切分文档 chunk"""

    def __init__(self, max_chars: int = settings.chunk_max_chars) -> None:
        self.max_chars = max_chars

    def chunk(self, text: str, doc_id: str) -> list[ChunkModel]:
        """将文档文本切分为 chunk 列表

        Args:
            text: 文档纯文本内容
            doc_id: 文档 ID

        Returns:
            list[ChunkModel]: 切分结果，chunk_index 从 0 递增
        """
        stripped = text.strip()
        if not stripped:
            return []

        # 全文不超长：整篇单 chunk
        if len(stripped) <= self.max_chars:
            return [ChunkModel(doc_id=doc_id, chunk_index=0, text=stripped)]

        headings = parse_headings(stripped)
        chunks: list[ChunkModel] = []

        if headings:
            # 结构化：先按标题切分 section，再按长度二次切分
            for section_text, section_path in self._split_sections(stripped, headings):
                for piece in self._split_by_length(section_text):
                    chunks.append(
                        ChunkModel(doc_id=doc_id, chunk_index=len(chunks), text=piece, section_path=section_path)
                    )
        else:
            # 无结构：直接按段落累加切分
            for piece in self._split_by_length(stripped):
                chunks.append(ChunkModel(doc_id=doc_id, chunk_index=len(chunks), text=piece))

        logger.info("文档切分完成", doc_id=doc_id, chunks_count=len(chunks), has_headings=bool(headings))
        return chunks

    def _split_sections(self, text: str, headings: list[Heading]) -> list[tuple[str, str]]:
        """按标题树切分 section，返回 (section 文本, section_path) 列表

        每个标题起点切分一个 section，section 文本含标题行本身；
        section_path 为祖先标题链（含自身标题），用 " / " 连接；
        首个标题前的引导正文归入无前缀 section（section_path 为空）。
        """
        lines = text.splitlines()
        sections: list[tuple[str, str]] = []

        # 首个标题前的引导内容
        preamble = "\n".join(lines[: headings[0].line_index]).strip()
        if preamble:
            sections.append((preamble, ""))

        # 维护祖先标题栈：遇到同级或更高级标题时弹栈
        stack: list[Heading] = []
        for i, heading in enumerate(headings):
            while stack and stack[-1].level >= heading.level:
                stack.pop()
            stack.append(heading)

            end = headings[i + 1].line_index if i + 1 < len(headings) else len(lines)
            section_text = "\n".join(lines[heading.line_index : end]).strip()
            section_path = " / ".join(h.title for h in stack)
            sections.append((section_text, section_path))

        return sections

    def _split_by_length(self, text: str) -> list[str]:
        """按 max_chars 切分文本：先按空行段落累加，单段落超长则硬切"""
        if len(text) <= self.max_chars:
            return [text]

        pieces: list[str] = []
        current = ""
        for paragraph in _PARAGRAPH_SPLIT_PATTERN.split(text):
            paragraph = paragraph.strip()
            if not paragraph:
                continue
            candidate = f"{current}\n\n{paragraph}" if current else paragraph
            if len(candidate) <= self.max_chars:
                current = candidate
                continue
            if current:
                pieces.append(current)
                current = ""
            # 单段落仍超长：按 max_chars 硬切
            if len(paragraph) > self.max_chars:
                pieces.extend(paragraph[i : i + self.max_chars] for i in range(0, len(paragraph), self.max_chars))
            else:
                current = paragraph
        if current:
            pieces.append(current)
        return pieces