"""Document chunking for semantic search.

Implements various chunking strategies to split documents into searchable chunks.

Performance:
- Uses optimized hierarchical merge algorithm
- Only tokenizes merged chunks, not the full document
- ~150ms for 250KB documents

Issue #1024: Entropy-aware chunking filters redundant/low-information chunks
before embedding, reducing storage costs and improving retrieval quality.
Based on SimpleMem paper (arXiv:2601.02553).
"""

import logging
import math
import re
from dataclasses import dataclass, field
from enum import StrEnum
from typing import TYPE_CHECKING, Any

logger = logging.getLogger(__name__)

if TYPE_CHECKING:
    import tiktoken as tiktoken_module

    # Removed: txtai handles this (Issue #2663)
    # from nexus.bricks.search.embeddings import EmbeddingProvider
    EmbeddingProvider = Any

    TIKTOKEN_AVAILABLE: bool
else:
    try:
        import tiktoken as tiktoken_module

        TIKTOKEN_AVAILABLE = True
    except ImportError:
        tiktoken_module = None  # type: ignore[assignment]
        TIKTOKEN_AVAILABLE = False


class ChunkStrategy(StrEnum):
    """Chunking strategy."""

    FIXED = "fixed"  # Fixed-size chunks
    SEMANTIC = "semantic"  # Semantic chunks (paragraphs/sections)
    OVERLAPPING = "overlapping"  # Overlapping fixed-size chunks
    MARKDOWN_AWARE = "markdown_aware"  # Structure-aware markdown chunks (Issue #3719)


@dataclass
class DocumentChunk:
    """A chunk of a document with source location metadata."""

    text: str
    chunk_index: int
    tokens: int
    start_offset: int  # Character offset in original document
    end_offset: int  # End character offset
    line_start: int | None = None  # Line number where chunk starts (1-indexed)
    line_end: int | None = None  # Line number where chunk ends (1-indexed)
    heading_prefix: str | None = None  # Heading hierarchy for embedding (Issue #3719)


def _offset_to_line(content: str, offset: int) -> int:
    """Convert character offset to line number (1-indexed).

    Args:
        content: Full document content
        offset: Character offset

    Returns:
        Line number (1-indexed)
    """
    if offset <= 0:
        return 1
    return content[:offset].count("\n") + 1


def _compute_line_numbers(content: str, start_offset: int, end_offset: int) -> tuple[int, int]:
    """Compute line numbers for a chunk.

    Args:
        content: Full document content
        start_offset: Start character offset
        end_offset: End character offset

    Returns:
        Tuple of (line_start, line_end)
    """
    line_start = _offset_to_line(content, start_offset)
    line_end = _offset_to_line(content, end_offset)
    return line_start, line_end


def _build_line_offsets(content: str) -> list[int]:
    """Pre-compute line start offsets for O(1) line number lookup.

    Args:
        content: Full document content

    Returns:
        List of character offsets where each line starts (0-indexed lines)
    """
    offsets = [0]
    for i, c in enumerate(content):
        if c == "\n":
            offsets.append(i + 1)
    return offsets


def _offset_to_line_fast(offset: int, line_offsets: list[int]) -> int:
    """Convert character offset to line number using pre-computed table.

    O(log n) lookup via binary search instead of O(n) counting.

    Args:
        offset: Character offset
        line_offsets: Pre-computed line offset table from _build_line_offsets()

    Returns:
        Line number (1-indexed)
    """
    import bisect

    if offset <= 0:
        return 1
    # bisect_right returns insertion point; that's the line number (1-indexed)
    return bisect.bisect_right(line_offsets, offset)


def _compute_line_numbers_fast(
    start_offset: int, end_offset: int, line_offsets: list[int]
) -> tuple[int, int]:
    """Compute line numbers using pre-computed offset table.

    Args:
        start_offset: Start character offset
        end_offset: End character offset
        line_offsets: Pre-computed line offset table

    Returns:
        Tuple of (line_start, line_end)
    """
    line_start = _offset_to_line_fast(start_offset, line_offsets)
    line_end = _offset_to_line_fast(end_offset, line_offsets)
    return line_start, line_end


# ---------------------------------------------------------------------------
# Markdown-aware chunking helpers (Issue #3719)
# ---------------------------------------------------------------------------

# Minimum tokens for a standalone chunk; smaller sections merge with siblings.
_MIN_CHUNK_TOKENS = 80
# Maximum characters for a heading prefix to prevent oversized embedding inputs.
_MAX_HEADING_PREFIX_CHARS = 200

# Regex for ATX headings (# through ######), allowing up to 3 leading spaces (CommonMark)
_HEADING_RE = re.compile(r"^ {0,3}(#{1,6})\s+(.*?)$", re.MULTILINE)
# Regex for setext headings: text line followed by === or --- underline
_SETEXT_RE = re.compile(r"^([^\n]+)\n {0,3}(=+|-+)\s*$", re.MULTILINE)
# Regex for code fence openers/closers (``` or ~~~, 3+ chars, up to 3 leading spaces)
_FENCE_RE = re.compile(r"^ {0,3}(`{3,}|~{3,})", re.MULTILINE)


@dataclass(slots=True)
class _HeadingInfo:
    """A heading found in markdown content (search-brick local)."""

    heading: str
    depth: int
    char_offset: int  # character offset of the heading line


