"""
RBSP Document Chunker.

Splits documents into overlapping chunks suitable for embedding.  Three
strategies: fixed-window, sentence-boundary, and paragraph-boundary.

Token counting uses the same regex as ``engine.py`` (``[a-zA-Z0-9_]+``)
to ensure token counts match BM25 and the inverted index.

Zero external dependencies.
"""

from __future__ import annotations

import re
from dataclasses import dataclass, field
from typing import Any, Dict, Iterable, List, Tuple


# ---------------------------------------------------------------------------
# Constants
# ---------------------------------------------------------------------------

STRATEGY_FIXED = "fixed"
STRATEGY_SENTENCE = "sentence"
STRATEGY_PARAGRAPH = "paragraph"

_VALID_STRATEGIES = frozenset({STRATEGY_FIXED, STRATEGY_SENTENCE, STRATEGY_PARAGRAPH})

# Matches the engine.py tokenizer exactly
_TOKEN_RE = re.compile(r"[a-zA-Z0-9_]+")

# Sentence boundary: split after .!? followed by whitespace
_SENTENCE_RE = re.compile(r"(?<=[.!?])\s+")

# Paragraph boundary: one or more blank lines
_PARAGRAPH_RE = re.compile(r"\n\s*\n")


# ---------------------------------------------------------------------------
# Data structures
# ---------------------------------------------------------------------------


@dataclass(frozen=True, slots=True)
class Chunk:
    """A single chunk produced by the document chunking pipeline.

    Attributes:
        text:         Raw text of this chunk.
        doc_id:       Parent document identifier.
        chunk_index:  0-based index within the document.
        start_offset: Character offset in original document.
        end_offset:   Character offset end in original document.
        token_count:  Number of tokens in this chunk.
        metadata:     Extensible metadata dict.
    """

    text: str
    doc_id: str
    chunk_index: int
    start_offset: int
    end_offset: int
    token_count: int
    metadata: Dict[str, Any] = field(default_factory=dict)


@dataclass(frozen=True, slots=True)
class ChunkResult:
    """Result of chunking an entire document.

    Attributes:
        doc_id:       Document identifier.
        chunks:       Ordered tuple of Chunk objects.
        total_tokens: Total tokens across all chunks.
        strategy:     Which strategy was used.
    """

    doc_id: str
    chunks: Tuple[Chunk, ...]
    total_tokens: int
    strategy: str


# ---------------------------------------------------------------------------
# Token counting (mirrors engine.py _TOKEN_RE exactly)
# ---------------------------------------------------------------------------


def _count_tokens(text: str) -> int:
    """Count tokens using the same regex as the engine tokenizer.

    Args:
        text: Input text.

    Returns:
        Number of tokens.
    """
    return len(_TOKEN_RE.findall(text))


def _token_boundaries(text: str) -> List[Tuple[int, int]]:
    """Return (start, end) character offsets for each token match.

    Args:
        text: Input text.

    Returns:
        List of (start, end) offset tuples.
    """
    return [(m.start(), m.end()) for m in _TOKEN_RE.finditer(text)]


# ---------------------------------------------------------------------------
# DocumentChunker
# ---------------------------------------------------------------------------


