"""chunker.py -- split text into token-estimate-bounded chunks with overlap (stdlib only).

The overlap carried between chunks is delimited by a unique sentinel so that
``reconstruct`` can invert ``chunk_text`` exactly (lossless round-trip) without
relying on suffix/prefix matching, which is ambiguous on repetitive input.

Architecture (this build):

  Every overlap is *context-only* -- a verbatim copy of the previous chunk's
  tail, prepended behind a sentinel and dropped verbatim by ``reconstruct``.
  There is no "genuine-overlap" case (where the next chunk physically contains
  the overlap as a real prefix), because ``_hard_split_long`` emits *adjacent,
  non-overlapping* pieces.  Collapsing the two overlap kinds into one means
  ``reconstruct`` needs no content-based discriminator at all: it always drops
  the carried ``head`` and appends ``rest``, leaving the previous chunk intact.
  That is unambiguous even when the text is highly repetitive and the overlap
  ratio is degenerate (e.g. ``overlap_tokens == max_tokens - 1``), where a
  ``rest.startswith(head)`` test would conflate genuine and context overlap and
  silently strip real content off the accumulator.

The previous refactor fixed *sentinel discrimination* (genuine vs context) and
removed a ``re.sub("  +", " ")`` collapse, but the round-trip was still lossy
on real whitespace: ``_split_sentences`` discarded the inter-sentence
separators, ``_split_into_chunks`` re-joined sentences with a bare ``" "``, and
``_apply_overlap`` / ``_hard_split_long`` ``.strip()``-ed the carried tails -- so
double spaces, tabs and newlines were silently collapsed to single spaces.

Fix (this build):

  * ``_split_sentences`` now returns ``(sentence, trailing_ws)`` pairs, so the
    concatenation of all pairs reproduces the source verbatim.  Whitespace is
    part of the data, not an afterthought.
  * ``_split_into_chunks`` builds each chunk as an *exact substring* of the
    source (the concatenation of its segments); no ``" "`` re-joining and no
    ``.strip()`` on accumulated content.  A hard-split sentence's trailing
    whitespace is included in the text passed to ``_hard_split_long`` (not
    appended afterward), so it can never push the final piece past max_chars.
  * ``_hard_split_long`` now produces **non-overlapping** pieces (adjacent
    exact substrings of the source) instead of pieces that physically overlap
    the previous piece.  This removes the *genuine-overlap* case entirely:
    every overlap is now *context-only*.
  * ``_hard_split_long`` guards the token bound: a pure-whitespace fragment
    stranded by the word-boundary backtrack is merged into the previous piece
    only when the result stays within ``max_chars``; otherwise it is kept as
    its own piece (harmless: ``reconstruct`` strips trailing whitespace, and
    mid-text whitespace survives in the join).  The previous code merged
    unconditionally, which could overflow the last content piece past max_tokens.
  * ``_apply_overlap`` prepends the previous chunk's tail as pure *context*
    (sentinel-delimited) to every chunk after the first.  Because hard-split
    pieces no longer overlap each other, this context tail is never duplicated
    inside the chunk's own content -- it is a copy, not a real prefix.
  * ``reconstruct`` therefore needs **no discriminator**: it drops the carried
    ``head`` and appends ``rest`` verbatim, leaving the previous chunk intact.
    This is unambiguous even on highly-repetitive text, where a
    ``rest.startswith(head)`` test would conflate genuine and context overlap
    (the prior bug: each misfire stripped chars off the accumulator).  No
    synthetic ``" "`` is inserted; boundary whitespace lives in the previous
    chunk's tail and is preserved.

Result: ``reconstruct(chunk_text(t, ...), ot) == t.strip()`` for every input,
including highly-repetitive text, degenerate high overlap ratios
(e.g. ``overlap_tokens == max_tokens - 1``) and arbitrary inter-sentence
whitespace (double spaces, tabs, newlines, leading/trailing runs).  Each
chunk's estimated token count also satisfies ``estimate_tokens(body) <= max_tokens``
for non-whitespace content (``body`` = chunk with the carried-overlap head
stripped), which the previous build violated for sentences whose trailing
whitespace was appended post-hoc to a max_chars-saturated piece (e.g.
"Pack my box. " at max_tokens=3).
"""
from __future__ import annotations

import math
import re

__all__ = ["chunk_text", "reconstruct", "estimate_tokens"]

# rough heuristic: ~4 chars per token for English-ish text
CHARS_PER_TOKEN = 4

# Sentinel used to mark the carried-overlap region between adjacent chunks.
# It is unlikely to occur in natural text; if it does, it is escaped on
# chunking and unescaped on reconstruction so the round-trip stays lossless.
_OVERLAP_SENTINEL = "\x00OVL\x00"