def _parse_headings_fence_aware(content: str) -> list[_HeadingInfo]:
    """Find ATX headings while ignoring those inside code fences.

    Builds a set of fenced character ranges, then filters headings
    that fall within them. Handles nested/mismatched fences gracefully.

    Returns:
        List of headings in document order.
    """
    # Build set of fenced ranges using a flat state machine.
    # CommonMark: a code fence is closed only by a fence line with the
    # same character and >= the opener length.  Shorter fence-like lines
    # inside an open fence are just literal content (no nesting).
    fenced: list[tuple[int, int]] = []
    active_fence: tuple[str, int, int] | None = None  # (char, length, start_offset)
    for m in _FENCE_RE.finditer(content):
        marker = m.group(1)
        fence_char = marker[0]
        fence_len = len(marker)
        if active_fence is not None:
            # Inside a fence — only close if same char and >= length
            if fence_char == active_fence[0] and fence_len >= active_fence[1]:
                fenced.append((active_fence[2], m.end()))
                active_fence = None
            # Otherwise: literal content inside the fence, ignore
        else:
            # Not inside a fence — this is an opener
            active_fence = (fence_char, fence_len, m.start())
    # Unclosed fence: treat opener to EOF as fenced
    if active_fence is not None:
        fenced.append((active_fence[2], len(content)))

    def _is_fenced(offset: int) -> bool:
        return any(s <= offset < e for s, e in fenced)

    headings: list[_HeadingInfo] = []

    # ATX headings (# through ######)
    for m in _HEADING_RE.finditer(content):
        if _is_fenced(m.start()):
            continue
        depth = len(m.group(1))
        text = m.group(2).strip()
        headings.append(_HeadingInfo(heading=text, depth=depth, char_offset=m.start()))

    # Detect YAML frontmatter range to exclude from setext detection.
    # Frontmatter `---` / `key: value` / `---` looks like setext H2.
    # Require at least one YAML key: value line to avoid thematic breaks.
    fm_end = 0
    stripped = content.lstrip()
    leading_ws = len(content) - len(stripped)
    if stripped.startswith("---"):
        close = stripped.find("\n---", 3)
        if close != -1:
            fm_body = stripped[4:close]
            has_yaml = any(":" in line for line in fm_body.split("\n") if line.strip())
            if has_yaml:
                fm_end = leading_ws + close + 4
            if fm_end < len(content) and content[fm_end] == "\n":
                fm_end += 1

    # Setext headings (underline with === or ---)
    for m in _SETEXT_RE.finditer(content):
        if _is_fenced(m.start()):
            continue
        # Skip matches inside frontmatter
        if m.start() < fm_end:
            continue
        text = m.group(1).strip()
        # Skip if the text line looks like a YAML frontmatter delimiter
        if text == "---" or not text:
            continue
        depth = 1 if m.group(2)[0] == "=" else 2
        headings.append(_HeadingInfo(heading=text, depth=depth, char_offset=m.start()))

    # Sort by position (ATX and setext may interleave)
    headings.sort(key=lambda h: h.char_offset)

    return headings


@dataclass
class _MdSegment:
    """Internal segment for markdown-aware chunking."""

    text: str
    char_start: int
    char_end: int
    heading_prefix: str
    tokens: int


def _safe_line_char_offset(line_offsets: list[int], line_num: int, content_len: int) -> int:
    """Get the character offset for a 0-indexed line number, clamped safely."""
    if line_num >= len(line_offsets):
        return content_len
    return line_offsets[line_num]


def _extract_file_name(file_path: str) -> str:
    """Extract just the filename from a path."""
    if "/" in file_path:
        return file_path.rsplit("/", 1)[-1]
    return file_path


def build_heading_hierarchy(
    sections: list[Any],
    section_index: int,
    file_name: str = "",
) -> str:
    """Build heading hierarchy prefix for embedding context.

    Walks backward through sections to find parent headings,
    building a path like ``[file.md > Auth > JWT Tokens]``.

    Args:
        sections: List of SectionInfo objects (from md_structure).
        section_index: Index of the current section.
        file_name: Filename to include at the start of the path.

    Returns:
        Bracket-notation heading path string.
    """
    if section_index < 0 or section_index >= len(sections):
        return f"[{file_name}]" if file_name else ""

    target = sections[section_index]
    parents: list[str] = []
    current_depth = target.depth

    for i in range(section_index - 1, -1, -1):
        if sections[i].depth < current_depth:
            parents.append(sections[i].heading)
            current_depth = sections[i].depth
            if current_depth <= 1:
                break

    parents.reverse()
    parts = parents + [target.heading]
    path = " > ".join(parts)
    result = f"[{file_name} > {path}]" if file_name else f"[{path}]"

    # Truncate to prevent oversized embedding inputs.  Keep the
    # file name and the deepest heading levels (most specific context).
    if len(result) > _MAX_HEADING_PREFIX_CHARS:
        tail = " > ".join(parts[-2:]) if len(parts) >= 2 else parts[-1]
        result = f"[{file_name} > ... > {tail}]" if file_name else f"[... > {tail}]"

    return result


def _merge_small_segments(
    segments: list[_MdSegment],
    min_tokens: int = _MIN_CHUNK_TOKENS,
    max_tokens: int = 1024,
) -> list[_MdSegment]:
    """Merge consecutive tiny segments using single-pass greedy algorithm.

    Segments with fewer than *min_tokens* are accumulated and merged with
    the next segment(s) until the accumulator exceeds *min_tokens* or
    adding more would exceed *max_tokens*.

    Args:
        segments: Ordered list of segments.
        min_tokens: Minimum token threshold for a standalone segment.
        max_tokens: Maximum token budget for a merged segment.

    Returns:
        New list of segments with tiny segments merged.
    """
    if not segments:
        return []

    merged: list[_MdSegment] = []
    acc_texts: list[str] = []
    acc_start: int = 0
    acc_end: int = 0
    acc_prefix: str = ""
    acc_tokens: int = 0

    def _emit_accumulator() -> None:
        if acc_texts:
            merged.append(
                _MdSegment(
                    text="\n\n".join(acc_texts),
                    char_start=acc_start,
                    char_end=acc_end,
                    heading_prefix=acc_prefix,
                    tokens=acc_tokens,
                )
            )

    for seg in segments:
        if not acc_texts:
            # Start new accumulator
            acc_texts = [seg.text]
            acc_start = seg.char_start
            acc_end = seg.char_end
            acc_prefix = seg.heading_prefix
            acc_tokens = seg.tokens

            # If segment is big enough on its own, emit immediately
            if seg.tokens >= min_tokens:
                _emit_accumulator()
                acc_texts = []
                acc_tokens = 0
        else:
            # Try to add to accumulator
            combined = acc_tokens + seg.tokens
            if combined <= max_tokens:
                acc_texts.append(seg.text)
                acc_end = seg.char_end
                acc_tokens = combined

                # If big enough now, emit
                if acc_tokens >= min_tokens:
                    _emit_accumulator()
                    acc_texts = []
                    acc_tokens = 0
            else:
                # Would exceed budget — emit accumulator, start fresh
                _emit_accumulator()
                acc_texts = [seg.text]
                acc_start = seg.char_start
                acc_end = seg.char_end
                acc_prefix = seg.heading_prefix
                acc_tokens = seg.tokens

                if seg.tokens >= min_tokens:
                    _emit_accumulator()
                    acc_texts = []
                    acc_tokens = 0

    # Emit any remaining
    _emit_accumulator()
    return merged


