MarkdownTextSplitter.py

"""Markdown Text Splitter for RAG Applications.

A production-ready Markdown text splitter with comprehensive syntax support,
designed for Retrieval-Augmented Generation (RAG) pipelines.

Features:
    - Token-based chunk size control with configurable overlap (using tiktoken)
    - Complete Markdown syntax support (tables, lists, blockquotes, code blocks)
    - Alternative header syntax (underline style: === and ---)
    - YAML frontmatter extraction to metadata
    - Thread-safe implementation (safe for concurrent use)
    - Semantic context preservation (keeps related content together)
    - Intelligent sentence splitting with abbreviation handling
    - Token-level fallback for edge cases (base64, minified JSON)

Example:
    >>> from MarkdownTextSplitter import MarkdownTextSplitter, split_markdown
    >>>
    >>> # Using the class (chunk_size and chunk_overlap are in tokens)
    >>> splitter = MarkdownTextSplitter(chunk_size=500, chunk_overlap=50)
    >>> chunks = splitter.split_text(markdown_content)
    >>>
    >>> # Using the convenience function
    >>> chunks = split_markdown(markdown_content, chunk_size=500)

Note:
    For full YAML frontmatter support (lists, nested objects), install PyYAML:
    ``pip install pyyaml``
"""

from __future__ import annotations

import re
from dataclasses import dataclass, field
from typing import TYPE_CHECKING, Iterator

import tiktoken
from langchain_core.documents import Document

# Default tiktoken encoding (compatible with GPT-4, GPT-3.5-turbo, etc.)
TIKTOKEN_ENCODING = "cl100k_base"

if TYPE_CHECKING:
    from typing import Any

__all__ = [
    "MarkdownTextSplitter",
    "split_markdown",
]

__version__ = "1.0.0"


# =============================================================================
# Internal Data Structures
# =============================================================================


@dataclass
class _ParserState:
    """Thread-safe parser state container.

    Using a dataclass instead of instance variables ensures thread safety
    since each call to split_text gets its own state instance.
    """

    chunks: list[Document] = field(default_factory=list)
    current_content: list[str] = field(default_factory=list)
    header_stack: list[tuple[int, str, str]] = field(default_factory=list)
    header_counts: dict[str, int] = field(
        default_factory=dict
    )  # Track header occurrences
    in_code_block: bool = False
    code_fence: str = ""
    code_fence_length: int = 0
    code_language: str = ""
    in_table: bool = False
    table_content: list[str] = field(default_factory=list)
    in_list: bool = False
    list_content: list[str] = field(default_factory=list)
    list_indent: int = 0
    in_blockquote: bool = False
    blockquote_content: list[str] = field(default_factory=list)
    frontmatter_metadata: dict[str, str] = field(default_factory=dict)


# =============================================================================
# Main Splitter Class
# =============================================================================