class DocumentChunker:
    """Splits documents into chunks suitable for embedding.

    Supports three strategies:
        - ``"fixed"``:     Fixed token-count windows with overlap.
        - ``"sentence"``:  Split on sentence boundaries, merge until
                           ``chunk_size`` approached. (DEFAULT)
        - ``"paragraph"``: Split on paragraph boundaries (blank lines),
                           merge short paragraphs.

    Token counting uses ``[a-zA-Z0-9_]+`` regex, matching engine.py exactly.
    Stemming is NOT applied (doesn't affect count).

    Args:
        chunk_size:     Target tokens per chunk (default 384).
        overlap:        Overlapping tokens between adjacent chunks (default 64).
        strategy:       One of ``"fixed"``, ``"sentence"``, ``"paragraph"``
                        (default ``"sentence"``).
        min_chunk_size: Minimum tokens for a chunk; smaller trailing chunks
                        merge into the previous chunk (default 32).
    """

    __slots__ = ("_chunk_size", "_overlap", "_strategy", "_min_chunk_size")

    def __init__(
        self,
        chunk_size: int = 384,
        overlap: int = 64,
        strategy: str = "sentence",
        min_chunk_size: int = 32,
    ) -> None:
        if chunk_size < 1:
            raise ValueError("chunk_size must be >= 1")
        if overlap < 0:
            raise ValueError("overlap must be >= 0")
        if overlap >= chunk_size:
            raise ValueError("overlap must be < chunk_size")
        if strategy not in _VALID_STRATEGIES:
            raise ValueError(
                f"strategy must be one of {sorted(_VALID_STRATEGIES)}, "
                f"got {strategy!r}"
            )
        if min_chunk_size < 1:
            raise ValueError("min_chunk_size must be >= 1")
        self._chunk_size = chunk_size
        self._overlap = overlap
        self._strategy = strategy
        self._min_chunk_size = min_chunk_size

    def chunk(self, text: str, doc_id: str = "") -> ChunkResult:
        """Chunk a single document.

        Args:
            text:   Full document text.
            doc_id: Identifier for the source document.

        Returns:
            ChunkResult with ordered Chunk objects.
        """
        if not text or not text.strip():
            empty = Chunk(
                text="", doc_id=doc_id, chunk_index=0,
                start_offset=0, end_offset=0, token_count=0,
            )
            return ChunkResult(
                doc_id=doc_id, chunks=(empty,),
                total_tokens=0, strategy=self._strategy,
            )

        if self._strategy == STRATEGY_FIXED:
            chunks = self._chunk_fixed(text, doc_id)
        elif self._strategy == STRATEGY_SENTENCE:
            chunks = self._chunk_sentence(text, doc_id)
        else:
            chunks = self._chunk_paragraph(text, doc_id)

        # Merge tiny trailing chunk into previous
        if len(chunks) > 1 and chunks[-1].token_count < self._min_chunk_size:
            last = chunks.pop()
            prev = chunks[-1]
            merged_text = text[prev.start_offset : last.end_offset]
            chunks[-1] = Chunk(
                text=merged_text,
                doc_id=doc_id,
                chunk_index=prev.chunk_index,
                start_offset=prev.start_offset,
                end_offset=last.end_offset,
                token_count=prev.token_count + last.token_count,
            )

        total = sum(c.token_count for c in chunks)
        return ChunkResult(
            doc_id=doc_id,
            chunks=tuple(chunks),
            total_tokens=total,
            strategy=self._strategy,
        )

    def chunk_many(
        self, documents: Iterable[Tuple[str, str]]
    ) -> List[ChunkResult]:
        """Chunk multiple (doc_id, text) pairs.

        Args:
            documents: Iterable of (doc_id, text) tuples.

        Returns:
            List of ChunkResult objects.
        """
        return [self.chunk(text, doc_id) for doc_id, text in documents]

    # -- Fixed strategy ----------------------------------------------------

    def _chunk_fixed(self, text: str, doc_id: str) -> List[Chunk]:
        """Split into fixed token-count windows."""
        boundaries = _token_boundaries(text)
        if not boundaries:
            return [Chunk(
                text="", doc_id=doc_id, chunk_index=0,
                start_offset=0, end_offset=0, token_count=0,
            )]

        chunks: List[Chunk] = []
        step = self._chunk_size - self._overlap
        i = 0
        idx = 0
        while i < len(boundaries):
            window = boundaries[i : i + self._chunk_size]
            start = window[0][0]
            end = window[-1][1]
            chunk_text = text[start:end]
            chunks.append(Chunk(
                text=chunk_text,
                doc_id=doc_id,
                chunk_index=idx,
                start_offset=start,
                end_offset=end,
                token_count=len(window),
            ))
            i += step
            idx += 1

        return chunks

    # -- Sentence strategy -------------------------------------------------

    def _chunk_sentence(self, text: str, doc_id: str) -> List[Chunk]:
        """Split on sentence boundaries, greedily merging."""
        parts = _SENTENCE_RE.split(text)
        if not parts:
            return [Chunk(
                text="", doc_id=doc_id, chunk_index=0,
                start_offset=0, end_offset=0, token_count=0,
            )]

        # Build (sentence_text, token_count, start_offset, end_offset)
        sentences: List[Tuple[str, int, int, int]] = []
        offset = 0
        for part in parts:
            start = text.find(part, offset)
            end = start + len(part)
            tc = _count_tokens(part)
            sentences.append((part, tc, start, end))
            offset = end

        chunks: List[Chunk] = []
        idx = 0
        si = 0
        while si < len(sentences):
            accumulated: List[int] = []
            running_tokens = 0

            while si < len(sentences):
                stc = sentences[si][1]
                if running_tokens + stc > self._chunk_size and accumulated:
                    break
                accumulated.append(si)
                running_tokens += stc
                si += 1

            start_off = sentences[accumulated[0]][2]
            end_off = sentences[accumulated[-1]][3]
            chunk_text = text[start_off:end_off]
            tc = sum(sentences[j][1] for j in accumulated)

            chunks.append(Chunk(
                text=chunk_text,
                doc_id=doc_id,
                chunk_index=idx,
                start_offset=start_off,
                end_offset=end_off,
                token_count=tc,
            ))
            idx += 1

            # Overlap: rewind sentence index so overlap tokens are repeated
            if self._overlap > 0 and si < len(sentences):
                overlap_tokens = 0
                rewind = 0
                for k in range(len(accumulated) - 1, -1, -1):
                    overlap_tokens += sentences[accumulated[k]][1]
                    rewind += 1
                    if overlap_tokens >= self._overlap:
                        break
                rewind_target = si - rewind
                # Don't rewind past where we started accumulating
                min_start = accumulated[0] + 1
                if rewind_target >= min_start:
                    si = rewind_target

        return chunks

    # -- Paragraph strategy ------------------------------------------------

    def _chunk_paragraph(self, text: str, doc_id: str) -> List[Chunk]:
        """Split on paragraph boundaries, greedily merging."""
        parts = _PARAGRAPH_RE.split(text)
        if not parts:
            return [Chunk(
                text="", doc_id=doc_id, chunk_index=0,
                start_offset=0, end_offset=0, token_count=0,
            )]

        # Build (para_text, token_count, start_offset, end_offset)
        paragraphs: List[Tuple[str, int, int, int]] = []
        offset = 0
        for part in parts:
            start = text.find(part, offset)
            end = start + len(part)
            tc = _count_tokens(part)
            paragraphs.append((part, tc, start, end))
            offset = end

        chunks: List[Chunk] = []
        idx = 0
        pi = 0
        while pi < len(paragraphs):
            accumulated: List[int] = []
            running_tokens = 0

            while pi < len(paragraphs):
                ptc = paragraphs[pi][1]
                if running_tokens + ptc > self._chunk_size and accumulated:
                    break
                accumulated.append(pi)
                running_tokens += ptc
                pi += 1

            start_off = paragraphs[accumulated[0]][2]
            end_off = paragraphs[accumulated[-1]][3]
            chunk_text = text[start_off:end_off]
            tc = sum(paragraphs[j][1] for j in accumulated)

            chunks.append(Chunk(
                text=chunk_text,
                doc_id=doc_id,
                chunk_index=idx,
                start_offset=start_off,
                end_offset=end_off,
                token_count=tc,
            ))
            idx += 1

            # Overlap: repeat last paragraph
            if self._overlap > 0 and accumulated and pi < len(paragraphs):
                last_para_tokens = paragraphs[accumulated[-1]][1]
                if last_para_tokens > 0:
                    pi = accumulated[-1]

        return chunks

    def __repr__(self) -> str:
        return (
            f"DocumentChunker(chunk_size={self._chunk_size}, "
            f"overlap={self._overlap}, strategy={self._strategy!r})"
        )