class DocumentChunker:
    """Document chunker for semantic search.

    Uses optimized hierarchical merge algorithm that only tokenizes
    merged chunks, not the full document. This is faster than tokenizing
    the entire document upfront.
    """

    encoding: Any  # tiktoken.Encoding or None

    def __init__(
        self,
        chunk_size: int = 1024,
        overlap_size: int = 128,
        strategy: ChunkStrategy = ChunkStrategy.FIXED,
        encoding_name: str = "cl100k_base",
    ):
        """Initialize document chunker.

        Args:
            chunk_size: Target chunk size in tokens
            overlap_size: Overlap size in tokens for overlapping strategy
            strategy: Chunking strategy to use
            encoding_name: Tiktoken encoding name (default: cl100k_base for GPT-4/Claude)
        """
        self.chunk_size = chunk_size
        self.overlap_size = overlap_size
        self.strategy = strategy
        self.encoding_name = encoding_name

        # Initialize tokenizer
        if tiktoken_module is not None:
            try:
                self.encoding = tiktoken_module.get_encoding(encoding_name)
            except Exception:
                # Fallback to approximate tokenization
                self.encoding = None
        else:
            self.encoding = None

    def _count_tokens(self, text: str) -> int:
        """Count tokens in text.

        Args:
            text: Text to count tokens in

        Returns:
            Number of tokens
        """
        if self.encoding is not None:
            return len(self.encoding.encode(text))
        else:
            # Rough approximation: 1 token ≈ 4 characters
            return len(text) // 4

    def _approx_tokens(self, text: str) -> int:
        """Fast approximate token count (1 token ≈ 4 chars).

        Args:
            text: Text to estimate tokens for

        Returns:
            Approximate number of tokens
        """
        return len(text) // 4

    def _fits_in_chunk(self, text: str) -> bool:
        """Check if text fits within chunk_size using smart approximation.

        Uses fast approximation for clear cases, only calls expensive
        tiktoken for boundary cases. This reduces tokenize calls by ~90%.

        Args:
            text: Text to check

        Returns:
            True if text fits within chunk_size
        """
        approx = self._approx_tokens(text)

        # Clear cases - skip expensive tokenization
        if approx < self.chunk_size * 0.7:
            return True  # Definitely fits
        if approx > self.chunk_size * 1.5:
            return False  # Definitely too big

        # Gray zone - use accurate count
        return self._count_tokens(text) <= self.chunk_size

    def _exceeds_chunk_size(self, text: str) -> bool:
        """Check if text exceeds chunk_size using smart approximation.

        Args:
            text: Text to check

        Returns:
            True if text exceeds chunk_size
        """
        return not self._fits_in_chunk(text)

    def chunk(
        self,
        content: str,
        file_path: str = "",
        compute_lines: bool = True,
    ) -> list[DocumentChunk]:
        """Chunk document into searchable chunks.

        Args:
            content: Document content to chunk
            file_path: Path to the file (used for file-type specific chunking)
            compute_lines: If True, compute line numbers for each chunk (default: True)

        Returns:
            List of document chunks with line numbers if compute_lines=True
        """
        if self.strategy == ChunkStrategy.FIXED:
            chunks = self._chunk_fixed(content)
        elif self.strategy == ChunkStrategy.SEMANTIC:
            chunks = self._chunk_semantic(content, file_path)
        elif self.strategy == ChunkStrategy.OVERLAPPING:
            chunks = self._chunk_overlapping(content)
        elif self.strategy == ChunkStrategy.MARKDOWN_AWARE:
            chunks = self._chunk_markdown_aware(content, file_path)
        else:
            raise ValueError(f"Unknown chunking strategy: {self.strategy}")

        # Add line numbers to chunks using pre-computed offset table (O(log n) per chunk)
        if compute_lines and chunks:
            line_offsets = _build_line_offsets(content)
            for chunk in chunks:
                line_start, line_end = _compute_line_numbers_fast(
                    chunk.start_offset, chunk.end_offset, line_offsets
                )
                chunk.line_start = line_start
                chunk.line_end = line_end

        return chunks

    def _chunk_fixed(self, content: str) -> list[DocumentChunk]:
        """Chunk document into fixed-size chunks using recursive splitting.

        Uses a recursive approach similar to LangChain's RecursiveCharacterTextSplitter:
        tries to split at semantic boundaries (paragraphs, sentences, words) while
        respecting the chunk_size limit.

        This is much more efficient than per-word tokenization as it only counts
        tokens on merged chunks, not individual words.

        Args:
            content: Document content

        Returns:
            List of chunks
        """
        # Separators in order of preference (most semantic to least)
        separators = ["\n\n", "\n", ". ", ", ", " "]
        return self._recursive_split(content, separators, start_offset=0)

    def _recursive_split(
        self, text: str, separators: list[str], start_offset: int
    ) -> list[DocumentChunk]:
        """Recursively split text using separators until chunks fit within chunk_size.

        Args:
            text: Text to split
            separators: List of separators to try, in order of preference
            start_offset: Starting offset in the original document

        Returns:
            List of chunks
        """
        if not text.strip():
            return []

        # Check if text already fits using smart approximation (avoids tokenization)
        if self._fits_in_chunk(text):
            # Only count tokens when we create the final chunk
            text_tokens = self._count_tokens(text)
            return [
                DocumentChunk(
                    text=text,
                    chunk_index=0,  # Will be renumbered later
                    tokens=text_tokens,
                    start_offset=start_offset,
                    end_offset=start_offset + len(text),
                )
            ]

        # Try each separator
        for sep in separators:
            if sep in text:
                return self._split_and_merge(text, sep, separators, start_offset)

        # Last resort: split by characters (shouldn't normally reach here)
        return self._split_by_chars(text, start_offset)

    def _split_and_merge(
        self, text: str, separator: str, separators: list[str], start_offset: int
    ) -> list[DocumentChunk]:
        """Split text by separator and merge segments to fit chunk_size.

        Args:
            text: Text to split
            separator: Separator to use
            separators: All separators for recursive splitting
            start_offset: Starting offset in the original document

        Returns:
            List of chunks
        """
        chunks: list[DocumentChunk] = []
        parts = text.split(separator)

        current_parts: list[str] = []
        current_approx_tokens = 0  # Use approximation for running total
        current_offset = start_offset
        sep_approx_tokens = self._approx_tokens(separator)

        for part in parts:
            if not part:
                continue

            part_approx_tokens = self._approx_tokens(part)

            # If single part exceeds chunk_size, recursively split it
            # Use smart check that may call tiktoken for boundary cases
            if self._exceeds_chunk_size(part):
                # First, finalize current chunk if any
                if current_parts:
                    chunk_text = separator.join(current_parts)
                    # Count actual tokens only when creating final chunk
                    actual_tokens = self._count_tokens(chunk_text)
                    chunks.append(
                        DocumentChunk(
                            text=chunk_text,
                            chunk_index=len(chunks),
                            tokens=actual_tokens,
                            start_offset=current_offset,
                            end_offset=current_offset + len(chunk_text),
                        )
                    )
                    current_offset += len(chunk_text) + len(separator)
                    current_parts = []
                    current_approx_tokens = 0

                # Recursively split the large part with remaining separators
                remaining_seps = separators[separators.index(separator) + 1 :]
                if remaining_seps:
                    sub_chunks = self._recursive_split(part, remaining_seps, current_offset)
                    for sub_chunk in sub_chunks:
                        sub_chunk.chunk_index = len(chunks)
                        chunks.append(sub_chunk)
                    current_offset += len(part) + len(separator)
                else:
                    # No more separators, split by chars
                    sub_chunks = self._split_by_chars(part, current_offset)
                    for sub_chunk in sub_chunks:
                        sub_chunk.chunk_index = len(chunks)
                        chunks.append(sub_chunk)
                    current_offset += len(part) + len(separator)
                continue

            # Check if adding this part would exceed chunk_size (use approximation)
            sep_to_add = sep_approx_tokens if current_parts else 0
            if (
                current_approx_tokens + sep_to_add + part_approx_tokens > self.chunk_size
                and current_parts
            ):
                # Finalize current chunk - count actual tokens now
                chunk_text = separator.join(current_parts)
                actual_tokens = self._count_tokens(chunk_text)
                chunks.append(
                    DocumentChunk(
                        text=chunk_text,
                        chunk_index=len(chunks),
                        tokens=actual_tokens,
                        start_offset=current_offset,
                        end_offset=current_offset + len(chunk_text),
                    )
                )
                current_offset += len(chunk_text) + len(separator)
                current_parts = []
                current_approx_tokens = 0

            current_parts.append(part)
            current_approx_tokens += part_approx_tokens + (
                sep_approx_tokens if len(current_parts) > 1 else 0
            )

        # Add remaining parts as final chunk
        if current_parts:
            chunk_text = separator.join(current_parts)
            final_tokens = self._count_tokens(chunk_text)
            chunks.append(
                DocumentChunk(
                    text=chunk_text,
                    chunk_index=len(chunks),
                    tokens=final_tokens,
                    start_offset=current_offset,
                    end_offset=current_offset + len(chunk_text),
                )
            )

        # Renumber chunk indices
        for i, chunk in enumerate(chunks):
            chunk.chunk_index = i

        return chunks

    def _split_by_chars(self, text: str, start_offset: int) -> list[DocumentChunk]:
        """Split text by characters when no separator works.

        Args:
            text: Text to split
            start_offset: Starting offset in the original document

        Returns:
            List of chunks
        """
        chunks: list[DocumentChunk] = []

        # Estimate chars per token (roughly 4 chars per token)
        chars_per_chunk = self.chunk_size * 4
        current_offset = start_offset

        for i in range(0, len(text), chars_per_chunk):
            chunk_text = text[i : i + chars_per_chunk]
            tokens = self._count_tokens(chunk_text)
            chunks.append(
                DocumentChunk(
                    text=chunk_text,
                    chunk_index=len(chunks),
                    tokens=tokens,
                    start_offset=current_offset,
                    end_offset=current_offset + len(chunk_text),
                )
            )
            current_offset += len(chunk_text)

        return chunks

    def _chunk_semantic(self, content: str, file_path: str) -> list[DocumentChunk]:
        """Chunk document semantically (by paragraphs/sections).

        Args:
            content: Document content
            file_path: Path to the file

        Returns:
            List of chunks
        """
        # Determine file type
        if file_path.endswith((".md", ".markdown")):
            return self._chunk_markdown(content)
        elif file_path.endswith((".py", ".js", ".ts", ".java", ".go", ".rs")):
            return self._chunk_code(content)
        else:
            return self._chunk_paragraphs(content)

    def _chunk_markdown(self, content: str) -> list[DocumentChunk]:
        """Chunk markdown by sections.

        Args:
            content: Markdown content

        Returns:
            List of chunks
        """
        chunks: list[DocumentChunk] = []
        # Split by headings
        sections = re.split(r"\n(?=#{1,6}\s)", content)
        current_offset = 0

        for section in sections:
            if not section.strip():
                continue

            tokens = self._count_tokens(section)

            # If section is too large, split it further
            if tokens > self.chunk_size:
                sub_chunks = self._chunk_paragraphs(section)
                for sub_chunk in sub_chunks:
                    sub_chunk.chunk_index = len(chunks)
                    sub_chunk.start_offset += current_offset
                    sub_chunk.end_offset += current_offset
                    chunks.append(sub_chunk)
            else:
                chunks.append(
                    DocumentChunk(
                        text=section,
                        chunk_index=len(chunks),
                        tokens=tokens,
                        start_offset=current_offset,
                        end_offset=current_offset + len(section),
                    )
                )

            current_offset += len(section) + 1  # +1 for newline

        return chunks

    def _chunk_code(self, content: str) -> list[DocumentChunk]:
        """Chunk code by functions/classes.

        Args:
            content: Code content

        Returns:
            List of chunks
        """
        # For now, use paragraph-based chunking
        # TODO: Implement AST-based chunking for better code structure preservation
        return self._chunk_paragraphs(content)

    def _chunk_paragraphs(self, content: str) -> list[DocumentChunk]:
        """Chunk by paragraphs.

        Args:
            content: Content to chunk

        Returns:
            List of chunks
        """
        chunks: list[DocumentChunk] = []
        paragraphs = content.split("\n\n")
        current_chunk: list[str] = []
        current_tokens = 0
        current_offset = 0

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

            para_tokens = self._count_tokens(para)

            if current_tokens + para_tokens > self.chunk_size and current_chunk:
                # Create chunk
                chunk_text = "\n\n".join(current_chunk)
                chunk_end = current_offset + len(chunk_text)
                chunks.append(
                    DocumentChunk(
                        text=chunk_text,
                        chunk_index=len(chunks),
                        tokens=current_tokens,
                        start_offset=current_offset,
                        end_offset=chunk_end,
                    )
                )
                current_offset = chunk_end + 2  # +2 for \n\n
                current_chunk = []
                current_tokens = 0

            current_chunk.append(para)
            current_tokens += para_tokens

        # Add remaining chunk
        if current_chunk:
            chunk_text = "\n\n".join(current_chunk)
            chunk_end = current_offset + len(chunk_text)
            chunks.append(
                DocumentChunk(
                    text=chunk_text,
                    chunk_index=len(chunks),
                    tokens=current_tokens,
                    start_offset=current_offset,
                    end_offset=chunk_end,
                )
            )

        return chunks

    def _chunk_overlapping(self, content: str) -> list[DocumentChunk]:
        """Chunk document with overlapping windows.

        Args:
            content: Document content

        Returns:
            List of chunks
        """
        chunks: list[DocumentChunk] = []
        words = content.split()
        current_offset = 0

        # Calculate step size (chunk_size - overlap_size)
        step_size = self.chunk_size - self.overlap_size

        i = 0
        while i < len(words):
            # Take chunk_size words
            chunk_words = words[i : i + self.chunk_size]
            chunk_text = " ".join(chunk_words)
            tokens = self._count_tokens(chunk_text)

            chunks.append(
                DocumentChunk(
                    text=chunk_text,
                    chunk_index=len(chunks),
                    tokens=tokens,
                    start_offset=current_offset,
                    end_offset=current_offset + len(chunk_text),
                )
            )

            # Move forward by step_size
            i += step_size
            current_offset += len(" ".join(words[i - step_size : i])) + 1

        return chunks

    def _chunk_markdown_aware(self, content: str, file_path: str) -> list[DocumentChunk]:
        """Chunk markdown using structural index with heading context prepend.

        Issue #3719: Uses the md_structure parser for heading boundaries,
        keeps code blocks and tables atomic, merges tiny sections, and
        attaches heading hierarchy prefix for embedding enrichment.

        Falls back to ``_chunk_semantic()`` for non-markdown files and to
        ``_chunk_fixed()`` when the parser produces no sections.

        Args:
            content: Document content (string).
            file_path: Path to the file.

        Returns:
            List of chunks with ``heading_prefix`` set.
        """
        # Non-markdown files: delegate to semantic strategy
        if not file_path.endswith((".md", ".markdown")):
            return self._chunk_semantic(content, file_path)

        # Parse headings with code-fence awareness (no cross-brick import)
        headings = _parse_headings_fence_aware(content)

        # Fallback if no headings found
        if not headings:
            return self._chunk_fixed(content)

        file_name = _extract_file_name(file_path)

        # ── Build raw segments from heading positions ────────────────
        segments: list[_MdSegment] = []

        # Detect YAML frontmatter (--- ... ---).
        # Validates that the block contains YAML-like content (at least
        # one "key:" line) to avoid false positives on thematic breaks.
        fm_end_offset = 0
        stripped = content.lstrip()
        leading_ws = len(content) - len(stripped)
        if stripped.startswith("---"):
            close = stripped.find("\n---", 3)
            if close != -1:
                fm_body = stripped[4:close]  # content between --- markers
                # Require at least one YAML key: value line
                has_yaml = any(":" in line for line in fm_body.split("\n") if line.strip())
                if has_yaml:
                    fm_end_offset = leading_ws + close + 4  # past closing ---
                    # Skip trailing newline after closing ---
                    if fm_end_offset < len(content) and content[fm_end_offset] == "\n":
                        fm_end_offset += 1
                    fm_text = content[:fm_end_offset].strip()
                    if fm_text:
                        segments.append(
                            _MdSegment(
                                text=fm_text,
                                char_start=0,
                                char_end=fm_end_offset,
                                heading_prefix=f"[{file_name} > frontmatter]",
                                tokens=self._count_tokens(fm_text),
                            )
                        )

        # Preamble (content between frontmatter end and first heading)
        first_heading_offset = headings[0].char_offset
        preamble_text = content[fm_end_offset:first_heading_offset].strip()
        if preamble_text:
            segments.append(
                _MdSegment(
                    text=preamble_text,
                    char_start=fm_end_offset,
                    char_end=first_heading_offset,
                    heading_prefix=f"[{file_name}]",
                    tokens=self._count_tokens(preamble_text),
                )
            )

        # Each heading → a segment (heading to next heading = non-overlapping)
        for i, h in enumerate(headings):
            sec_start = h.char_offset
            sec_end = headings[i + 1].char_offset if i + 1 < len(headings) else len(content)
            sec_text = content[sec_start:sec_end]
            if not sec_text.strip():
                continue

            heading_prefix = build_heading_hierarchy(headings, i, file_name)
            segments.append(
                _MdSegment(
                    text=sec_text,
                    char_start=sec_start,
                    char_end=sec_end,
                    heading_prefix=heading_prefix,
                    tokens=self._count_tokens(sec_text),
                )
            )

        if not segments:
            return self._chunk_fixed(content)

        # ── Separate frontmatter/preamble from heading segments ─────
        # Frontmatter and preamble are emitted as standalone chunks —
        # they must not be merged with heading segments, otherwise the
        # merged chunk inherits the wrong heading_prefix.
        standalone: list[_MdSegment] = []
        heading_segments: list[_MdSegment] = []
        for seg in segments:
            if (
                seg.heading_prefix.endswith("frontmatter]")
                or seg.heading_prefix == f"[{file_name}]"
            ):
                standalone.append(seg)
            else:
                heading_segments.append(seg)

        # ── Merge tiny heading segments (single-pass greedy, O(n)) ──
        merged = standalone + _merge_small_segments(
            heading_segments,
            min_tokens=_MIN_CHUNK_TOKENS,
            max_tokens=self.chunk_size,
        )

        # ── Build final chunks, splitting oversized segments ────────
        chunks: list[DocumentChunk] = []
        for seg in merged:
            prefix_tokens = (
                self._count_tokens(seg.heading_prefix + " ") if seg.heading_prefix else 0
            )
            effective_budget = max(self.chunk_size - prefix_tokens, self.chunk_size // 2)

            if seg.tokens <= effective_budget:
                chunks.append(
                    DocumentChunk(
                        text=seg.text,
                        chunk_index=len(chunks),
                        tokens=seg.tokens,
                        start_offset=seg.char_start,
                        end_offset=seg.char_end,
                        heading_prefix=seg.heading_prefix,
                    )
                )
            else:
                # Oversized section: split around code fences to keep
                # fenced blocks intact.  Splits text into alternating
                # prose / fenced-block fragments, then sub-splits only
                # the prose parts with the generic fixed splitter.
                sub_chunks = self._split_preserving_fences(
                    seg.text,
                    seg.char_start,
                    effective_budget,
                )
                for sc in sub_chunks:
                    sc.heading_prefix = seg.heading_prefix
                    chunks.append(sc)

        # Renumber chunk indices
        for i, chunk in enumerate(chunks):
            chunk.chunk_index = i

        logger.info(
            "[MARKDOWN-AWARE] path=%s headings=%d chunks=%d",
            file_path,
            len(headings),
            len(chunks),
        )

        return chunks

    def _split_preserving_fences(
        self,
        text: str,
        base_offset: int,
        budget: int,
    ) -> list[DocumentChunk]:
        """Split oversized section while keeping fenced code blocks intact.

        Splits *text* into alternating prose / fenced-block fragments.
        Prose fragments are sub-split with the fixed splitter at *budget*.
        Fenced blocks up to 1.5× budget are emitted whole; larger blocks
        are sub-split as a last resort.
        """
        # Find fenced ranges within this text (flat state machine,
        # same logic as _parse_headings_fence_aware).
        fenced_ranges: list[tuple[int, int]] = []
        active: tuple[str, int, int] | None = None  # (char, length, start)
        for m in _FENCE_RE.finditer(text):
            marker = m.group(1)
            fence_char = marker[0]
            fence_len = len(marker)
            if active is not None:
                if fence_char == active[0] and fence_len >= active[1]:
                    fenced_ranges.append((active[2], m.end()))
                    active = None
            else:
                active = (fence_char, fence_len, m.start())
        if active is not None:
            fenced_ranges.append((active[2], len(text)))

        if not fenced_ranges:
            # No fences — safe to use generic splitter
            return self._sub_split_prose(text, base_offset, budget)

        # Build alternating fragments: prose, fence, prose, fence, ...
        chunks: list[DocumentChunk] = []
        pos = 0
        allowance = int(budget * 1.5)

        for fence_start, fence_end in fenced_ranges:
            # Prose before this fence
            if pos < fence_start:
                prose = text[pos:fence_start]
                if prose.strip():
                    chunks.extend(self._sub_split_prose(prose, base_offset + pos, budget))

            # The fenced block itself
            block_text = text[fence_start:fence_end]
            block_tokens = self._count_tokens(block_text)
            if block_tokens <= allowance:
                # Keep intact
                chunks.append(
                    DocumentChunk(
                        text=block_text,
                        chunk_index=0,
                        tokens=block_tokens,
                        start_offset=base_offset + fence_start,
                        end_offset=base_offset + fence_end,
                    )
                )
            else:
                # Truly enormous block — last resort split
                chunks.extend(self._sub_split_prose(block_text, base_offset + fence_start, budget))
            pos = fence_end

        # Trailing prose after last fence
        if pos < len(text):
            trailing = text[pos:]
            if trailing.strip():
                chunks.extend(self._sub_split_prose(trailing, base_offset + pos, budget))

        return chunks

    def _sub_split_prose(
        self,
        text: str,
        base_offset: int,
        budget: int,
    ) -> list[DocumentChunk]:
        """Sub-split a prose fragment using a fixed-budget chunker."""
        sub_chunker = DocumentChunker(
            chunk_size=budget,
            overlap_size=self.overlap_size,
            strategy=ChunkStrategy.FIXED,
            encoding_name=self.encoding_name,
        )
        sub_chunker.encoding = self.encoding
        sub_chunks = sub_chunker._chunk_fixed(text)
        for sc in sub_chunks:
            sc.start_offset += base_offset
            sc.end_offset += base_offset
        return sub_chunks


@dataclass
class EntropyFilterResult:
    """Result of entropy-aware filtering.

    Issue #1024: Tracks filtering statistics for benchmarking.
    """

    chunks: list[DocumentChunk]
    original_count: int
    filtered_count: int
    scores: list[float] = field(default_factory=list)

    @property
    def reduction_percent(self) -> float:
        """Calculate chunk reduction percentage."""
        if self.original_count == 0:
            return 0.0
        return ((self.original_count - self.filtered_count) / self.original_count) * 100


def _cosine_similarity(vec1: list[float], vec2: list[float]) -> float:
    """Calculate cosine similarity between two vectors.

    Args:
        vec1: First vector
        vec2: Second vector

    Returns:
        Cosine similarity (-1 to 1)
    """
    if len(vec1) != len(vec2):
        return 0.0

    dot_product = sum(a * b for a, b in zip(vec1, vec2, strict=False))
    norm1 = math.sqrt(sum(a * a for a in vec1))
    norm2 = math.sqrt(sum(b * b for b in vec2))

    if norm1 == 0 or norm2 == 0:
        return 0.0

    return dot_product / (norm1 * norm2)


class EntropyAwareChunker:
    """Entropy-aware chunker that filters redundant/low-information chunks.

    Issue #1024: Implements SimpleMem's entropy-aware filtering formula:
        H(W_t) = α · |ℰ_new|/|W_t| + (1-α) · (1 - cos(E(W_t), E(H_prev)))

    Where:
    - |ℰ_new|: New entities in current chunk not seen in history
    - |W_t|: Current chunk size (word count)
    - cos(E(W_t), E(H_prev)): Semantic similarity to previous chunks
    - α: Balance between entity novelty and semantic divergence

    Reference: SimpleMem paper (arXiv:2601.02553)

    Usage:
        chunker = EntropyAwareChunker(
            redundancy_threshold=0.35,
            alpha=0.5,
            embedding_provider=embedding_provider,
        )
        result = await chunker.chunk_with_filtering(content)
        print(f"Filtered {result.reduction_percent:.1f}% redundant chunks")
    """

    # Regex patterns for lightweight entity extraction
    # Matches capitalized words/phrases, numbers, URLs, emails, etc.
    ENTITY_PATTERNS = [
        # Capitalized words (proper nouns) - at least 2 chars
        r"\b[A-Z][a-z]+(?:\s+[A-Z][a-z]+)*\b",
        # ALL CAPS acronyms (2+ chars)
        r"\b[A-Z]{2,}\b",
        # Numbers with context (dates, versions, IDs)
        r"\b\d{4}[-/]\d{2}[-/]\d{2}\b",  # Dates
        r"\bv?\d+\.\d+(?:\.\d+)?\b",  # Versions
        r"\b\d+(?:,\d{3})*(?:\.\d+)?\b",  # Large numbers
        # Technical identifiers
        r"\b[a-z_][a-z0-9_]*(?:_[a-z0-9]+)+\b",  # snake_case
        r"\b[a-z]+(?:[A-Z][a-z]+)+\b",  # camelCase
        # URLs and emails (simplified)
        r"https?://[^\s]+",
        r"\b[\w.-]+@[\w.-]+\.\w+\b",
    ]

    def __init__(
        self,
        redundancy_threshold: float = 0.35,
        alpha: float = 0.5,
        embedding_provider: "EmbeddingProvider | None" = None,
        base_chunker: DocumentChunker | None = None,
        history_window: int = 5,
    ):
        """Initialize entropy-aware chunker.

        Args:
            redundancy_threshold: Filter chunks with score below this (default: 0.35)
                                  SimpleMem paper uses τ_redundant = 0.35
            alpha: Balance between entity novelty (α) and semantic novelty (1-α)
                   Default 0.5 gives equal weight to both signals
            embedding_provider: Provider for semantic embeddings (optional)
                                If None, only entity novelty is used (α=1.0 effectively)
            base_chunker: Base chunker to use for initial chunking
                          Default: DocumentChunker with SEMANTIC strategy
            history_window: Number of previous chunks to consider for novelty
                            Default: 5 chunks
        """
        self.redundancy_threshold = redundancy_threshold
        self.alpha = alpha
        self.embedding_provider = embedding_provider
        self.history_window = history_window

        # Use provided chunker or create default
        self.base_chunker = base_chunker or DocumentChunker(
            chunk_size=1024,
            strategy=ChunkStrategy.SEMANTIC,
            overlap_size=128,
        )

        # Compile regex patterns for efficiency
        self._entity_pattern = re.compile(
            "|".join(f"({p})" for p in self.ENTITY_PATTERNS),
            re.MULTILINE,
        )

    def extract_entities(self, text: str) -> set[str]:
        """Extract entities from text using lightweight regex patterns.

        This is a fast, dependency-free alternative to full NER.
        Captures proper nouns, technical terms, dates, versions, etc.

        Args:
            text: Text to extract entities from

        Returns:
            Set of extracted entity strings (normalized to lowercase)
        """
        entities = set()

        for match in self._entity_pattern.finditer(text):
            entity = match.group(0).strip()
            if entity and len(entity) >= 2:
                # Normalize to lowercase for comparison
                entities.add(entity.lower())

        return entities

    def _entity_novelty_score(
        self,
        chunk_entities: set[str],
        history_entities: set[str],
        chunk_word_count: int,
    ) -> float:
        """Calculate entity novelty score.

        Formula: |ℰ_new| / |W_t|

        Args:
            chunk_entities: Entities in current chunk
            history_entities: Entities seen in previous chunks
            chunk_word_count: Number of words in current chunk

        Returns:
            Entity novelty score (0.0 to 1.0)
        """
        if chunk_word_count == 0:
            return 0.0

        new_entities = chunk_entities - history_entities
        # Normalize by word count, cap at 1.0
        return min(len(new_entities) / max(chunk_word_count, 1), 1.0)

    def _semantic_novelty_score(
        self,
        chunk_embedding: list[float],
        history_embeddings: list[list[float]],
    ) -> float:
        """Calculate semantic novelty score.

        Formula: 1 - max(cos(E(W_t), E(H_i)) for H_i in history)

        Args:
            chunk_embedding: Embedding of current chunk
            history_embeddings: Embeddings of previous chunks

        Returns:
            Semantic novelty score (0.0 to 1.0)
        """
        if not history_embeddings:
            return 1.0  # First chunk is fully novel

        # Find maximum similarity to any history chunk
        max_similarity = max(
            _cosine_similarity(chunk_embedding, hist_emb) for hist_emb in history_embeddings
        )

        # Novelty = 1 - similarity
        return 1.0 - max(0.0, min(max_similarity, 1.0))

    def information_score(
        self,
        chunk_entities: set[str],
        history_entities: set[str],
        chunk_word_count: int,
        chunk_embedding: list[float] | None = None,
        history_embeddings: list[list[float]] | None = None,
    ) -> float:
        """Calculate information density score for a chunk.

        Formula: H(W_t) = α · entity_novelty + (1-α) · semantic_novelty

        Args:
            chunk_entities: Entities in current chunk
            history_entities: Entities seen in previous chunks
            chunk_word_count: Number of words in current chunk
            chunk_embedding: Embedding of current chunk (optional)
            history_embeddings: Embeddings of previous chunks (optional)

        Returns:
            Information score (0.0 to 1.0), higher = more informative
        """
        entity_score = self._entity_novelty_score(
            chunk_entities, history_entities, chunk_word_count
        )

        # If no embeddings, use only entity score
        if chunk_embedding is None or history_embeddings is None:
            return entity_score

        semantic_score = self._semantic_novelty_score(chunk_embedding, history_embeddings)

        # Weighted combination
        return self.alpha * entity_score + (1 - self.alpha) * semantic_score

    async def chunk_with_filtering(
        self,
        content: str,
        file_path: str = "",
        compute_lines: bool = True,
    ) -> EntropyFilterResult:
        """Chunk content and filter redundant chunks.

        Args:
            content: Document content to chunk
            file_path: Path to file (for file-type specific chunking)
            compute_lines: If True, compute line numbers for chunks

        Returns:
            EntropyFilterResult with filtered chunks and statistics
        """
        # First, chunk using base chunker
        all_chunks = self.base_chunker.chunk(content, file_path, compute_lines)

        if not all_chunks:
            return EntropyFilterResult(
                chunks=[],
                original_count=0,
                filtered_count=0,
                scores=[],
            )

        # Get embeddings if provider available
        embeddings: list[list[float]] | None = None
        if self.embedding_provider:
            chunk_texts = [chunk.text for chunk in all_chunks]
            embeddings = await self.embedding_provider.embed_texts(chunk_texts)

        # Filter chunks based on information score
        filtered_chunks: list[DocumentChunk] = []
        scores: list[float] = []
        history_entities: set[str] = set()
        history_embeddings: list[list[float]] = []

        for i, chunk in enumerate(all_chunks):
            # Extract entities from current chunk
            chunk_entities = self.extract_entities(chunk.text)
            chunk_word_count = len(chunk.text.split())

            # Get embedding if available
            chunk_embedding = embeddings[i] if embeddings else None

            # Calculate information score
            score = self.information_score(
                chunk_entities=chunk_entities,
                history_entities=history_entities,
                chunk_word_count=chunk_word_count,
                chunk_embedding=chunk_embedding,
                history_embeddings=history_embeddings[-self.history_window :]
                if history_embeddings
                else None,
            )
            scores.append(score)

            # Keep chunk if above threshold OR if it's the first chunk
            # The first chunk is always kept since there's no history to compare against
            is_first_chunk = i == 0
            if is_first_chunk or score >= self.redundancy_threshold:
                # Update chunk index to be sequential in filtered list
                chunk.chunk_index = len(filtered_chunks)
                filtered_chunks.append(chunk)

            # Always update history (even for filtered chunks)
            # This ensures we don't re-include similar content later
            history_entities.update(chunk_entities)
            if chunk_embedding:
                history_embeddings.append(chunk_embedding)

        logger.info(
            "[ENTROPY-CHUNKER] Filtered %d/%d chunks (%.1f%% reduction), "
            "threshold=%.2f, avg_score=%.3f",
            len(all_chunks) - len(filtered_chunks),
            len(all_chunks),
            ((len(all_chunks) - len(filtered_chunks)) / len(all_chunks) * 100) if all_chunks else 0,
            self.redundancy_threshold,
            sum(scores) / len(scores) if scores else 0,
        )

        return EntropyFilterResult(
            chunks=filtered_chunks,
            original_count=len(all_chunks),
            filtered_count=len(filtered_chunks),
            scores=scores,
        )

    def chunk_with_filtering_sync(
        self,
        content: str,
        file_path: str = "",
        compute_lines: bool = True,
    ) -> EntropyFilterResult:
        """Synchronous version of chunk_with_filtering (entity-only, no embeddings).

        Use this when you don't need semantic novelty scoring.

        Args:
            content: Document content to chunk
            file_path: Path to file (for file-type specific chunking)
            compute_lines: If True, compute line numbers for chunks

        Returns:
            EntropyFilterResult with filtered chunks and statistics
        """
        # First, chunk using base chunker
        all_chunks = self.base_chunker.chunk(content, file_path, compute_lines)

        if not all_chunks:
            return EntropyFilterResult(
                chunks=[],
                original_count=0,
                filtered_count=0,
                scores=[],
            )

        # Filter chunks based on entity-only information score
        filtered_chunks: list[DocumentChunk] = []
        scores: list[float] = []
        history_entities: set[str] = set()

        for i, chunk in enumerate(all_chunks):
            # Extract entities from current chunk
            chunk_entities = self.extract_entities(chunk.text)
            chunk_word_count = len(chunk.text.split())

            # Calculate entity-only information score (no embeddings)
            score = self._entity_novelty_score(chunk_entities, history_entities, chunk_word_count)
            scores.append(score)

            # Keep chunk if above threshold OR if it's the first chunk
            # The first chunk is always kept since there's no history to compare against
            is_first_chunk = i == 0
            if is_first_chunk or score >= self.redundancy_threshold:
                chunk.chunk_index = len(filtered_chunks)
                filtered_chunks.append(chunk)

            # Always update history
            history_entities.update(chunk_entities)

        return EntropyFilterResult(
            chunks=filtered_chunks,
            original_count=len(all_chunks),
            filtered_count=len(filtered_chunks),
            scores=scores,
        )
