Chunking

"""Structure-aware chunking with explicit source lineage.

Every returned chunk retains exact character and token offsets into the source
document.  The richer :class:`ChunkLineage` sits beside the package's existing
``Chunk`` model so current retrievers remain interoperable while ingestion and
citation systems gain the provenance they need.
"""

import bisect
import hashlib
import re
from dataclasses import dataclass
from typing import Iterable, List, Optional, Sequence, Tuple

from .models import Chunk, Document
from .text import TOKEN_RE


SENTENCE_SPAN_RE = re.compile(r"[^.!?\n]+(?:[.!?]+(?=\s|$)|(?=\n)|$)", re.MULTILINE)
HEADING_RE = re.compile(r"^(#{1,6})[ \t]+(.+?)[ \t]*#*[ \t]*$", re.MULTILINE)


@dataclass(frozen=True)
class ChunkLineage:
    """Coordinates that map a chunk back to its source and hierarchy."""

    document_id: str
    chunk_id: str
    strategy: str
    ordinal: int
    start_char: int
    end_char: int
    start_token: int
    end_token: int
    content_hash: str
    section_path: Tuple[str, ...] = ()
    parent_chunk_id: Optional[str] = None


@dataclass(frozen=True)
class LineagedChunk:
    """An ordinary retrieval chunk paired with exact provenance."""

    chunk: Chunk
    lineage: ChunkLineage


@dataclass(frozen=True)
class ParentChildChunks:
    """Two retrieval granularities joined through ``parent_chunk_id``."""

    parents: Tuple[LineagedChunk, ...]
    children: Tuple[LineagedChunk, ...]


@dataclass(frozen=True)
class _Span:
    start: int
    end: int
    section_path: Tuple[str, ...] = ()


def _trim_span(text: str, start: int, end: int) -> Optional[_Span]:
    while start < end and text[start].isspace():
        start += 1
    while end > start and text[end - 1].isspace():
        end -= 1
    return _Span(start, end) if end > start else None


def _sentence_spans(text: str, offset: int = 0) -> List[_Span]:
    spans: List[_Span] = []
    for match in SENTENCE_SPAN_RE.finditer(text):
        span = _trim_span(text, match.start(), match.end())
        if span is not None:
            spans.append(_Span(span.start + offset, span.end + offset))
    return spans


def _split_oversized_spans(text: str, spans: Iterable[_Span], max_tokens: int) -> List[_Span]:
    output: List[_Span] = []
    for span in spans:
        matches = list(TOKEN_RE.finditer(text, span.start, span.end))
        if len(matches) <= max_tokens:
            if matches:
                output.append(span)
            continue
        for start in range(0, len(matches), max_tokens):
            group = matches[start : start + max_tokens]
            output.append(_Span(group[0].start(), group[-1].end(), span.section_path))
    return output


def _pack_sentence_spans(
    text: str,
    spans: Sequence[_Span],
    max_tokens: int,
    overlap_sentences: int,
) -> List[_Span]:
    atoms = _split_oversized_spans(text, spans, max_tokens)
    counts = [len(list(TOKEN_RE.finditer(text, span.start, span.end))) for span in atoms]
    packed: List[_Span] = []
    start = 0
    while start < len(atoms):
        end = start
        token_count = 0
        while end < len(atoms) and token_count + counts[end] <= max_tokens:
            token_count += counts[end]
            end += 1
        if end == start:
            end += 1
        packed.append(
            _Span(
                atoms[start].start,
                atoms[end - 1].end,
                atoms[start].section_path,
            )
        )
        if end == len(atoms):
            break
        start = max(start + 1, end - overlap_sentences)
    return packed


def _token_boundaries(document: Document) -> Tuple[List[int], List[int]]:
    matches = list(TOKEN_RE.finditer(document.text))
    return [match.start() for match in matches], [match.end() for match in matches]


def _record(
    document: Document,
    span: _Span,
    strategy: str,
    ordinal: int,
    token_starts: Sequence[int],
    token_ends: Sequence[int],
    parent_chunk_id: Optional[str] = None,
) -> LineagedChunk:
    start_token = bisect.bisect_right(token_ends, span.start)
    end_token = bisect.bisect_left(token_starts, span.end)
    chunk_id = f"{document.id}::{strategy}::{ordinal:04d}"
    chunk_text = document.text[span.start : span.end]
    digest = hashlib.sha256(chunk_text.encode("utf-8")).hexdigest()
    lineage = ChunkLineage(
        document_id=document.id,
        chunk_id=chunk_id,
        strategy=strategy,
        ordinal=ordinal,
        start_char=span.start,
        end_char=span.end,
        start_token=start_token,
        end_token=end_token,
        content_hash=digest,
        section_path=span.section_path,
        parent_chunk_id=parent_chunk_id,
    )
    metadata = dict(document.metadata)
    metadata.update(
        {
            "chunk_strategy": strategy,
            "content_hash": digest,
            "end_char": span.end,
            "parent_chunk_id": parent_chunk_id,
            "section_path": span.section_path,
            "start_char": span.start,
        }
    )
    chunk = Chunk(
        id=chunk_id,
        document_id=document.id,
        text=chunk_text,
        start_token=start_token,
        end_token=end_token,
        title=document.title,
        source=document.source,
        metadata=metadata,
    )
    return LineagedChunk(chunk, lineage)