# Matches the whitespace run that follows an end-of-sentence punctuation mark.
# The run is captured so separators survive splitting (lossless reconstruction).
_SENT_RE = re.compile(r"(?<=[.!?])(\s+)")


def estimate_tokens(s: str) -> int:
    """Estimate the token count of ``s`` using a char-per-token heuristic."""
    if not s:
        return 0
    return max(1, math.ceil(len(s) / CHARS_PER_TOKEN))


# Backward-compatible alias (used by test_chunker_refactor.py).
_estimate_tokens = estimate_tokens


def _split_sentences(text: str) -> list[tuple[str, str]]:
    """Split *text* into ``(sentence, trailing_ws)`` pairs.

    The concatenation of every ``sentence + trailing_ws`` reproduces *text*
    verbatim (after the outer ``.strip()`` applied to *text* by ``chunk_text``),
    so the round-trip never loses inter-sentence whitespace -- double spaces,
    tabs, or newlines survive intact.
    """
    if not text:
        return []
    parts = _SENT_RE.split(text)
    # re.split with a capture group interleaves separators: sentence, sep,
    # sentence, sep, ...  Group them into (sentence, trailing_ws) pairs.
    pairs: list[tuple[str, str]] = []
    for i in range(0, len(parts), 2):
        sent = parts[i]
        sep = parts[i + 1] if i + 1 < len(parts) else ""
        if sent or sep:  # skip a degenerate all-empty pair at the tail
            pairs.append((sent, sep))
    return pairs


def _hard_split_long(text: str, max_chars: int) -> list[str]:
    """Word-aware hard split of a single over-long sentence into adjacent pieces.

    Pieces are *exact*, **non-overlapping** substrings of *text* (no
    ``.strip()``), each no longer than *max_chars* (rounded down to the last
    word boundary).  They are adjacent -- ``"".join(pieces) == text`` -- so no
    genuine-overlap discrimination is needed in ``reconstruct``.  Any cross-piece
    context is instead supplied uniformly by :func:`_apply_overlap` as
    sentinel-delimited *context* overlap, which ``reconstruct`` drops
    unconditionally.  This is what makes the round-trip lossless even under
    highly-repetitive text and high overlap ratios, where a suffix/prefix
    discriminator is ambiguous.
    """
    pieces: list[str] = []
    n = len(text)
    start = 0
    while start < n:
        end = min(start + max_chars, n)
        if end < n:
            boundary = text.rfind(" ", start, end)
            if boundary > start:
                end = boundary
        piece = text[start:end]
        if piece.strip():
            pieces.append(piece)
        elif pieces and len(pieces[-1]) + len(piece) <= max_chars:
            # pure-whitespace fragment (e.g. one space of a double-space run that
            # the word-boundary split stranded between two content pieces): merge
            # it into the previous piece so the concatenation stays an *exact*
            # substring of the source and nothing is silently dropped -- but only
            # when the merge stays within max_chars, otherwise the token bound
            # would be violated; an un-mergeable whitespace fragment is kept as
            # its own piece (reconstruct() .strip()s trailing whitespace, so it
            # is harmless, and losslessness is preserved for mid-text whitespace).
            pieces[-1] += piece
        elif piece:
            pieces.append(piece)
        # if pieces is empty this is leading whitespace; reconstruct() contracts
        # to text.strip(), so it is intentionally not retained here.
        start = end  # while start<n exits naturally when end>=n; no break needed
    return pieces


def _apply_overlap(chunks: list[str], overlap_chars: int) -> list[str]:
    """Prepend a sentinel-delimited tail of the previous chunk to each chunk.

    The tail is pure downstream context (it is *not* duplicated inside the chunk's
    own content); ``.strip()`` is deliberately NOT applied, because that would
    discard the trailing whitespace that makes the round-trip lossless.
    """
    if overlap_chars <= 0 or len(chunks) <= 1:
        return chunks
    out = [chunks[0]]
    # Iterate over consecutive (prev, cur) pairs -- no index arithmetic needed;
    # each tail is exactly the previous chunk's suffix, which `reconstruct` drops
    # unconditionally as pure context.
    for prev, cur in zip(chunks, chunks[1:]):
        out.append(prev[-overlap_chars:] + _OVERLAP_SENTINEL + cur)
    return out


