"""
Markdown 结构化分块策略

按 Markdown 标题层级分割文档，保留文档结构信息
"""
import re
import logging
from typing import Dict, Any, List, Optional

from .base import BaseChunker, ChunkResult

logger = logging.getLogger(__name__)


class MarkdownChunker(BaseChunker):
    """
    Markdown 分块器
    
    按标题层级分割 Markdown 文档，每个标题下的内容作为一个分块
    如果单个标题下的内容超过 chunk_size，则使用递归分割
    """

    def chunk(self, text: str, metadata: Dict[str, Any] = None) -> List[ChunkResult]:
        """按 Markdown 标题分块"""
        if not text or not text.strip():
            return []

        text = self._clean_text(text)
        metadata = metadata or {}

        # 按标题分割
        sections = self._split_by_headers(text)

        # 处理每个 section
        raw_chunks = []
        for section in sections:
            header = section.get('header', '')
            content = section.get('content', '')
            level = section.get('level', 0)

            if not content.strip():
                continue

            # 组合标题和内容
            full_text = f"{header}\n{content}" if header else content

            if len(full_text) <= self.chunk_size:
                raw_chunks.append({
                    'content': full_text.strip(),
                    'metadata': {
                        **metadata,
                        'header': header,
                        'header_level': level,
                    }
                })
            else:
                # 内容超长，递归分割
                sub_chunks = self._split_long_section(content, header)
                for i, sub in enumerate(sub_chunks):
                    raw_chunks.append({
                        'content': sub.strip(),
                        'metadata': {
                            **metadata,
                            'header': header,
                            'header_level': level,
                            'sub_chunk': i,
                        }
                    })

        # 合并过小的分块
        merged = self._merge_small_section_chunks(raw_chunks)

        # 构建结果
        results = []
        for i, item in enumerate(merged):
            if item['content'].strip():
                results.append(ChunkResult(
                    content=item['content'],
                    position=i,
                    metadata=item.get('metadata', {}),
                ))

        return results

    def _split_by_headers(self, text: str) -> List[Dict[str, Any]]:
        """按 Markdown 标题分割"""
        # 匹配 Markdown 标题: # Title, ## Title, ### Title 等
        header_pattern = re.compile(r'^(#{1,6})\s+(.+)$', re.MULTILINE)

        sections = []
        last_end = 0
        last_header = ''
        last_level = 0

        for match in header_pattern.finditer(text):
            # 保存前一个 section 的内容
            if last_end > 0 or match.start() > 0:
                content = text[last_end:match.start()]
                if content.strip() or last_header:
                    sections.append({
                        'header': last_header,
                        'content': content.strip(),
                        'level': last_level,
                    })

            last_header = match.group(0)
            last_level = len(match.group(1))
            last_end = match.end()

        # 最后一个 section
        remaining = text[last_end:]
        if remaining.strip() or last_header:
            sections.append({
                'header': last_header,
                'content': remaining.strip(),
                'level': last_level,
            })

        # 如果没有找到任何标题，整个文本作为一个 section
        if not sections:
            sections.append({
                'header': '',
                'content': text.strip(),
                'level': 0,
            })

        return sections

    def _split_long_section(self, content: str, header: str = '') -> List[str]:
        """分割超长的 section 内容"""
        from .recursive import RecursiveChunker

        chunker = RecursiveChunker(
            chunk_size=self.chunk_size,
            chunk_overlap=self.chunk_overlap,
        )
        results = chunker.chunk(content)

        chunks = []
        for i, result in enumerate(results):
            # 第一个分块带上标题
            if i == 0 and header:
                chunks.append(f"{header}\n{result.content}")
            else:
                chunks.append(result.content)

        return chunks if chunks else [content]

    def _merge_small_section_chunks(self, chunks: List[Dict], min_size: int = 80) -> List[Dict]:
        """合并过小的 section 分块"""
        if not chunks:
            return []

        merged = []
        buffer = None

        for chunk in chunks:
            if buffer is None:
                buffer = chunk
            elif len(buffer['content']) < min_size and len(buffer['content']) + len(chunk['content']) <= self.chunk_size:
                buffer['content'] = buffer['content'] + "\n\n" + chunk['content']
            else:
                merged.append(buffer)
                buffer = chunk

        if buffer:
            if merged and len(buffer['content']) < min_size:
                merged[-1]['content'] = merged[-1]['content'] + "\n\n" + buffer['content']
            else:
                merged.append(buffer)

        return merged