"""Markdown 按标题分块"""

from __future__ import annotations

import re

from .base import BaseChunker, Chunk, estimate_tokens

# 匹配 Markdown 标题行
_HEADING_RE = re.compile(r"^(#{1,6})\s+(.+)$", re.MULTILINE)


class MarkdownChunker(BaseChunker):
    """按 Markdown 标题层级分割"""

    def chunk(self, content: str, file_path: str) -> list[Chunk]:
        content = _strip_frontmatter(content)

        if not content.strip():
            return []

        sections = self._split_by_heading(content)
        if not sections:
            return [Chunk(content=content.strip(), chunk_index=0)]

        chunks: list[Chunk] = []
        for sec in sections:
            if estimate_tokens(sec["content"]) > self.MAX_CHUNK_TOKENS:
                chunks.extend(self._split_oversized(
                    sec["content"], sec["heading"], sec["parent"], len(chunks)
                ))
            else:
                chunks.append(Chunk(
                    content=sec["content"],
                    section=sec["heading"],
                    parent_section=sec["parent"],
                    chunk_index=len(chunks),
                    tags=_extract_heading_tags(sec["heading"]),
                ))

        return chunks

    def _split_by_heading(self, content: str) -> list[dict]:
        """按标题分割，跳过代码块内的 #"""
        lines = content.split("\n")
        sections: list[dict] = []
        current_lines: list[str] = []
        current_heading: str | None = None
        parent_heading: str | None = None
        current_level = 0
        in_code_block = False

        for line in lines:
            # 跟踪代码块状态
            if line.strip().startswith("```"):
                in_code_block = not in_code_block

            if in_code_block:
                current_lines.append(line)
                continue

            match = _HEADING_RE.match(line)
            if match:
                # 保存上一段
                if current_lines or current_heading:
                    text = "\n".join(current_lines).strip()
                    if text or current_heading:
                        sections.append({
                            "heading": current_heading,
                            "parent": parent_heading,
                            "content": text,
                        })

                level = len(match.group(1))
                heading_text = f"{match.group(1)} {match.group(2)}"

                # 更新 parent 追踪
                if level <= current_level and current_level > 0:
                    # 同级或更高级标题，parent 不变
                    pass
                elif level > current_level and current_heading:
                    parent_heading = current_heading

                if level == 1:
                    parent_heading = None

                current_heading = heading_text
                current_level = level
                current_lines = []
            else:
                current_lines.append(line)

        # 最后一段
        text = "\n".join(current_lines).strip()
        if text or current_heading:
            sections.append({
                "heading": current_heading,
                "parent": parent_heading,
                "content": text,
            })

        # 过滤空内容
        return [s for s in sections if s["content"].strip()]


def _strip_frontmatter(content: str) -> str:
    """去除 YAML frontmatter"""
    if content.startswith("---"):
        end = content.find("---", 3)
        if end != -1:
            return content[end + 3:].strip()
    return content


def _extract_heading_tags(heading: str | None) -> list[str]:
    """从标题提取标签"""
    if not heading:
        return []
    # 去掉 # 符号，分词
    text = re.sub(r"^#+\s*", "", heading)
    words = re.findall(r"[\w\u4e00-\u9fff]+", text.lower())
    return [w for w in words if len(w) > 1]