class MarkdownTextSplitter:
    """A production-ready Markdown text splitter with comprehensive syntax support.

    This splitter is designed for RAG (Retrieval-Augmented Generation) pipelines
    and provides several advantages over basic text splitters:

    1. **Token-based chunk size control**: Supports ``chunk_size`` and ``chunk_overlap``
       parameters measured in tokens (using tiktoken) to prevent oversized chunks
       that exceed LLM token limits.

    2. **Complete Markdown support**: Properly handles tables, lists, blockquotes,
       code blocks, and alternative header syntax.

    3. **YAML frontmatter**: Extracts frontmatter metadata and propagates it
       to all chunks for filtering and context.

    4. **Thread-safe**: Uses local state instead of instance variables during
       parsing, allowing safe concurrent use in web applications.

    5. **Semantic preservation**: Keeps related content together (e.g., table rows,
       list items) and provides rich metadata for context reconstruction.

    6. **Intelligent splitting**: Uses paragraph -> sentence -> word -> token
       fallback chain with abbreviation-aware sentence detection.

    Attributes:
        chunk_size: Maximum chunk size in tokens (None for header-only splits).
        chunk_overlap: Number of tokens to overlap between chunks.
        encoding_name: Tiktoken encoding name (default: cl100k_base).
        strip_headers: Whether to exclude headers from chunk content.
        preserve_frontmatter: Whether to extract YAML frontmatter as metadata.
        keep_separator: Whether to keep horizontal rules in output.
        semantic_chunking: Whether to keep semantic units together.

    Example:
        >>> splitter = MarkdownTextSplitter(
        ...     chunk_size=500,
        ...     chunk_overlap=50,
        ...     headers_to_split_on=[
        ...         ("#", "title"),
        ...         ("##", "section"),
        ...     ]
        ... )
        >>> chunks = splitter.split_text(markdown_text)
        >>> for chunk in chunks:
        ...     print(chunk.metadata)  # Contains header hierarchy
    """

    # -------------------------------------------------------------------------
    # Class-level compiled regex patterns for performance
    # -------------------------------------------------------------------------

    _HEADER_PATTERN = re.compile(r"^(#{1,6})\s+(.+)$")
    _ALT_HEADER_1_PATTERN = re.compile(r"^=+\s*$")
    _ALT_HEADER_2_PATTERN = re.compile(r"^-+\s*$")
    _CODE_FENCE_PATTERN = re.compile(r"^(`{3,}|~{3,})(.*)$")
    _HORIZONTAL_RULE_PATTERN = re.compile(r"^(\*{3,}|-{3,}|_{3,})\s*$")
    _TABLE_ROW_PATTERN = re.compile(r"^\|.*\|.*$")
    _TABLE_SEPARATOR_PATTERN = re.compile(r"^\|[\s\-:|]+\|$")
    _LIST_ITEM_PATTERN = re.compile(r"^(\s*)([-*+]|\d+\.)\s+(.*)$")
    _BLOCKQUOTE_PATTERN = re.compile(r"^(\s*>)+\s?(.*)$")
    _FRONTMATTER_FENCE = re.compile(r"^-{3}\s*$")
    _INLINE_CODE_PATTERN = re.compile(r"`[^`]+`")

    # Common abbreviations that shouldn't trigger sentence splits
    _ABBREVIATIONS = frozenset(
        {
            "mr",
            "mrs",
            "ms",
            "dr",
            "prof",
            "sr",
            "jr",
            "vs",
            "etc",
            "inc",
            "ltd",
            "corp",
            "eg",
            "ie",
            "al",
            "cf",
            "vol",
            "no",
            "fig",
            "eq",
            "dept",
            "univ",
            "est",
            "approx",
            "govt",
            "avg",
            "min",
            "max",
        }
    )

    # -------------------------------------------------------------------------
    # Initialization
    # -------------------------------------------------------------------------

    def __init__(
        self,
        headers_to_split_on: list[tuple[str, str]] | None = None,
        chunk_size: int | None = None,
        chunk_overlap: int = 0,
        encoding_name: str = TIKTOKEN_ENCODING,
        strip_headers: bool = True,
        preserve_frontmatter: bool = True,
        keep_separator: bool = True,
        semantic_chunking: bool = True,
        deduplicate_headers: bool = False,
        strict_mode: bool = False,
        split_code_blocks: bool = False,
        filter_separators: bool = False,
        merge_small_chunks: bool = False,
        min_chunk_size: int = 50,
        isolate_code_blocks: bool = False,
        isolate_tables: bool = False,
    ) -> None:
        """Initialize the Markdown text splitter.

        Args:
            headers_to_split_on: List of tuples ``(header_marker, metadata_key)``.
                Example: ``[("#", "title"), ("##", "section")]``.
                Defaults to standard Markdown headers (# through ######).
            chunk_size: Maximum size of chunks in tokens. If ``None``,
                chunks are only split on headers (original behavior).
            chunk_overlap: Number of tokens to overlap between chunks.
                Only used when ``chunk_size`` is set.
            encoding_name: Tiktoken encoding name for token counting.
                Defaults to "cl100k_base" (GPT-4, GPT-3.5-turbo compatible).
            strip_headers: Whether to exclude headers from chunk content.
                Headers are always included in metadata regardless.
            preserve_frontmatter: Whether to extract YAML frontmatter as
                metadata instead of discarding it.
            keep_separator: Whether to keep horizontal rules (---) in output.
            semantic_chunking: Whether to keep semantic units together (tables,
                code blocks) even if they exceed ``chunk_size`` slightly.
            deduplicate_headers: Whether to add numeric suffixes to duplicate
                header names for unique identification in RAG systems.
                Example: "Conclusion" -> "Conclusion", "Conclusion (2)", etc.
            strict_mode: Enable all strict token compliance features:
                split_code_blocks, filter_separators, and merge_small_chunks.
            split_code_blocks: Split code blocks that exceed chunk_size.
                Overrides semantic_chunking for code blocks only.
            filter_separators: Remove separator-only chunks (like ---).
            merge_small_chunks: Merge consecutive small chunks that share
                the same header hierarchy.
            min_chunk_size: Minimum chunk size in tokens for merging.
                Chunks below this threshold will be merged with adjacent chunks
                if they share the same headers. Only used when merge_small_chunks
                is enabled.
            isolate_code_blocks: If True, each code block (delimited by ```)
                becomes a separate chunk, even in header-only mode.
            isolate_tables: If True, each table becomes a separate chunk,
                even in header-only mode.
        """
        self.chunk_size = chunk_size
        self.chunk_overlap = chunk_overlap
        self.encoding_name = encoding_name
        self.strip_headers = strip_headers
        self.preserve_frontmatter = preserve_frontmatter
        self.keep_separator = keep_separator
        self.semantic_chunking = semantic_chunking
        self.deduplicate_headers = deduplicate_headers
        self.isolate_code_blocks = isolate_code_blocks
        self.isolate_tables = isolate_tables

        # Strict mode enables all strict features
        self.strict_mode = strict_mode
        self.split_code_blocks = split_code_blocks or strict_mode
        self.filter_separators = filter_separators or strict_mode
        self.merge_small_chunks = merge_small_chunks or strict_mode
        self.min_chunk_size = min_chunk_size

        # Initialize tiktoken encoding for token counting
        self._encoding = tiktoken.get_encoding(self.encoding_name)

        if headers_to_split_on:
            self._splittable_headers = dict(headers_to_split_on)
        else:
            self._splittable_headers = {
                "#": "Header 1",
                "##": "Header 2",
                "###": "Header 3",
                "####": "Header 4",
                "#####": "Header 5",
                "######": "Header 6",
            }

    # -------------------------------------------------------------------------
    # Token Counting Helpers
    # -------------------------------------------------------------------------

    def _count_tokens(self, text: str) -> int:
        """Count the number of tokens in a text string.

        Args:
            text: The text to count tokens for.

        Returns:
            Number of tokens in the text.
        """
        if not text:
            return 0
        return len(self._encoding.encode(text))

    def _get_last_n_tokens_text(self, text: str, n_tokens: int) -> str:
        """Get the text corresponding to the last n tokens.

        Args:
            text: The source text.
            n_tokens: Number of tokens to extract from the end.

        Returns:
            The text corresponding to the last n tokens.
        """
        if not text or n_tokens <= 0:
            return ""
        tokens = self._encoding.encode(text)
        if len(tokens) <= n_tokens:
            return text
        return self._encoding.decode(tokens[-n_tokens:])

    def _wrap_code_fence(
        self,
        content: str,
        opening_fence: str,
        closing_fence: str,
    ) -> str:
        """Wrap content in code fences.

        Args:
            content: The code content to wrap.
            opening_fence: The opening fence line (e.g., "```python\\n").
            closing_fence: The closing fence line (e.g., "```\\n").

        Returns:
            Content wrapped in code fences.
        """
        return opening_fence + content.rstrip("\n") + "\n" + closing_fence

    def _accumulate_parts(
        self,
        parts: list[str],
        max_tokens: int,
        separator: str,
        metadata: dict[str, Any],
        start_index: int = 0,
        split_type: str | None = None,
        wrapper: tuple[str, str] | None = None,
        oversized_handler: Any | None = None,
    ) -> list[Document]:
        """Accumulate parts into chunks respecting token limits.

        This is the common pattern used across multiple splitting methods.

        Args:
            parts: List of text parts to accumulate.
            max_tokens: Maximum tokens per chunk.
            separator: Separator between parts (e.g., "\\n\\n", " ").
            metadata: Base metadata for chunks.
            start_index: Starting chunk index.
            split_type: Optional split type for metadata.
            wrapper: Optional (opening, closing) tuple for wrapping content.
            oversized_handler: Optional callable(part, metadata, index) -> list[Document]
                for handling parts that exceed max_tokens.

        Returns:
            List of Document chunks.
        """
        result = []
        current_chunk = ""
        chunk_index = start_index

        for part in parts:
            part = part.strip() if separator == "\n\n" else part
            if not part:
                continue

            part_tokens = self._count_tokens(part)

            # Handle oversized parts
            if part_tokens > max_tokens:
                if current_chunk:
                    result.append(
                        self._create_chunk(
                            current_chunk, metadata, chunk_index, split_type, wrapper
                        )
                    )
                    chunk_index += 1
                    current_chunk = ""

                if oversized_handler:
                    sub_chunks = oversized_handler(part, metadata, chunk_index)
                    result.extend(sub_chunks)
                    chunk_index += len(sub_chunks)
                else:
                    # Default: add as-is with warning
                    result.append(
                        self._create_chunk(
                            part, metadata, chunk_index, split_type, wrapper
                        )
                    )
                    chunk_index += 1
                continue

            # Try to accumulate
            test_content = current_chunk + (separator if current_chunk else "") + part
            if self._count_tokens(test_content) <= max_tokens:
                current_chunk = test_content
            else:
                if current_chunk:
                    result.append(
                        self._create_chunk(
                            current_chunk, metadata, chunk_index, split_type, wrapper
                        )
                    )
                    chunk_index += 1

                    # Apply overlap if configured
                    if self.chunk_overlap > 0:
                        overlap = self._get_last_n_tokens_text(
                            current_chunk, self.chunk_overlap
                        )
                        candidate = overlap + separator + part
                        # Fall back to no overlap if it would exceed chunk_size
                        if max_tokens and self._count_tokens(candidate) > max_tokens:
                            current_chunk = part
                        else:
                            current_chunk = candidate
                    else:
                        current_chunk = part
                else:
                    current_chunk = part

        if current_chunk:
            result.append(
                self._create_chunk(
                    current_chunk, metadata, chunk_index, split_type, wrapper
                )
            )

        return result

    def _create_chunk(
        self,
        content: str,
        metadata: dict[str, Any],
        chunk_index: int,
        split_type: str | None = None,
        wrapper: tuple[str, str] | None = None,
    ) -> Document:
        """Create a Document chunk with consistent metadata.

        Args:
            content: The chunk content.
            metadata: Base metadata.
            chunk_index: Index of this chunk.
            split_type: Optional split type for metadata.
            wrapper: Optional (opening, closing) tuple for wrapping.

        Returns:
            Document with content and metadata.
        """
        if wrapper:
            content = self._wrap_code_fence(content, wrapper[0], wrapper[1])

        chunk_metadata = {**metadata, "chunk_index": chunk_index}
        if split_type:
            chunk_metadata["split_type"] = split_type

        return Document(page_content=content, metadata=chunk_metadata)

    # -------------------------------------------------------------------------
    # Public Methods
    # -------------------------------------------------------------------------

    def split_text(self, text: str) -> list[Document]:
        """Split Markdown text into chunks with metadata.

        This method is thread-safe and can be called concurrently.

        Args:
            text: The Markdown text to split.

        Returns:
            List of Document objects with ``page_content`` and ``metadata``.
            Metadata includes header hierarchy and content type information.
        """
        state = _ParserState()

        lines = text.splitlines(keepends=True)
        lines = self._extract_frontmatter(lines, state)
        self._process_lines(lines, state)
        if state.in_code_block:
            self._finalize_code_block(state)
        else:
            self._finalize_current_block(state)

        if self.chunk_size:
            state.chunks = self._apply_chunk_size_limits(state.chunks)

        # Apply strict mode post-processing
        if self.filter_separators:
            state.chunks = self._filter_separator_chunks(state.chunks)

        if self.merge_small_chunks and self.chunk_size:
            state.chunks = self._merge_small_chunks_by_headers(state.chunks)

        return state.chunks

    def split_documents(self, documents: list[Document]) -> list[Document]:
        """Split multiple documents, preserving original metadata.

        Args:
            documents: List of Document objects to split.

        Returns:
            List of split documents with merged metadata (original + chunk).
        """
        result = []
        for doc in documents:
            chunks = self.split_text(doc.page_content)
            for chunk in chunks:
                merged_metadata = {**doc.metadata, **chunk.metadata}
                result.append(
                    Document(
                        page_content=chunk.page_content,
                        metadata=merged_metadata,
                    )
                )
        return result

    def lazy_split_text(self, text: str) -> Iterator[Document]:
        """Split text and yield chunks one at a time.

        Note: The full document is processed internally by split_text() first.
        This method provides a generator interface for convenience when iterating
        over results, but does not reduce peak memory usage.

        Args:
            text: The Markdown text to split.

        Yields:
            Document objects one at a time.
        """
        for chunk in self.split_text(text):
            yield chunk

    # -------------------------------------------------------------------------
    # Frontmatter Handling
    # -------------------------------------------------------------------------

    def _extract_frontmatter(
        self,
        lines: list[str],
        state: _ParserState,
    ) -> list[str]:
        """Extract YAML frontmatter from the beginning of the document."""
        if not lines:
            return lines

        first_line = lines[0].strip()
        if not self._FRONTMATTER_FENCE.match(first_line):
            return lines

        for i, line in enumerate(lines[1:], start=1):
            if self._FRONTMATTER_FENCE.match(line.strip()):
                frontmatter_lines = lines[1:i]
                if self.preserve_frontmatter:
                    state.frontmatter_metadata = self._parse_yaml(frontmatter_lines)
                return lines[i + 1 :]

        return lines

    def _parse_yaml(self, lines: list[str]) -> dict[str, str]:
        """Parse YAML frontmatter into metadata dictionary.

        Attempts to use PyYAML if available for full YAML support.
        Falls back to a simple key: value parser otherwise.
        """
        yaml_text = "\n".join(lines)

        # Try PyYAML first if available
        try:
            import yaml

            parsed = yaml.safe_load(yaml_text)
            if isinstance(parsed, dict):
                return {
                    str(k): str(v) if not isinstance(v, (list, dict)) else repr(v)
                    for k, v in parsed.items()
                    if v is not None
                }
        except ImportError:
            pass
        except Exception:
            pass

        # Simple fallback parser
        metadata: dict[str, str] = {}
        current_key: str | None = None
        multiline_value: list[str] = []

        for line in lines:
            stripped = line.strip()

            if not stripped or stripped.startswith("#"):
                continue

            if ":" in stripped and not stripped.startswith("-"):
                if current_key and multiline_value:
                    metadata[current_key] = " ".join(multiline_value).strip()
                    multiline_value = []

                key, _, value = stripped.partition(":")
                key = key.strip()
                value = value.strip()

                if not key:
                    continue

                if value.startswith("[") and value.endswith("]"):
                    metadata[key] = value
                    current_key = None
                elif value and value[0] in "\"'":
                    metadata[key] = value.strip("\"'")
                    current_key = None
                elif value:
                    metadata[key] = value
                    current_key = None
                else:
                    current_key = key
            elif current_key:
                multiline_value.append(stripped)

        if current_key and multiline_value:
            metadata[current_key] = " ".join(multiline_value).strip()

        return metadata

    # -------------------------------------------------------------------------
    # Line Processing
    # -------------------------------------------------------------------------

    def _process_lines(self, lines: list[str], state: _ParserState) -> None:
        """Process all lines and build chunks.

        Handles code blocks, tables, lists, blockquotes, headers, and regular text.
        In header-only mode (chunk_size is None), only splits on headers unless
        isolate_code_blocks or isolate_tables is enabled.
        """
        header_only_mode = self.chunk_size is None

        i = 0
        while i < len(lines):
            line = lines[i]
            raw_line = line
            stripped = line.strip()

            # Handle code blocks (highest priority)
            if state.in_code_block:
                i = self._handle_code_block_content(lines, i, state)
                continue

            # Check for code block start
            code_match = self._CODE_FENCE_PATTERN.match(stripped)
            if code_match:
                # Finalize current block if NOT in header-only mode OR if isolating code blocks
                if not header_only_mode or self.isolate_code_blocks:
                    self._finalize_current_block(state)
                state.in_code_block = True
                state.code_fence = code_match.group(1)[0]
                state.code_fence_length = len(code_match.group(1))
                state.code_language = code_match.group(2).strip()
                if header_only_mode and not self.isolate_code_blocks:
                    # In header-only mode (without isolation), add code fence to current content
                    state.current_content.append(raw_line)
                else:
                    state.current_content = [raw_line]
                i += 1
                continue

            # Handle tables
            if self._is_table_line(stripped, state):
                if not state.in_table:
                    # Finalize current block if NOT in header-only mode OR if isolating tables
                    if not header_only_mode or self.isolate_tables:
                        self._finalize_current_block(state)
                    state.in_table = True
                if header_only_mode and not self.isolate_tables:
                    # In header-only mode (without isolation), add to current content
                    state.current_content.append(raw_line)
                else:
                    state.table_content.append(raw_line)
                i += 1
                continue
            elif state.in_table:
                state.in_table = False
                if not header_only_mode or self.isolate_tables:
                    self._finalize_table(state)
                else:
                    state.table_content = []

            # Handle lists
            list_match = self._LIST_ITEM_PATTERN.match(line)
            if list_match or (
                state.in_list and self._is_list_continuation(line, state)
            ):
                if not state.in_list:
                    if not header_only_mode:
                        self._finalize_current_block(state)
                    state.in_list = True
                    state.list_indent = len(list_match.group(1)) if list_match else 0
                if header_only_mode:
                    state.current_content.append(raw_line)
                else:
                    state.list_content.append(raw_line)
                i += 1
                continue
            elif state.in_list:
                state.in_list = False
                state.list_indent = 0
                if not header_only_mode:
                    self._finalize_list(state)
                else:
                    state.list_content = []

            # Handle blockquotes
            blockquote_match = self._BLOCKQUOTE_PATTERN.match(line)
            if blockquote_match:
                if not state.in_blockquote:
                    if not header_only_mode:
                        self._finalize_current_block(state)
                    state.in_blockquote = True
                if header_only_mode:
                    state.current_content.append(raw_line)
                else:
                    state.blockquote_content.append(raw_line)
                i += 1
                continue
            elif state.in_blockquote:
                state.in_blockquote = False
                if not header_only_mode:
                    self._finalize_blockquote(state)
                else:
                    state.blockquote_content = []

            # Handle horizontal rules
            if self._HORIZONTAL_RULE_PATTERN.match(stripped) and not state.in_list:
                if header_only_mode:
                    # In header-only mode, treat separator as regular content
                    state.current_content.append(raw_line)
                else:
                    self._finalize_current_block(state)
                    if self.keep_separator:
                        state.current_content = [raw_line]
                        self._finalize_current_block(state, is_separator=True)
                i += 1
                continue

            # Handle headers
            # Use protected_line (inline code stripped) for detection to avoid
            # false matches on # inside inline code, but extract text from
            # the original stripped line to preserve inline code content
            protected_line = self._INLINE_CODE_PATTERN.sub("", stripped)
            header_match = self._HEADER_PATTERN.match(protected_line)

            if header_match and self._HEADER_PATTERN.match(stripped):
                actual_match = self._HEADER_PATTERN.match(stripped)
                self._handle_header(
                    state,
                    len(header_match.group(1)),
                    actual_match.group(2),
                    raw_line,
                )
                i += 1
                continue

            # Alternative header (underline style)
            if i + 1 < len(lines):
                next_stripped = lines[i + 1].strip()
                if self._ALT_HEADER_1_PATTERN.match(next_stripped) and stripped:
                    self._handle_header(state, 1, stripped, raw_line)
                    i += 2
                    continue
                elif (
                    self._ALT_HEADER_2_PATTERN.match(next_stripped)
                    and stripped
                    and not self._HORIZONTAL_RULE_PATTERN.match(next_stripped)
                ):
                    self._handle_header(state, 2, stripped, raw_line)
                    i += 2
                    continue

            # Regular text
            state.current_content.append(raw_line)
            i += 1

    def _handle_code_block_content(
        self,
        lines: list[str],
        start_idx: int,
        state: _ParserState,
    ) -> int:
        """Handle content inside a code block.

        Processes lines until closing fence is found, respecting header-only mode
        and code block isolation settings.
        """
        line = lines[start_idx]
        stripped = line.strip()
        header_only_mode = self.chunk_size is None

        close_match = self._CODE_FENCE_PATTERN.match(stripped)
        if close_match:
            fence_char = close_match.group(1)[0]
            fence_len = len(close_match.group(1))

            if fence_char == state.code_fence and fence_len >= state.code_fence_length:
                state.current_content.append(line)
                if header_only_mode and not self.isolate_code_blocks:
                    # In header-only mode (without isolation), keep code block as part of current content
                    state.in_code_block = False
                    state.code_fence = ""
                    state.code_fence_length = 0
                    state.code_language = ""
                else:
                    # Either not in header-only mode OR isolating code blocks
                    self._finalize_code_block(state)
                return start_idx + 1

        state.current_content.append(line)
        return start_idx + 1

    def _handle_header(
        self,
        state: _ParserState,
        level: int,
        text: str,
        raw_line: str,
    ) -> None:
        """Handle a header line.

        Finalizes current block if header should split, updates header stack,
        and applies deduplication if enabled.
        """
        header_key = self._splittable_headers.get("#" * level)

        # Only finalize (create new chunk) if this header level should split
        if header_key:
            self._finalize_current_block(state)

            # Apply header deduplication if enabled
            display_text = text
            if self.deduplicate_headers:
                # Create unique key combining header level and text
                dedup_key = f"{header_key}:{text}"
                count = state.header_counts.get(dedup_key, 0) + 1
                state.header_counts[dedup_key] = count
                if count > 1:
                    display_text = f"{text} ({count})"

            state.header_stack = [
                (lvl, key, val) for lvl, key, val in state.header_stack if lvl < level
            ]
            state.header_stack.append((level, header_key, display_text))

        if not self.strip_headers:
            state.current_content.append(raw_line)

    # -------------------------------------------------------------------------
    # Block Detection Helpers
    # -------------------------------------------------------------------------

    def _is_table_line(self, line: str, state: _ParserState) -> bool:
        """Check if line is part of a Markdown table with strict detection.

        Uses multiple heuristics: pipe delimiters, cell count, separator patterns,
        and cell length ratios to avoid false positives.
        """
        if state.in_table:
            return bool(
                self._TABLE_ROW_PATTERN.match(line)
                or self._TABLE_SEPARATOR_PATTERN.match(line)
            )

        if not line.startswith("|") or not line.endswith("|"):
            return False

        if line.count("|") < 3:
            return False

        if self._TABLE_SEPARATOR_PATTERN.match(line):
            return True

        cells = line.strip("|").split("|")
        if len(cells) < 2:
            return False

        cell_lengths = [len(c.strip()) for c in cells]
        if cell_lengths:
            max_len = max(cell_lengths)
            min_len = min(cell_lengths)
            if min_len > 0 and max_len / min_len > 10:
                return False

        return True

    def _is_list_continuation(self, line: str, state: _ParserState) -> bool:
        """Check if line continues a list.

        Returns True for blank lines or lines with greater indentation than
        the list's base indent level.
        """
        if not line.strip():
            return True
        indent = len(line) - len(line.lstrip())
        return indent > state.list_indent

    # -------------------------------------------------------------------------
    # Block Finalization
    # -------------------------------------------------------------------------

    def _finalize_current_block(
        self,
        state: _ParserState,
        is_separator: bool = False,
    ) -> None:
        """Finalize the current content block and create a chunk."""
        if not state.current_content:
            return

        content = "".join(state.current_content)

        if not content or content.isspace():
            state.current_content = []
            return

        metadata = self._build_metadata(state, is_separator=is_separator)
        state.chunks.append(Document(page_content=content, metadata=metadata))
        state.current_content = []

    def _finalize_code_block(self, state: _ParserState) -> None:
        """Finalize a code block."""
        content = "".join(state.current_content)
        metadata = self._build_metadata(state, block_type="code")
        if state.code_language:
            metadata["language"] = state.code_language

        state.chunks.append(Document(page_content=content, metadata=metadata))

        state.current_content = []
        state.in_code_block = False
        state.code_fence = ""
        state.code_fence_length = 0
        state.code_language = ""

    def _finalize_table(self, state: _ParserState) -> None:
        """Finalize a table block."""
        if not state.table_content:
            state.in_table = False
            return

        content = "".join(state.table_content)
        metadata = self._build_metadata(state, block_type="table")
        state.chunks.append(Document(page_content=content, metadata=metadata))

        state.table_content = []
        state.in_table = False

    def _finalize_list(self, state: _ParserState) -> None:
        """Finalize a list block."""
        if not state.list_content:
            state.in_list = False
            return

        while state.list_content and not state.list_content[-1].strip():
            state.list_content.pop()

        if not state.list_content:
            state.in_list = False
            return

        content = "".join(state.list_content)
        metadata = self._build_metadata(state, block_type="list")
        state.chunks.append(Document(page_content=content, metadata=metadata))

        state.list_content = []
        state.in_list = False
        state.list_indent = 0

    def _finalize_blockquote(self, state: _ParserState) -> None:
        """Finalize a blockquote block."""
        if not state.blockquote_content:
            state.in_blockquote = False
            return

        content = "".join(state.blockquote_content)
        metadata = self._build_metadata(state, block_type="blockquote")
        state.chunks.append(Document(page_content=content, metadata=metadata))

        state.blockquote_content = []
        state.in_blockquote = False

    def _build_metadata(
        self,
        state: _ParserState,
        block_type: str | None = None,
        is_separator: bool = False,
    ) -> dict[str, Any]:
        """Build metadata dictionary for a chunk.

        Includes frontmatter metadata, current header hierarchy, and block type
        information (code, table, list, blockquote, separator).
        """
        metadata: dict[str, Any] = {}

        if state.frontmatter_metadata:
            metadata.update(state.frontmatter_metadata)

        for _, key, value in state.header_stack:
            metadata[key] = value

        if block_type:
            metadata["type"] = block_type
        elif is_separator:
            metadata["type"] = "separator"

        return metadata

    # -------------------------------------------------------------------------
    # Chunk Size Management
    # -------------------------------------------------------------------------

    def _apply_chunk_size_limits(self, chunks: list[Document]) -> list[Document]:
        """Apply chunk size limits (in tokens), splitting oversized chunks.

        Handles code blocks, tables, and text differently based on semantic_chunking
        and split_code_blocks settings. Preserves semantic units when possible.
        """
        result = []

        for chunk in chunks:
            content = chunk.page_content

            if self._count_tokens(content) <= self.chunk_size:
                result.append(chunk)
                continue

            block_type = chunk.metadata.get("type")

            # Handle code blocks: split if split_code_blocks is enabled
            if block_type == "code":
                if self.split_code_blocks:
                    result.extend(self._split_code_block(chunk))
                elif self.semantic_chunking:
                    oversized_chunk = Document(
                        page_content=content,
                        metadata={
                            **chunk.metadata,
                            "warning": "oversized_semantic_block",
                        },
                    )
                    result.append(oversized_chunk)
                else:
                    result.extend(self._split_large_chunk(chunk))
                continue

            # Handle tables: keep semantic if enabled
            if block_type == "table" and self.semantic_chunking:
                oversized_chunk = Document(
                    page_content=content,
                    metadata={**chunk.metadata, "warning": "oversized_semantic_block"},
                )
                result.append(oversized_chunk)
                continue

            result.extend(self._split_large_chunk(chunk))

        return result

    def _split_large_chunk(self, chunk: Document) -> list[Document]:
        """Split a large chunk into smaller pieces with overlap (token-based).

        Uses paragraph -> sentence -> word -> token fallback chain for splitting.
        Applies chunk_overlap when configured.
        """
        content = chunk.page_content
        metadata = chunk.metadata
        result = []

        paragraphs = re.split(r"\n\s*\n", content)
        current_chunk = ""
        chunk_index = 0

        for para in paragraphs:
            para = para.strip()
            if not para:
                continue

            if self._count_tokens(para) > self.chunk_size:
                if current_chunk:
                    result.append(
                        Document(
                            page_content=current_chunk,
                            metadata={**metadata, "chunk_index": chunk_index},
                        )
                    )
                    chunk_index += 1
                    current_chunk = ""

                sentence_chunks = self._split_by_sentences(para, metadata, chunk_index)
                result.extend(sentence_chunks)
                chunk_index += len(sentence_chunks)
                continue

            test_content = current_chunk + ("\n\n" if current_chunk else "") + para

            if self._count_tokens(test_content) <= self.chunk_size:
                current_chunk = test_content
            else:
                if current_chunk:
                    result.append(
                        Document(
                            page_content=current_chunk,
                            metadata={**metadata, "chunk_index": chunk_index},
                        )
                    )
                    chunk_index += 1

                    if self.chunk_overlap > 0:
                        overlap_text = self._get_last_n_tokens_text(
                            current_chunk, self.chunk_overlap
                        )
                        candidate = overlap_text + "\n\n" + para
                        # Fall back to no overlap if it would exceed chunk_size
                        if self._count_tokens(candidate) > self.chunk_size:
                            current_chunk = para
                        else:
                            current_chunk = candidate
                    else:
                        current_chunk = para
                else:
                    current_chunk = para

        if current_chunk:
            if self._count_tokens(current_chunk) > self.chunk_size:
                sentence_chunks = self._split_by_sentences(
                    current_chunk, metadata, chunk_index
                )
                result.extend(sentence_chunks)
            else:
                result.append(
                    Document(
                        page_content=current_chunk,
                        metadata={**metadata, "chunk_index": chunk_index},
                    )
                )

        return result

    def _split_by_sentences(
        self,
        text: str,
        metadata: dict[str, Any],
        start_index: int,
    ) -> list[Document]:
        """Split text by sentences with abbreviation-aware detection (token-based)."""
        sentences = self._smart_sentence_split(text)
        result = []
        current_chunk = ""
        chunk_index = start_index

        for sentence in sentences:
            if self._count_tokens(sentence) > self.chunk_size:
                if current_chunk:
                    result.append(
                        Document(
                            page_content=current_chunk,
                            metadata={**metadata, "chunk_index": chunk_index},
                        )
                    )
                    chunk_index += 1

                token_chunks = self._split_by_tokens(sentence, metadata, chunk_index)
                result.extend(token_chunks)
                chunk_index += len(token_chunks)
                current_chunk = ""
                continue

            test_content = current_chunk + (" " if current_chunk else "") + sentence

            if self._count_tokens(test_content) <= self.chunk_size:
                current_chunk = test_content
            else:
                if current_chunk:
                    result.append(
                        Document(
                            page_content=current_chunk,
                            metadata={**metadata, "chunk_index": chunk_index},
                        )
                    )
                    chunk_index += 1

                if self.chunk_overlap > 0 and current_chunk:
                    overlap = self._get_last_n_tokens_text(
                        current_chunk, self.chunk_overlap
                    )
                    candidate = overlap + " " + sentence
                    # Fall back to no overlap if it would exceed chunk_size
                    if self._count_tokens(candidate) > self.chunk_size:
                        current_chunk = sentence
                    else:
                        current_chunk = candidate
                else:
                    current_chunk = sentence

        if current_chunk:
            result.append(
                Document(
                    page_content=current_chunk,
                    metadata={**metadata, "chunk_index": chunk_index},
                )
            )

        return result

    def _smart_sentence_split(self, text: str) -> list[str]:
        """Split text into sentences, handling abbreviations properly.

        Uses regex-based sentence boundary detection with abbreviation awareness
        (Mr., Dr., etc.) and checks for lowercase letter following period.
        """
        sentences = []
        current = ""

        parts = re.split(r"([.!?]+)\s+", text)

        i = 0
        while i < len(parts):
            next_part = parts[i + 1] if i + 1 < len(parts) else ""
            if i + 1 < len(parts) and re.match(r"^[.!?]+$", next_part):
                current += parts[i] + parts[i + 1]

                words = parts[i].split()
                last_word = words[-1].lower().rstrip(".") if words else ""

                if last_word in self._ABBREVIATIONS:
                    current += " "
                    i += 2
                    continue

                if i + 2 < len(parts) and parts[i + 2] and parts[i + 2][0].islower():
                    current += " "
                    i += 2
                    continue

                sentences.append(current.strip())
                current = ""
                i += 2
            else:
                current += parts[i]
                i += 1

        if current.strip():
            sentences.append(current.strip())

        return sentences if sentences else [text]

    def _split_by_tokens(
        self,
        text: str,
        metadata: dict[str, Any],
        start_index: int,
    ) -> list[Document]:
        """Last-resort token-level splitting for oversized content."""
        result = []
        chunk_index = start_index

        if " " in text:
            # Try word-boundary splitting first
            words = text.split(" ")
            current = ""

            for word in words:
                test = current + (" " if current else "") + word
                if self._count_tokens(test) <= self.chunk_size:
                    current = test
                else:
                    if current:
                        result.append(
                            Document(
                                page_content=current,
                                metadata={
                                    **metadata,
                                    "chunk_index": chunk_index,
                                    "split_type": "word_boundary",
                                },
                            )
                        )
                        chunk_index += 1

                    if self._count_tokens(word) > self.chunk_size:
                        # Word itself exceeds chunk size, split by tokens
                        word_tokens = self._encoding.encode(word)
                        step = max(1, self.chunk_size - self.chunk_overlap)
                        for j in range(0, len(word_tokens), step):
                            sub_tokens = word_tokens[j : j + self.chunk_size]
                            sub_text = self._encoding.decode(sub_tokens)
                            result.append(
                                Document(
                                    page_content=sub_text,
                                    metadata={
                                        **metadata,
                                        "chunk_index": chunk_index,
                                        "split_type": "hard_token",
                                    },
                                )
                            )
                            chunk_index += 1
                        current = ""
                    else:
                        current = word

            if current:
                result.append(
                    Document(
                        page_content=current,
                        metadata={
                            **metadata,
                            "chunk_index": chunk_index,
                            "split_type": "word_boundary",
                        },
                    )
                )
        else:
            # No spaces, split directly by tokens
            tokens = self._encoding.encode(text)
            step = max(1, self.chunk_size - self.chunk_overlap)
            for j in range(0, len(tokens), step):
                chunk_tokens = tokens[j : j + self.chunk_size]
                chunk_text = self._encoding.decode(chunk_tokens)
                result.append(
                    Document(
                        page_content=chunk_text,
                        metadata={
                            **metadata,
                            "chunk_index": chunk_index,
                            "split_type": "hard_token",
                        },
                    )
                )
                chunk_index += 1

        return result

    # -------------------------------------------------------------------------
    # Strict Mode Features
    # -------------------------------------------------------------------------

    def _extract_code_fences(
        self,
        content: str,
    ) -> tuple[str, str, str]:
        """Extract code fence markers and content from a code block.

        Args:
            content: The full code block content including fences.

        Returns:
            Tuple of (opening_fence, closing_fence, code_content).
        """
        lines = content.splitlines(keepends=True)
        if not lines:
            return "", "", ""

        opening_fence = ""
        closing_fence = ""
        code_lines = []

        for i, line in enumerate(lines):
            stripped = line.strip()
            if i == 0 and self._CODE_FENCE_PATTERN.match(stripped):
                opening_fence = line
            elif i == len(lines) - 1 and self._CODE_FENCE_PATTERN.match(stripped):
                closing_fence = line
            else:
                code_lines.append(line)

        return opening_fence, closing_fence, "".join(code_lines)

    def _get_effective_chunk_size(
        self,
        opening_fence: str,
        closing_fence: str,
    ) -> int:
        """Calculate effective chunk size accounting for fence overhead.

        Args:
            opening_fence: The opening fence line.
            closing_fence: The closing fence line.

        Returns:
            Effective chunk size in tokens.
        """
        fence_tokens = self._count_tokens(opening_fence + closing_fence)
        return self.chunk_size - fence_tokens

    def _split_code_block(self, chunk: Document) -> list[Document]:
        """Split a code block into smaller chunks while preserving structure.

        Strategy: Attempts to split on logical boundaries (blank lines, function/class
        definitions) first. Falls back to line-by-line splitting if necessary.
        Preserves code fences in each resulting chunk.
        """
        opening_fence, closing_fence, code_content = self._extract_code_fences(
            chunk.page_content
        )
        if not code_content:
            return [chunk]

        effective_size = self._get_effective_chunk_size(opening_fence, closing_fence)
        code_parts = re.split(r"\n\s*\n", code_content)
        wrapper = (opening_fence, closing_fence)

        def handle_oversized(part: str, meta: dict, idx: int) -> list[Document]:
            return self._split_code_by_lines(
                part, meta, idx, opening_fence, closing_fence
            )

        result = self._accumulate_parts(
            parts=code_parts,
            max_tokens=effective_size,
            separator="\n\n",
            metadata=chunk.metadata,
            start_index=0,
            split_type="code_block",
            wrapper=wrapper,
            oversized_handler=handle_oversized,
        )

        return result if result else [chunk]

    def _split_code_by_lines(
        self,
        code: str,
        metadata: dict[str, Any],
        start_index: int,
        opening_fence: str,
        closing_fence: str,
    ) -> list[Document]:
        """Split code by individual lines when logical splits aren't possible."""
        effective_size = self._get_effective_chunk_size(opening_fence, closing_fence)
        lines = code.splitlines(keepends=True)
        wrapper = (opening_fence, closing_fence)

        # For line-by-line splitting, don't strip parts and use empty separator
        result = []
        current_chunk = ""
        chunk_index = start_index

        for line in lines:
            test_content = current_chunk + line
            if self._count_tokens(test_content) <= effective_size:
                current_chunk = test_content
            else:
                if current_chunk:
                    result.append(
                        self._create_chunk(
                            current_chunk, metadata, chunk_index, "code_lines", wrapper
                        )
                    )
                    chunk_index += 1
                current_chunk = line

        if current_chunk:
            result.append(
                self._create_chunk(
                    current_chunk, metadata, chunk_index, "code_lines", wrapper
                )
            )

        return result

    def _filter_separator_chunks(self, chunks: list[Document]) -> list[Document]:
        """Filter out separator-only chunks (horizontal rules like ---).

        Removes chunks marked as type='separator' or matching horizontal rule pattern.
        Also filters empty or zero-token chunks.
        """
        result = []

        for chunk in chunks:
            content = chunk.page_content.strip()

            # Check if chunk is separator-only
            if chunk.metadata.get("type") == "separator":
                continue

            # Also check content pattern for separators
            if self._HORIZONTAL_RULE_PATTERN.match(content):
                continue

            # Skip very small whitespace-only or trivial chunks
            if not content or self._count_tokens(content) == 0:
                continue

            result.append(chunk)

        return result

    def _merge_small_chunks_by_headers(
        self,
        chunks: list[Document],
    ) -> list[Document]:
        """Merge consecutive small chunks that share the same header hierarchy.

        This reduces fragmentation by combining chunks that are below min_chunk_size
        and have identical header metadata.
        """
        if not chunks:
            return chunks

        result = []
        current_merged: Document | None = None

        def get_header_key(metadata: dict[str, Any]) -> tuple:
            """Extract header values as a tuple for comparison."""
            headers = []
            for key in sorted(metadata.keys()):
                if key.startswith("Header"):
                    headers.append((key, metadata[key]))
            return tuple(headers)

        def can_merge(chunk1: Document, chunk2: Document) -> bool:
            """Check if two chunks can be merged."""
            # Don't merge if either has special types
            type1 = chunk1.metadata.get("type")
            type2 = chunk2.metadata.get("type")

            if type1 in ("code", "table") or type2 in ("code", "table"):
                return False

            # Check if headers match
            return get_header_key(chunk1.metadata) == get_header_key(chunk2.metadata)

        def merge_chunks(chunk1: Document, chunk2: Document) -> Document:
            """Merge two chunks into one."""
            # Combine content with appropriate separator
            content1 = chunk1.page_content.rstrip()
            content2 = chunk2.page_content.lstrip()

            # Use double newline for paragraph separation
            merged_content = content1 + "\n\n" + content2

            # Merge metadata (keep headers from first, update any chunk_index)
            merged_metadata = {**chunk1.metadata}
            if "chunk_index" in merged_metadata:
                del merged_metadata["chunk_index"]
            if "merged_count" in chunk1.metadata:
                merged_metadata["merged_count"] = chunk1.metadata["merged_count"] + 1
            else:
                merged_metadata["merged_count"] = 2

            return Document(page_content=merged_content, metadata=merged_metadata)

        for chunk in chunks:
            chunk_tokens = self._count_tokens(chunk.page_content)

            if current_merged is None:
                current_merged = chunk
                continue

            current_tokens = self._count_tokens(current_merged.page_content)

            # Try to merge if current is small and headers match
            if (
                current_tokens < self.min_chunk_size
                and can_merge(current_merged, chunk)
                and current_tokens + chunk_tokens <= self.chunk_size
            ):
                current_merged = merge_chunks(current_merged, chunk)
            else:
                # Finalize current and start new
                result.append(current_merged)
                current_merged = chunk

        # Don't forget the last chunk
        if current_merged is not None:
            result.append(current_merged)

        return result