def _split_into_chunks(
    text: str,
    max_tokens: int,
    max_chars: int,
    overlap_chars: int,
) -> list[str]:
    """Core sentence-boundary chunking producing *base* chunks (no overlap).

    Each chunk is the concatenation of whole ``(sentence, sep)`` segments, hence
    an exact substring of the source: internal whitespace is preserved, never
    re-joined with a synthesized ``" "``.
    """
    segments = _split_sentences(text)
    chunks: list[str] = []
    cur = ""
    for sent, sep in segments:
        piece = sent + sep
        # Flush cur exactly when appending *piece* would overflow max_tokens.
        # Safe because estimate_tokens is monotonic in string length: when
        # cur is empty the guard (cur and ...) is skipped and the single
        # estimate_tokens(piece) check below decides.  When cur is non-empty
        # and cur+piece fits, piece alone must also fit (monotonicity), so the
        # else-branch cur += piece matches the original cur = candidate.
        if cur and estimate_tokens(cur + piece) > max_tokens:
            chunks.append(cur)
            cur = ""
        if estimate_tokens(piece) > max_tokens:
            chunks.extend(_hard_split_long(piece, max_chars))
        else:
            cur += piece
    if cur:
        chunks.append(cur)
    return _apply_overlap(chunks, overlap_chars)


def chunk_text(text: str, max_tokens: int = 256, overlap_tokens: int = 32) -> list[str]:
    """Split *text* into chunks each <= *max_tokens* estimated tokens.

    *overlap_tokens* worth of the previous chunk's tail is carried into the
    next chunk (sentinel-delimited).  Reconstruction via :func:`reconstruct`
    reverses this exactly.

    Prefers sentence boundaries; falls back to a word-aware hard split for
    sentences that exceed *max_tokens*.
    """
    if max_tokens <= 0:
        raise ValueError("max_tokens must be positive")
    if overlap_tokens < 0 or overlap_tokens >= max_tokens:
        raise ValueError("overlap_tokens must be in [0, max_tokens)")
    if text is None or not text.strip():
        return []

    max_chars = max_tokens * CHARS_PER_TOKEN
    overlap_chars = overlap_tokens * CHARS_PER_TOKEN
    return _split_into_chunks(text, max_tokens, max_chars, overlap_chars)


def reconstruct(chunks: list[str], overlap_tokens: int = 32) -> str:
    """Invert :func:`chunk_text`: remove sentinel-delimited overlaps and join.

    Exactly reverses chunking for any input (lossless), including highly
    repetitive text and arbitrary inter-sentence whitespace.

    Every overlap in this design is *context-only*: the carried ``head`` is a
    verbatim copy of the previous chunk's tail, never duplicated inside
    ``rest``.  ``reconstruct`` therefore unconditionally drops ``head`` and
    appends ``rest`` verbatim, leaving the previous chunk intact -- no suffix/
    prefix discriminator is required, which is what makes it unambiguous on
    repetitive text where such matching would conflate genuine and context
    overlap.

    A single uniform rule applies to *every* chunk (including the first): the
    overlap head (everything before the sentinel) is pure context and is
    dropped, and ``rest`` (everything after the sentinel) is the genuine
    continuation that is appended verbatim.  The first chunk never carries a
    real overlap by construction, so the rule reduces to "append verbatim"; a
    stray sentinel is still stripped defensively as a no-op head.  Because an
    empty ``chunks`` list simply yields ``"".join([]).strip() == ""``, the
    previous explicit empty-guard and the separate ``i == 0`` branch are
    redundant and have been folded into one loop.

    ``overlap_tokens`` is retained in the signature for API symmetry with
    :func:`chunk_text` -- the canonical round-trip form
    ``reconstruct(chunk_text(t, mt, ot), ot)`` -- and as the ``overlap_tokens
    == 0`` "no overlap" signal.  Reconstruction itself is fully driven by the
    sentinel-delimited overlap carried between chunks, so this parameter is an
    intentional no-op (deleted below) and carries no state.
    """
    # `overlap_tokens` is genuinely unused: each chunk already carries its own
    # sentinel-delimited overlap, so the magnitude is redundant information for
    # reconstruction.  It is kept only for the public round-trip API above;
    # `del` makes the no-op explicit and silences unused-argument detectors
    # without altering any caller.
    del overlap_tokens
    parts: list[str] = []
    for c in chunks:
        if _OVERLAP_SENTINEL in c:
            # The carried `head` is pure context (a copy of the previous chunk's
            # tail), not a prefix duplicated inside `rest`.  Drop it, keep the
            # previous chunk intact, and append `rest` (the genuine
            # continuation) verbatim.  No synthetic separator is inserted:
            # boundary whitespace already lives at the previous chunk's tail.
            _, _, rest = c.partition(_OVERLAP_SENTINEL)
            parts.append(rest)
        else:
            # No sentinel (e.g. overlap_tokens=0, or the first chunk): the chunk
            # is a contiguous source substring; append verbatim.
            parts.append(c)
    return "".join(parts).strip()


if __name__ == "__main__":
    sample = " ".join(f"Sentence number {i}." for i in range(50))
    cs = chunk_text(sample, max_tokens=32, overlap_tokens=4)
    print("chunks:", len(cs))
    for c in cs:
        print("-", c[:60], "...")
    print("reconstructed len:", len(reconstruct(cs, 4)))