def sentence_chunks(
    document: Document,
    max_tokens: int = 120,
    overlap_sentences: int = 1,
) -> Tuple[LineagedChunk, ...]:
    """Pack complete sentences up to a token budget.

    Oversized individual sentences are split at token boundaries.  Overlap is
    expressed in sentences, avoiding the mid-sentence starts caused by blind
    fixed token windows.
    """

    if max_tokens <= 0:
        raise ValueError("max_tokens must be positive")
    if overlap_sentences < 0:
        raise ValueError("overlap_sentences must be non-negative")
    spans = _pack_sentence_spans(
        document.text,
        _sentence_spans(document.text),
        max_tokens,
        overlap_sentences,
    )
    token_starts, token_ends = _token_boundaries(document)
    return tuple(
        _record(document, span, "sentence", ordinal, token_starts, token_ends)
        for ordinal, span in enumerate(spans)
    )


def _markdown_sections(document: Document) -> List[_Span]:
    headings = list(HEADING_RE.finditer(document.text))
    if not headings:
        span = _trim_span(document.text, 0, len(document.text))
        path = (document.title,) if document.title else ()
        return [] if span is None else [_Span(span.start, span.end, path)]

    sections: List[_Span] = []
    preamble = _trim_span(document.text, 0, headings[0].start())
    if preamble is not None:
        path = (document.title,) if document.title else ()
        sections.append(_Span(preamble.start, preamble.end, path))
    path_stack: List[str] = []
    for index, heading in enumerate(headings):
        level = len(heading.group(1))
        name = heading.group(2).strip()
        path_stack = path_stack[: level - 1]
        path_stack.append(name)
        end = headings[index + 1].start() if index + 1 < len(headings) else len(document.text)
        span = _trim_span(document.text, heading.start(), end)
        if span is not None:
            sections.append(_Span(span.start, span.end, tuple(path_stack)))
    return sections


def section_chunks(document: Document, max_tokens: int = 240) -> Tuple[LineagedChunk, ...]:
    """Chunk Markdown by heading boundaries, splitting only oversized sections."""

    if max_tokens <= 0:
        raise ValueError("max_tokens must be positive")
    spans: List[_Span] = []
    for section in _markdown_sections(document):
        tokens = list(TOKEN_RE.finditer(document.text, section.start, section.end))
        if not tokens:
            continue
        for start in range(0, len(tokens), max_tokens):
            group = tokens[start : start + max_tokens]
            spans.append(_Span(group[0].start(), group[-1].end(), section.section_path))
    token_starts, token_ends = _token_boundaries(document)
    return tuple(
        _record(document, span, "section", ordinal, token_starts, token_ends)
        for ordinal, span in enumerate(spans)
    )


def parent_child_chunks(
    document: Document,
    parent_max_tokens: int = 240,
    child_max_tokens: int = 80,
    child_overlap_sentences: int = 1,
) -> ParentChildChunks:
    """Build small retrieval children linked to larger context parents.

    Parents respect Markdown section boundaries.  Children respect sentence
    boundaries inside each parent and retain document-global offsets, allowing
    retrieval by a child followed by expansion to its parent for generation.
    """

    if parent_max_tokens <= 0 or child_max_tokens <= 0:
        raise ValueError("token budgets must be positive")
    if child_overlap_sentences < 0:
        raise ValueError("child_overlap_sentences must be non-negative")
    if child_max_tokens > parent_max_tokens:
        raise ValueError("child_max_tokens must not exceed parent_max_tokens")
    base_parents = section_chunks(document, parent_max_tokens)
    token_starts, token_ends = _token_boundaries(document)
    parents: List[LineagedChunk] = []
    children: List[LineagedChunk] = []
    for parent_ordinal, base in enumerate(base_parents):
        parent_span = _Span(
            base.lineage.start_char,
            base.lineage.end_char,
            base.lineage.section_path,
        )
        parent = _record(
            document,
            parent_span,
            "parent",
            parent_ordinal,
            token_starts,
            token_ends,
        )
        parents.append(parent)
        local_spans = [
            _Span(span.start, span.end, parent.lineage.section_path)
            for span in _sentence_spans(
                parent.chunk.text,
                offset=parent.lineage.start_char,
            )
        ]
        packed = _pack_sentence_spans(
            document.text,
            local_spans,
            child_max_tokens,
            child_overlap_sentences,
        )
        for child_span in packed:
            children.append(
                _record(
                    document,
                    child_span,
                    "child",
                    len(children),
                    token_starts,
                    token_ends,
                    parent.chunk.id,
                )
            )
    return ParentChildChunks(tuple(parents), tuple(children))


# Verb-first aliases make the functions easy to discover alongside
# ``text.chunk_document`` while preserving concise noun-first public names.
chunk_by_sentence = sentence_chunks
chunk_by_section = section_chunks
chunk_parent_child = parent_child_chunks