# =============================================================================
# Convenience Function
# =============================================================================


def split_markdown(
    text: str,
    chunk_size: int | None = None,
    chunk_overlap: int = 0,
    encoding_name: str = TIKTOKEN_ENCODING,
    **kwargs: Any,
) -> list[Document]:
    """Convenience function to split Markdown text.

    This is a shorthand for creating a MarkdownTextSplitter and calling
    ``split_text()`` on it.

    Args:
        text: Markdown text to split.
        chunk_size: Maximum chunk size in tokens. If ``None``,
            splits only on headers.
        chunk_overlap: Overlap between chunks in tokens.
        encoding_name: Tiktoken encoding name for token counting.
            Defaults to "cl100k_base" (GPT-4, GPT-3.5-turbo compatible).
        **kwargs: Additional arguments passed to MarkdownTextSplitter.

    Returns:
        List of Document chunks with content and metadata.

    Example:
        >>> chunks = split_markdown(text, chunk_size=500, chunk_overlap=50)
        >>> for chunk in chunks:
        ...     print(f"[{chunk.metadata.get('Header 1', 'No header')}]")
        ...     print(chunk.page_content[:100])
    """
    splitter = MarkdownTextSplitter(
        chunk_size=chunk_size,
        chunk_overlap=chunk_overlap,
        encoding_name=encoding_name,
        **kwargs,
    )
    return splitter.split_text(text)