"""Chunking objects not specific to a particular chunking strategy."""

from __future__ import annotations

import collections
import copy
import uuid
from functools import cached_property
from typing import Any, Callable, DefaultDict, Iterable, Iterator, cast

import regex
from lxml.etree import ParserError, tostring
from lxml.html import fragment_fromstring
from typing_extensions import Self, TypeAlias

from unstructured.common.html_table import HtmlCell, HtmlRow, HtmlTable
from unstructured.documents.elements import (
    CodeSnippet,
    CompositeElement,
    ConsolidationStrategy,
    Element,
    ElementMetadata,
    Table,
    TableChunk,
    Title,
)
from unstructured.logger import logger

# ================================================================================================
# MODEL
# ================================================================================================

CHUNK_MAX_CHARS_DEFAULT: int = 500
"""Hard-max chunk-length when no explicit value specified in `max_characters` argument.

Provided for reference only, for example so the ingest CLI can advertise the default value in its
UI. External chunking-related functions (e.g. in ingest or decorators) should use
`max_characters: int | None = None` and not apply this default themselves. Only
`ChunkingOptions.max_characters` should apply a default value.
"""

CHUNK_MULTI_PAGE_DEFAULT: bool = True
"""When False, respect page-boundaries (no two elements from different page in same chunk).

Only operative for "by_title" chunking strategy.
"""

BoundaryPredicate: TypeAlias = Callable[[Element], bool]
"""Detects when element represents crossing a semantic boundary like section or page."""

TextAndHtml: TypeAlias = tuple[str, str]


class TokenCounter:
    """Token counting using tiktoken for token-based chunking.

    Lazily imports tiktoken only when token counting is first used.
    """

    def __init__(self, tokenizer: str):
        self._tokenizer_name = tokenizer

    @cached_property
    def _encoder(self):
        """Lazily initialize the tiktoken encoder."""
        import tiktoken

        try:
            # -- try as model name first (e.g., "gpt-4") --
            return tiktoken.encoding_for_model(self._tokenizer_name)
        except KeyError:
            # -- fall back to encoding name (e.g., "cl100k_base") --
            return tiktoken.get_encoding(self._tokenizer_name)

    def count(self, text: str) -> int:
        """Return the number of tokens in `text`."""
        return len(self._encoder.encode(text))

    def validate(self) -> None:
        """Resolve the tokenizer now, raising if it is unknown or tiktoken is not installed.

        The encoder is otherwise resolved on the first `count()`. The list-form chunkers drive
        the whole pipeline during the call, so they hit that immediately, but the `iter_*` forms
        return before anything is counted. Forcing resolution during option validation keeps
        both forms raising at the same call site.
        """
        _ = self._encoder


# ================================================================================================
# CHUNKING OPTIONS
# ================================================================================================


class ChunkingOptions:
    """Specifies parameters of optional chunking behaviors.

    Parameters
    ----------
    max_characters
        Hard-maximum text-length of chunk. A chunk longer than this will be split mid-text and be
        emitted as two or more chunks. Mutually exclusive with `max_tokens`.
    max_tokens
        Hard-maximum token count of chunk. A chunk with more tokens than this will be split mid-text
        and be emitted as two or more chunks. Requires `tokenizer` to be specified. Mutually
        exclusive with `max_characters`.
    new_after_n_chars
        Preferred approximate chunk size. A chunk composed of elements totalling this size or
        greater is considered "full" and will not be enlarged by adding another element, even if it
        will fit within the remaining `max_characters` for that chunk. Defaults to `max_characters`
        when not specified, which effectively disables this behavior. Specifying 0 for this
        argument causes each element to appear in a chunk by itself (although an element with text
        longer than `max_characters` will be still be split into two or more chunks).
    new_after_n_tokens
        Token-based equivalent of `new_after_n_chars`. Preferred approximate chunk size in tokens.
        Requires `tokenizer` and `max_tokens` to be specified.
    combine_text_under_n_chars
        Provides a way to "recombine" small chunks formed by breaking on a semantic boundary. Only
        relevant for a chunking strategy that specifies higher-level semantic boundaries to be
        respected, like "section" or "page". Recursively combines two adjacent pre-chunks when the
        first pre-chunk is smaller than this threshold. "Recursively" here means the resulting
        pre-chunk can be combined with the next pre-chunk if it is still under the length threshold.
        Defaults to `max_characters` which combines chunks whenever space allows. Specifying 0 for
        this argument suppresses combining of small chunks. Note this value is "capped" at the
        `new_after_n_chars` value since a value higher than that would not change this parameter's
        effect.
    overlap
        Specifies the length of a string ("tail") to be drawn from each chunk and prefixed to the
        next chunk as a context-preserving mechanism. By default, this only applies to split-chunks
        where an oversized element is divided into multiple chunks by text-splitting.
    overlap_all
        Default: `False`. When `True`, apply overlap between "normal" chunks formed from whole
        elements and not subject to text-splitting. Use this with caution as it entails a certain
        level of "pollution" of otherwise clean semantic chunk boundaries.
    repeat_table_headers
        Default: `True`. When `True`, repeated table-header behavior is enabled for chunked table
        continuations. Specify `False` to opt out and preserve legacy table-chunk behavior.
    isolate_table
        Default: `True`. When `True`, `Table` and `TableChunk` elements are always staged in
        their own pre-chunk and never combined with adjacent non-table elements. Specify
        `False` to allow tables to share pre-chunks with adjacent elements (the pre-#4307
        behavior), which is sometimes useful when downstream consumers expect mixed-content
        composite chunks.
    text_splitting_separators
        A sequence of strings like `("\n", " ")` to be used as target separators during
        text-splitting. Text-splitting only applies to splitting an oversized element into two or
        more chunks. These separators are tried in the specified order until one is found in the
        string to be split. The default separator is `""` which matches between any two characters.
        This separator should not be specified in this sequence because it is always the separator
        of last-resort. Note that because the separator is removed during text-splitting, only
        whitespace character sequences are suitable.
    tokenizer
        The tokenizer to use for token-based chunking. Can be either an encoding name (e.g.,
        "cl100k_base") or a model name (e.g., "gpt-4"). Required when using `max_tokens`.
    """

    def __init__(self, **kwargs: Any):
        self._kwargs = kwargs

    @classmethod
    def new(cls, **kwargs: Any) -> Self:
        """Return instance or raises `ValueError` on invalid arguments like overlap > max_chars."""
        self = cls(**kwargs)
        self._validate()
        return self

    @cached_property
    def boundary_predicates(self) -> tuple[BoundaryPredicate, ...]:
        """The semantic-boundary detectors to be applied to break pre-chunks.

        Overridden by sub-typs to provide semantic-boundary isolation behaviors.
        """
        return ()

    @cached_property
    def combine_text_under_n_chars(self) -> int:
        """Combine two consecutive text pre-chunks if first is smaller than this and both will fit.

        Default applied here is `0` which essentially disables chunk combining. Must be overridden
        by subclass where combining behavior is supported.
        """
        arg_value = self._kwargs.get("combine_text_under_n_chars")
        return arg_value if arg_value is not None else 0

    @cached_property
    def hard_max(self) -> int:
        """The maximum size for a chunk (in characters or tokens depending on mode).

        A pre-chunk will only exceed this size when it contains exactly one element which by itself
        exceeds this size. Such a pre-chunk is subject to mid-text splitting later in the chunking
        process.
        """
        if self.use_token_counting:
            # -- token-based chunking: max_tokens is required and validated --
            return self._kwargs["max_tokens"]

        arg_value = self._kwargs.get("max_characters")
        return arg_value if arg_value is not None else CHUNK_MAX_CHARS_DEFAULT

    @cached_property
    def include_orig_elements(self) -> bool:
        """When True, add original elements from pre-chunk to `.metadata.orig_elements` of chunk.

        Default value is `True`.
        """
        arg_value = self._kwargs.get("include_orig_elements")
        return True if arg_value is None else bool(arg_value)

    @cached_property
    def repeat_table_headers(self) -> bool:
        """When True, repeat detected table headers in continuation table chunks.

        Default value is `True`.
        """
        arg_value = self._kwargs.get("repeat_table_headers")
        return True if arg_value is None else bool(arg_value)

    @cached_property
    def skip_table_chunking(self) -> bool:
        """When True, Table elements are passed through without chunking.

        Default value is `False`.
        """
        arg_value = self._kwargs.get("skip_table_chunking")
        return False if arg_value is None else bool(arg_value)

    @cached_property
    def isolate_table(self) -> bool:
        """When True, `Table`/`TableChunk` elements are staged in their own pre-chunk.

        Default value is `True`. When `False`, table-family elements are allowed to share a
        pre-chunk with adjacent non-table elements (and may be merged by `PreChunkCombiner`),
        restoring the pre-#4307 behavior.
        """
        arg_value = self._kwargs.get("isolate_table")
        return True if arg_value is None else bool(arg_value)

    @cached_property
    def inter_chunk_overlap(self) -> int:
        """Characters of overlap to add between chunks.

        This applies only to boundaries between chunks formed from whole elements and not to
        text-splitting boundaries that arise from splitting an oversized element.
        """
        overlap_all_arg = self._kwargs.get("overlap_all")
        return self.overlap if overlap_all_arg else 0

    @cached_property
    def overlap(self) -> int:
        """The number of characters to overlap text when splitting chunks mid-text.

        The actual overlap will not exceed this number of characters but may be less as required to
        respect splitting-character boundaries.
        """
        overlap_arg = self._kwargs.get("overlap")
        return overlap_arg or 0

    @cached_property
    def soft_max(self) -> int:
        """A pre-chunk of this size or greater is considered full.

        Note that while a value of `0` is valid, it essentially disables chunking by putting
        each element into its own chunk.
        """
        hard_max = self.hard_max

        if self.use_token_counting:
            new_after_n_tokens_arg = self._kwargs.get("new_after_n_tokens")
            # -- default value is == max_tokens --
            if new_after_n_tokens_arg is None:
                return hard_max
            # -- new_after_n_tokens > max_tokens behaves the same as ==max_tokens --
            if new_after_n_tokens_arg > hard_max:
                return hard_max
            return new_after_n_tokens_arg

        new_after_n_chars_arg = self._kwargs.get("new_after_n_chars")

        # -- default value is == max_characters --
        if new_after_n_chars_arg is None:
            return hard_max

        # -- new_after_n_chars > max_characters behaves the same as ==max_characters --
        if new_after_n_chars_arg > hard_max:
            return hard_max

        # -- otherwise, give them what they asked for --
        return new_after_n_chars_arg

    @cached_property
    def split(self) -> Callable[[str], tuple[str, str]]:
        """A text-splitting function suitable for splitting the text of an oversized pre-chunk.

        The function is pre-configured with the chosen chunking window size and any other applicable
        options specified by the caller as part of this chunking-options instance.
        """
        return _TextSplitter(self)

    @cached_property
    def text_separator(self) -> str:
        """The string to insert between elements when concatenating their text for a chunk.

        Right now this is just "\n\n" (a blank line in plain text), but having this here rather
        than as a module-level constant provides a way for us to easily make it user-configurable
        in future if we want to.
        """
        return "\n\n"

    @cached_property
    def text_splitting_separators(self) -> tuple[str, ...]:
        """Sequence of text-splitting target strings to be used in order of preference."""
        text_splitting_separators_arg = self._kwargs.get("text_splitting_separators")
        return (
            ("\n", " ")
            if text_splitting_separators_arg is None
            else tuple(text_splitting_separators_arg)
        )

    @cached_property
    def token_counter(self) -> TokenCounter | None:
        """The token counter for token-based chunking, or None for character-based chunking."""
        tokenizer = self._kwargs.get("tokenizer")
        return TokenCounter(tokenizer) if tokenizer else None

    @cached_property
    def use_token_counting(self) -> bool:
        """True when token-based chunking is configured, False for character-based."""
        return self._kwargs.get("max_tokens") is not None

    def measure(self, text: str) -> int:
        """Return the size of `text` in the configured units (characters or tokens)."""
        if self.use_token_counting and self.token_counter:
            return self.token_counter.count(text)
        return len(text)

    def _validate(self) -> None:
        """Raise ValueError if requestion option-set is invalid."""
        max_tokens = self._kwargs.get("max_tokens")
        max_characters = self._kwargs.get("max_characters")
        tokenizer = self._kwargs.get("tokenizer")

        # -- max_tokens and max_characters are mutually exclusive --
        if max_tokens is not None and max_characters is not None:
            raise ValueError(
                "'max_tokens' and 'max_characters' are mutually exclusive;"
                " specify one or the other, not both"
            )

        # -- max_tokens requires tokenizer. An empty string is rejected along with `None`: it is
        # -- not `None` so it would pass this check, but `token_counter` is `None` for any falsey
        # -- tokenizer, which silently sends `measure()` down the character-counting path and
        # -- enforces `max_tokens` as a count of characters.
        if max_tokens is not None and not tokenizer:
            raise ValueError("'tokenizer' is required when using 'max_tokens'")

        # -- max_tokens must be positive --
        if max_tokens is not None and max_tokens <= 0:
            raise ValueError(f"'max_tokens' argument must be > 0, got {max_tokens}")

        # -- new_after_n_tokens requires max_tokens --
        new_after_n_tokens = self._kwargs.get("new_after_n_tokens")
        if new_after_n_tokens is not None and max_tokens is None:
            raise ValueError("'new_after_n_tokens' requires 'max_tokens' to be specified")

        # -- new_after_n_tokens must be non-negative --
        if new_after_n_tokens is not None and new_after_n_tokens < 0:
            raise ValueError(
                f"'new_after_n_tokens' argument must be >= 0, got {new_after_n_tokens}"
            )

        # -- chunking window must have positive length --
        hard_max = self.hard_max
        if hard_max <= 0:
            raise ValueError(f"'max_characters' argument must be > 0, got {hard_max}")

        # -- a negative value for `new_after_n_chars` is assumed to be a mistake the caller will
        # -- want to know about
        new_after_n_chars = self._kwargs.get("new_after_n_chars")
        if new_after_n_chars is not None and new_after_n_chars < 0:
            raise ValueError(f"'new_after_n_chars' argument must be >= 0, got {new_after_n_chars}")

        # -- `skip_table_chunking` requires `isolate_table` because the pass-through path only
        # -- fires when the pre-chunk contains a single `Table` element. With isolation disabled,
        # -- tables can fold into `CompositeElement` alongside neighbors and the skip would be
        # -- silently ignored, breaking the contract.
        if self.skip_table_chunking and not self.isolate_table:
            raise ValueError(
                "'skip_table_chunking=True' requires 'isolate_table=True' (the default);"
                " tables cannot be passed through unchanged while also sharing a pre-chunk with"
                " adjacent elements"
            )

        # -- overlap must be less than max-chars or the chunk text will never be consumed --
        if self.overlap >= hard_max:
            raise ValueError(
                f"'overlap' argument must be less than `max_characters`,"
                f" got {self.overlap} >= {hard_max}"
            )

        # -- an unknown tokenizer is an invalid option, so surface it here with the rest rather
        # -- than on the first token count. Otherwise `chunk_*()` raises during the call (it
        # -- counts immediately) while `iter_chunk*()` raises whenever the caller first advances
        # -- the generator it returned. Only checked when token counting is actually in use,
        # -- matching `measure()`; a `tokenizer` passed without `max_tokens` goes unused. The
        # -- check above guarantees a counter exists whenever token counting is in use.
        if self.use_token_counting and self.token_counter is not None:
            self.token_counter.validate()


# ================================================================================================
# TABLE ISOLATION (SHARED PRECHECKS)
# ================================================================================================
# Tables are always staged alone in a pre-chunk so downstream splitting can emit `Table` /
# `TableChunk` elements instead of folding them into `CompositeElement` with surrounding text.
# See GitHub issue #3921 and the chunking strategy docs for rationale.


def _element_is_table_family(element: Element) -> bool:
    """True when ``element`` is a `Table` or a concrete subtype such as `TableChunk`.

    Subclasses share the same isolation contract: they must not share a pre-chunk with arbitrary
    text elements, and two table-bearing sequences must not be merged by `PreChunkCombiner`.
    """
    return isinstance(element, Table)


def _elements_contain_table_family(elements: Iterable[Element]) -> bool:
    """True when ``elements`` already includes at least one table-family element."""
    return any(_element_is_table_family(e) for e in elements)


def _table_isolation_forbids_side_by_side_merge(
    left: Iterable[Element],
    right: Iterable[Element],
) -> bool:
    """True when a proposed merge of two element streams must be rejected for table isolation.

    If either side already contains a table, the combiner must flush before accepting the other
    side. This keeps `combine_text_under_n_chars` from concatenating a table pre-chunk with
    neighboring narrative pre-chunks.
    """
    return _elements_contain_table_family(left) or _elements_contain_table_family(right)


# ================================================================================================
# PRE-CHUNKER
# ================================================================================================


class PreChunker:
    """Gathers sequential elements into pre-chunks as length constraints allow.

    The pre-chunker's responsibilities are:

    - **Segregate semantic units.** Identify semantic unit boundaries and segregate elements on
      either side of those boundaries into different sections. In this case, the primary indicator
      of a semantic boundary is a `Title` element. A page-break (change in page-number) is also a
      semantic boundary when `multipage_sections` is `False`.

    - **Minimize chunk count for each semantic unit.** Group the elements within a semantic unit
      into sections as big as possible without exceeding the chunk window size.

    - **Minimize chunks that must be split mid-text.** Precompute the text length of each section
      and only produce a section that exceeds the chunk window size when there is a single element
      with text longer than that window.

    A Table element is placed into a section by itself. CheckBox elements are dropped.

    The "by-title" strategy specifies breaking on section boundaries; a `Title` element indicates
    a new "section", hence the "by-title" designation.
    """

    def __init__(self, elements: Iterable[Element], opts: ChunkingOptions):
        self._elements = elements
        self._opts = opts

    @classmethod
    def iter_pre_chunks(
        cls, elements: Iterable[Element], opts: ChunkingOptions
    ) -> Iterator[PreChunk]:
        """Generate pre-chunks from the element-stream provided on construction."""
        return cls(elements, opts)._iter_pre_chunks()

    def _iter_pre_chunks(self) -> Iterator[PreChunk]:
        """Generate pre-chunks from the element-stream provided on construction.

        A *pre-chunk* is the largest sub-sequence of elements that will both fit within the
        chunking window and respects the semantic boundary rules of the chunking strategy. When a
        single element exceeds the chunking window size it is placed in a pre-chunk by itself and
        is subject to mid-text splitting in the second phase of the chunking process.
        """
        pre_chunk_builder = PreChunkBuilder(self._opts)

        for element in self._elements:
            # -- start new pre-chunk when necessary to uphold segregation guarantees --
            if (
                # -- start new pre-chunk when necessary to uphold segregation guarantees --
                self._is_in_new_semantic_unit(element)
                # -- or when next element won't fit --
                or not pre_chunk_builder.will_fit(element)
            ):
                yield from pre_chunk_builder.flush()

            # -- add this element to the work-in-progress (WIP) pre-chunk --
            pre_chunk_builder.add_element(element)

        # -- flush "tail" pre-chunk, any partially-filled pre-chunk after last element is
        # -- processed
        yield from pre_chunk_builder.flush()

    @cached_property
    def _boundary_predicates(self) -> tuple[BoundaryPredicate, ...]:
        """The semantic-boundary detectors to be applied to break pre-chunks."""
        return self._opts.boundary_predicates

    def _is_in_new_semantic_unit(self, element: Element) -> bool:
        """True when `element` begins a new semantic unit such as a section or page."""
        # -- all detectors need to be called to update state and avoid double counting
        # -- boundaries that happen to coincide, like Table and new section on same element.
        # -- Using `any()` would short-circuit on first True.
        semantic_boundaries = [pred(element) for pred in self._boundary_predicates]
        return any(semantic_boundaries)


class PreChunkBuilder:
    """An element accumulator suitable for incrementally forming a pre-chunk.

    Provides the trial method `.will_fit()` a pre-chunker can use to determine whether it should add
    the next element in the element stream.

    `.flush()` is used to build a PreChunk object from the accumulated elements. This method
    returns an iterator that generates zero-or-one `PreChunk` object and is used like so:

        yield from builder.flush()

    If no elements have been accumulated, no `PreChunk` instance is generated. Flushing the builder
    clears the elements it contains so it is ready to build the next pre-chunk.
    """

    def __init__(self, opts: ChunkingOptions) -> None:
        self._opts = opts
        self._separator_len = len(opts.text_separator)
        self._elements: list[Element] = []

        # -- overlap is only between pre-chunks so starts empty --
        self._overlap_prefix: str = ""
        # -- only includes non-empty element text, e.g. PageBreak.text=="" is not included --
        self._text_segments: list[str] = []
        # -- combined length of text-segments, not including separators --
        self._text_len: int = 0

    def add_element(self, element: Element) -> None:
        """Add `element` to this section."""
        # -- do not prefix a table-only pre-chunk with narrative overlap from the prior chunk --
        if (
            self._opts.isolate_table
            and len(self._elements) == 0
            and _element_is_table_family(element)
        ):
            self._overlap_prefix = ""
            self._text_segments = []
            self._text_len = 0

        self._elements.append(element)
        if element.text:
            self._text_segments.append(element.text)
            # -- only track char-based length; token-based length computed on demand --
            self._text_len += len(element.text)

    def flush(self) -> Iterator[PreChunk]:
        """Generate zero-or-one `PreChunk` object and clear the accumulator.

        Suitable for use to emit a PreChunk when the maximum size has been reached or a semantic
        boundary has been reached. Also to clear out a terminal pre-chunk at the end of an element
        stream.
        """
        elements = self._elements

        if not elements:
            return

        # -- copy element list, don't use original or it may change contents as builder proceeds --
        pre_chunk = PreChunk(elements, self._overlap_prefix, self._opts)
        # -- clear builder before yield so we're not sensitive to the timing of how/when this
        # -- iterator is exhausted and can add elements for the next pre-chunk immediately.
        overlap_for_next = pre_chunk.overlap_tail
        # -- table tails must not prefix the following narrative pre-chunk (overlap_all) --
        if (
            self._opts.isolate_table
            and len(elements) == 1
            and _element_is_table_family(elements[0])
        ):
            overlap_for_next = ""
        self._reset_state(overlap_for_next)
        yield pre_chunk

    def will_fit(self, element: Element) -> bool:
        """True when `element` can be added to this prechunk without violating its limits.

        There are several limits:
        - A `Table` element will never fit with any other element. It will only fit in an empty
          pre-chunk.
        - No element will fit in a pre-chunk that already contains a `Table` element.
        - A text-element will not fit in a pre-chunk that already exceeds the soft-max
          (aka. new_after_n_chars/new_after_n_tokens).
        - A text-element will not fit when together with the elements already present it would
          exceed the hard-max (aka. max_characters/max_tokens).
        """
        if self._opts.isolate_table:
            # -- a `Table` can only start a pre-chunk; it is never appended to a non-empty
            # -- pre-chunk --
            if _element_is_table_family(element):
                return len(self._elements) == 0

            # -- no non-table element may share a pre-chunk with a `Table` --
            if _elements_contain_table_family(self._elements):
                return False

        # -- an empty pre-chunk will accept any element (including an oversized-element) --
        if len(self._elements) == 0:
            return True
        # -- a pre-chunk that already exceeds the soft-max is considered "full" --
        if self._text_length > self._opts.soft_max:
            return False
        # -- don't add an element if it would increase total size beyond the hard-max --
        # -- for token counting, compute what the new total would be --
        if self._opts.use_token_counting:
            new_text = self._opts.text_separator.join(
                self._text_segments + ([element.text] if element.text else [])
            )
            return self._opts.measure(new_text) <= self._opts.hard_max
        # -- for character counting, use the efficient incremental approach --
        return not self._remaining_space < len(element.text or "")

    @property
    def _remaining_space(self) -> int:
        """Maximum text-length of an element that can be added without exceeding maxlen."""
        # -- include length of trailing separator that will go before next element text --
        separators_len = self._separator_len * len(self._text_segments)
        return self._opts.hard_max - self._text_len - separators_len

    def _reset_state(self, overlap_prefix: str) -> None:
        """Set working-state values back to "empty", ready to accumulate next pre-chunk."""
        self._overlap_prefix = overlap_prefix
        self._elements.clear()
        self._text_segments = [overlap_prefix] if overlap_prefix else []
        self._text_len = len(overlap_prefix)

    @property
    def _text_length(self) -> int:
        """Size of the text in this pre-chunk (in characters or tokens depending on mode).

        This value represents the chunk-size that would result if this pre-chunk was flushed in its
        current state. In particular, it does not include the length of a trailing separator (since
        that would only appear if an additional element was added).

        Not suitable for judging remaining space, use `.remaining_space` for that value.
        """
        # -- for token counting, compute the actual token count of the joined text --
        if self._opts.use_token_counting:
            if not self._text_segments:
                return 0
            text = self._opts.text_separator.join(self._text_segments)
            return self._opts.measure(text)

        # -- for character counting, use the efficient incremental approach --
        # -- number of text separators present in joined text of elements. This includes only
        # -- separators *between* text segments, not one at the end. Note there are zero separators
        # -- for both 0 and 1 text-segments.
        n = len(self._text_segments)
        separator_count = n - 1 if n else 0
        return self._text_len + (separator_count * self._separator_len)


# ================================================================================================
# PRE-CHUNK
# ================================================================================================


class PreChunk:
    """Sequence of elements staged to form a single chunk.

    This object is purposely immutable.
    """

    def __init__(
        self, elements: Iterable[Element], overlap_prefix: str, opts: ChunkingOptions
    ) -> None:
        self._elements = list(elements)
        self._overlap_prefix = overlap_prefix
        self._opts = opts

    def __eq__(self, other: Any) -> bool:
        if not isinstance(other, PreChunk):
            return False
        return self._overlap_prefix == other._overlap_prefix and self._elements == other._elements

    def can_combine(self, pre_chunk: PreChunk) -> bool:
        """True when `pre_chunk` can be combined with this one without exceeding size limits."""
        if self._opts.isolate_table and _table_isolation_forbids_side_by_side_merge(
            self._elements, pre_chunk._elements
        ):
            return False
        if len(self._text) >= self._opts.combine_text_under_n_chars:
            return False
        # -- avoid duplicating length computations by doing a trial-combine which is just as
        # -- efficient and definitely more robust than hoping two different computations of combined
        # -- length continue to get the same answer as the code evolves. Only possible because
        # -- `.combine()` is non-mutating.
        combined_len = len(self.combine(pre_chunk)._text)

        return combined_len <= self._opts.hard_max

    def combine(self, other_pre_chunk: PreChunk) -> PreChunk:
        """Return new `PreChunk` that combines this and `other_pre_chunk`."""
        # -- combined pre-chunk gets the overlap-prefix of the first pre-chunk. The second overlap
        # -- is automatically incorporated at the end of the first chunk, where it originated.
        return PreChunk(
            self._elements + other_pre_chunk._elements,
            overlap_prefix=self._overlap_prefix,
            opts=self._opts,
        )

    def iter_chunks(self) -> Iterator[CompositeElement | Table | TableChunk]:
        """Form this pre-chunk into one or more chunk elements maxlen or smaller.

        When the total size of the pre-chunk will fit in the chunking window, a single chunk it
        emitted. When this prechunk contains an oversized element (always isolated), it is split
        into two or more chunks that each fit the chunking window.
        """

        # -- a one-table-only pre-chunk is handled specially, by `TablePreChunk`, mainly because
        # -- it may need to be split into multiple `TableChunk` elements and that operation is
        # -- quite specialized.
        if len(self._elements) == 1 and isinstance(self._elements[0], Table):
            if self._opts.skip_table_chunking:
                yield self._elements[0]
            else:
                yield from _TableChunker.iter_chunks(
                    self._elements[0], self._overlap_prefix, self._opts
                )
        else:
            yield from _Chunker.iter_chunks(self._elements, self._text, self._opts)

    @cached_property
    def overlap_tail(self) -> str:
        """The portion of this chunk's text to be repeated as a prefix in the next chunk.

        This value is the empty-string ("") when either the `.overlap` length option is `0` or
        `.overlap_all` is `False`. When there is a text value, it is stripped of both leading and
        trailing whitespace.
        """
        overlap = self._opts.inter_chunk_overlap
        return self._text[-overlap:].strip() if overlap else ""

    def _iter_text_segments(self) -> Iterator[str]:
        """Generate overlap text and each element text segment in order.

        Empty text segments are not included. CodeSnippet elements preserve their
        original whitespace (including newlines) to maintain code formatting.
        """
        if self._overlap_prefix:
            yield self._overlap_prefix
        for e in self._elements:
            if e.text and len(e.text):
                # -- preserve all whitespace for code snippets to maintain formatting --
                if isinstance(e, CodeSnippet):
                    yield e.text
                else:
                    text = " ".join(e.text.strip().split())
                    if text:
                        yield text

    @cached_property
    def _text(self) -> str:
        """The concatenated text of all elements in this pre-chunk, including any overlap.

        Whitespace is normalized to a single space. The text of each element is separated from
        that of the next by a blank line ("\n\n").
        """
        return self._opts.text_separator.join(self._iter_text_segments())


# ================================================================================================
# CHUNKING HELPER/SPLITTERS
# ================================================================================================


class _Chunker:
    """Forms chunks from a pre-chunk other than one containing only a `Table`.

    Produces zero-or-more `CompositeElement` objects.
    """

    def __init__(self, elements: Iterable[Element], text: str, opts: ChunkingOptions) -> None:
        self._elements = list(elements)
        self._text = text
        self._opts = opts

    @classmethod
    def iter_chunks(
        cls, elements: Iterable[Element], text: str, opts: ChunkingOptions
    ) -> Iterator[CompositeElement]:
        """Form zero or more chunks from `elements`.

        One `CompositeElement` is produced when all `elements` will fit. Otherwise there is a
        single `Text`-subtype element and chunks are formed by splitting.
        """
        return cls(elements, text, opts)._iter_chunks()

    def _iter_chunks(self) -> Iterator[CompositeElement]:
        """Form zero or more chunks from `elements`."""
        # -- a pre-chunk containing no text (maybe only a PageBreak element for example) does not
        # -- generate any chunks.
        if not self._text:
            return

        # -- `split()` is the text-splitting function used to split an oversized element --
        split = self._opts.split

        # -- emit first chunk --
        s, remainder = split(self._text)
        yield CompositeElement(text=s, metadata=self._chunk_metadata())

        # -- an oversized pre-chunk will have a remainder, split that up into additional chunks.
        # -- Note these get continuation_metadata which includes is_continuation=True.
        while remainder:
            s, remainder = split(remainder)
            yield CompositeElement(text=s, metadata=self._continuation_metadata())

    @cached_property
    def _all_metadata_values(self) -> dict[str, list[Any]]:
        """Collection of all populated metadata values across elements.

        The resulting dict has one key for each `ElementMetadata` field that had a non-None value in
        at least one of the elements in this pre-chunk. The value of that key is a list of all those
        populated values, in element order, for example:

            {
                "filename": ["sample.docx", "sample.docx"],
                "languages": [["lat"], ["lat", "eng"]]
                ...
            }

        This preprocessing step provides the input for a specified consolidation strategy that will
        resolve the list of values for each field to a single consolidated value.
        """

        def iter_populated_fields(metadata: ElementMetadata) -> Iterator[tuple[str, Any]]:
            """(field_name, value) pair for each non-None field in single `ElementMetadata`."""
            return (
                (field_name, value)
                for field_name, value in metadata.known_fields.items()
                if value is not None
            )

        field_values: DefaultDict[str, list[Any]] = collections.defaultdict(list)

        # -- collect all non-None field values in a list for each field, in element-order --
        for e in self._elements:
            for field_name, value in iter_populated_fields(e.metadata):
                field_values[field_name].append(value)

        return dict(field_values)

    @cached_property
    def _consolidated_metadata(self) -> ElementMetadata:
        """Metadata applicable to this pre-chunk as a single chunk.

        Formed by applying consolidation rules to all metadata fields across the elements of this
        pre-chunk.

        For the sake of consistency, the same rules are applied (for example, for dropping values)
        to a single-element pre-chunk too, even though metadata for such a pre-chunk is already
        "consolidated".
        """
        consolidated_metadata = ElementMetadata(**self._meta_kwargs)
        if self._opts.include_orig_elements:
            consolidated_metadata.orig_elements = self._orig_elements
        return consolidated_metadata

    def _continuation_metadata(self) -> ElementMetadata:
        """Fresh metadata for one second-or-later text-split chunk of the pre-chunk.

        The same metadata as the first text-split chunk but includes `.is_continuation = True`.
        Unused for non-oversized pre-chunks since those are not subject to text-splitting.

        A new object is produced on each call (not cached) because each continuation chunk needs
        its own copy: `enrichment_origins` is a mutable dict-of-lists that a downstream additive
        enrichment may mutate in place, and a shared object would let one chunk's mutation leak
        into its siblings.
        """
        continuation_metadata = self._chunk_metadata()
        continuation_metadata.is_continuation = True
        return continuation_metadata

    def _chunk_metadata(self) -> ElementMetadata:
        """Fresh metadata for one text-split chunk of the pre-chunk."""
        # -- we need to make a copy, otherwise adding a field would also change metadata value
        # -- already assigned to another chunk. A shallow copy suffices for scalar fields, but
        # -- `enrichment_origins` is mutable, so deep-copy it. (Deep-copying the whole metadata is
        # -- avoided because it may carry the full `orig_elements`.)
        metadata = copy.copy(self._consolidated_metadata)
        if metadata.enrichment_origins is not None:
            metadata.enrichment_origins = copy.deepcopy(metadata.enrichment_origins)
        return metadata

    @cached_property
    def _meta_kwargs(self) -> dict[str, Any]:
        """The consolidated metadata values as a dict suitable for constructing ElementMetadata.

        This is where consolidation strategies are actually applied. The output is suitable for use
        in constructing an `ElementMetadata` object like `ElementMetadata(**self._meta_kwargs)`.
        """
        CS = ConsolidationStrategy
        field_consolidation_strategies = ConsolidationStrategy.field_consolidation_strategies()

        def iter_kwarg_pairs() -> Iterator[tuple[str, Any]]:
            """Generate (field-name, value) pairs for each field in consolidated metadata."""
            for field_name, values in self._all_metadata_values.items():
                strategy = field_consolidation_strategies.get(field_name)
                if strategy is CS.FIRST:
                    yield field_name, values[0]
                # -- concatenate lists from each element that had one, in order --
                elif strategy is CS.LIST_CONCATENATE:
                    yield field_name, sum(values, cast("list[Any]", []))
                # -- union lists from each element, preserving order of appearance --
                elif strategy is CS.LIST_UNIQUE:
                    # -- Python 3.7+ maintains dict insertion order --
                    ordered_unique_keys = {key: None for val_list in values for key in val_list}
                    yield field_name, list(ordered_unique_keys.keys())
                elif strategy is CS.STRING_CONCATENATE:
                    yield field_name, " ".join(val.strip() for val in values)
                # -- merge dict-of-list values: union keys, per key concatenate then dedupe
                # -- records, preserving first-seen order --
                elif strategy is CS.DICT_LIST_UNIQUE:
                    merged: dict[str, list[Any]] = {}
                    for value in values:
                        for key, records in value.items():
                            seen = merged.setdefault(key, [])
                            seen_ids = {tuple(sorted(r.items())) for r in seen}
                            for record in records:
                                record_id = tuple(sorted(record.items()))
                                if record_id not in seen_ids:
                                    seen_ids.add(record_id)
                                    seen.append(record)
                    yield field_name, merged
                elif strategy is CS.DROP:
                    continue
                else:  # pragma: no cover
                    # -- not likely to hit this since we have a test in `text_elements.py` that
                    # -- ensures every ElementMetadata fields has an assigned strategy.
                    raise NotImplementedError(
                        f"metadata field {repr(field_name)} has no defined consolidation strategy"
                    )

        return dict(iter_kwarg_pairs())

    @cached_property
    def _orig_elements(self) -> list[Element]:
        """The `.metadata.orig_elements` value for chunks formed from this pre-chunk."""

        def iter_orig_elements():
            for e in self._elements:
                if e.metadata.orig_elements is None:
                    yield e
                    continue
                # -- make copy of any element we're going to mutate because these elements don't
                # -- belong to us (the user may have downstream purposes for them).
                orig_element = copy.copy(e)
                # -- prevent recursive .orig_elements when element is a chunk (has orig-elements of
                # -- its own)
                orig_element.metadata.orig_elements = None
                yield orig_element

        return list(iter_orig_elements())


class _TableChunker:
    """Responsible for forming chunks, especially splits, from a single-table pre-chunk.

    Table splitting is specialized because we recursively split on an even row, cell, text
    boundary. This object encapsulate those details.
    """

    def __init__(self, table: Table, overlap_prefix: str, opts: ChunkingOptions) -> None:
        self._table = table
        self._overlap_prefix = overlap_prefix
        self._opts = opts

    @classmethod
    def iter_chunks(
        cls, table: Table, overlap_prefix: str, opts: ChunkingOptions
    ) -> Iterator[Table | TableChunk]:
        """Split this pre-chunk into `Table` or `TableChunk` objects maxlen or smaller."""
        return cls(table, overlap_prefix, opts)._iter_chunks()

    def _iter_chunks(self) -> Iterator[Table | TableChunk]:
        """Split this pre-chunk into `Table` or `TableChunk` objects maxlen or smaller."""
        # -- A table with no non-whitespace text produces no chunks --
        if not self._table_text:
            return

        # -- only text-split a table when it's longer than the chunking window --
        maxlen = self._opts.hard_max
        measure = self._opts.measure
        text_size = measure(self._text_with_overlap)
        html_size = measure(self._html) if self._html else 0

        if text_size <= maxlen and html_size <= maxlen:
            # -- use the compactified html for .text_as_html, even though we're not splitting --
            metadata = self._metadata
            metadata.text_as_html = self._html or None
            # -- note the overlap-prefix is prepended to its text --
            yield Table(text=self._text_with_overlap, metadata=metadata)
            return

        # -- When there's no HTML, split it like a normal element. Also fall back to text-only
        # -- chunks when `max_characters` is less than 50 (or in token mode, less than 15 tokens).
        # -- `.text_as_html` metadata is impractical for a chunking window that small because the
        # -- 33 characters of HTML overhead for each chunk would produce many very small chunks.
        min_html_threshold = 15 if self._opts.use_token_counting else 50
        if not self._html or self._opts.hard_max < min_html_threshold:
            yield from self._iter_text_only_table_chunks()
            return

        # -- otherwise, form splits with "synchronized" text and html --
        yield from self._iter_text_and_html_table_chunks()

    @cached_property
    def _html(self) -> str:
        """The compactified HTML for this table when it has text-as-HTML.

        The empty string when table-structure has not been captured, perhaps because
        `infer_table_structure` was set `False` in the partitioning call.
        """
        if not (html_table := self._html_table):
            return ""

        return html_table.html

    @cached_property
    def _html_table(self) -> HtmlTable | None:
        """The `lxml` HTML element object for this table.

        `None` when the `Table` element has no `.metadata.text_as_html`.
        """
        if (text_as_html := self._table.metadata.text_as_html) is None:
            return None

        text_as_html = text_as_html.strip()
        if not text_as_html:  # pragma: no cover
            return None

        try:
            return HtmlTable.from_html_text(text_as_html)
        except (ParserError, ValueError):
            logger.warning(
                "Could not parse text_as_html for table element; skipping HTML-based chunking."
                " text_as_html: %s",
                text_as_html[:100] + "..." if len(text_as_html) > 100 else text_as_html,
            )
            return None

    @cached_property
    def _leading_header_row_count(self) -> int:
        """Number of contiguous leading rows that should be treated as table headers."""
        html_table = self._html_table
        if html_table is None:
            return 0

        count = 0
        for row in html_table.iter_rows():
            if not row.is_header:
                break
            count += 1
        return count

    def _iter_text_and_html_table_chunks(self) -> Iterator[TableChunk]:
        """Split table into chunks where HTML corresponds exactly to text.

        `.metadata.text_as_html` for each chunk is a parsable `<table>` HTML fragment.
        """
        if (html_table := self._html_table) is None:  # pragma: no cover
            raise ValueError("this method is undefined for a table having no .text_as_html")

        header_row_count = self._leading_header_row_count if self._opts.repeat_table_headers else 0
        splitter = _HtmlTableSplitter(
            html_table,
            self._opts,
            header_row_count=header_row_count,
        )
        yield from self._make_table_chunks(
            _HtmlTableSplitter.iter_subtables(
                html_table,
                self._opts,
                header_row_count=header_row_count,
            ),
            num_carried_over_header_rows=splitter.carried_over_header_row_count,
        )

    def _iter_text_only_table_chunks(self) -> Iterator[TableChunk]:
        """Split oversized text-only table (no text-as-html) into chunks.

        `.metadata.text_as_html` is optional, not included when `infer_table_structure` is
        `False`.
        """

        def _iter_text_splits() -> Iterator[tuple[str, None]]:
            text_remainder = self._text_with_overlap
            split = self._opts.split
            while text_remainder:
                # -- split off the next chunk-worth of characters into a TableChunk --
                chunk_text, text_remainder = split(text_remainder)
                yield chunk_text, None

        yield from self._make_table_chunks(_iter_text_splits())

    def _make_table_chunks(
        self,
        text_html_pairs: Iterator[tuple[str, str | None]],
        num_carried_over_header_rows: int = 0,
    ) -> Iterator[TableChunk]:
        """Form `TableChunk` objects from (text, html) pairs.

        Handles `is_continuation` and chunk sequencing metadata (`table_id`, `chunk_index`)
        so the original table can be reconstructed from its chunks. Carries
        `num_carried_over_header_rows` so synthetic repeated header rows can be removed.
        """
        table_id = str(uuid.uuid4())
        carried_header_row_count = max(0, num_carried_over_header_rows)

        for chunk_index, (text, html) in enumerate(text_html_pairs):
            metadata = self._metadata
            if html is not None:
                metadata.text_as_html = html
            else:
                metadata.text_as_html = None
            # -- second and later chunks get `.metadata.is_continuation = True` --
            metadata.is_continuation = (chunk_index > 0) or None
            metadata.num_carried_over_header_rows = (
                carried_header_row_count if chunk_index > 0 else 0
            )

            chunk = TableChunk(text=text, metadata=metadata)
            chunk.metadata.table_id = table_id
            chunk.metadata.chunk_index = chunk_index
            yield chunk

    @property
    def _metadata(self) -> ElementMetadata:
        """The base `.metadata` value for chunks formed from this pre-chunk.

        The term "base" here means that other metadata fields will be added, depending on the
        chunk. In particular, `.metadata.text_as_html` will be different for each text-split chunk
        and `.metadata.is_continuation` must be added for second-and-later text-split chunks.

        Note this is a fresh copy of the metadata on each call since it will need to be mutated
        differently for each chunk formed from this pre-chunk.
        """
        CS = ConsolidationStrategy
        metadata = copy.deepcopy(self._table.metadata)

        # -- drop metadata fields not appropriate for chunks, in particular
        # -- parent_id's will not reliably point to an existing element
        drop_field_names = [
            field_name
            for field_name, strategy in CS.field_consolidation_strategies().items()
            if strategy is CS.DROP
        ]
        for field_name in drop_field_names:
            setattr(metadata, field_name, None)

        if self._opts.include_orig_elements:
            metadata.orig_elements = self._orig_elements
        return metadata

    @cached_property
    def _orig_elements(self) -> list[Element]:
        """The `.metadata.orig_elements` value for chunks formed from this pre-chunk.

        Note this is not just the `Table` element, it must be adjusted to strip out any
        `.metadata.orig_elements` value it may have when it is itself a chunk and not a direct
        product of partitioning.
        """
        # -- make a copy because we're going to mutate the `Table` element and it doesn't belong to
        # -- us (the user may have downstream purposes for it).
        orig_table = copy.deepcopy(self._table)
        # -- prevent recursive .orig_elements when `Table` element is a chunk --
        orig_table.metadata.orig_elements = None
        return [orig_table]

    @cached_property
    def _table_text(self) -> str:
        """The text in this table, not including any overlap-prefix or extra whitespace."""
        if not self._table.text:
            return ""
        return " ".join(self._table.text.split())

    @cached_property
    def _text_with_overlap(self) -> str:
        """The text for this chunk, including the overlap-prefix when present."""
        overlap_prefix = self._overlap_prefix
        table_text = "" if not self._table.text else self._table.text.strip()
        # -- use row-separator between overlap and table-text --
        return overlap_prefix + "\n" + table_text if overlap_prefix else table_text


# ================================================================================================
# HTML SPLITTERS
# ================================================================================================


class _HtmlTableSplitter:
    """Produces (text, html) pairs for a `<table>` HtmlElement.

    Each chunk contains a whole number of rows whenever possible. An oversized row is split on an
    even cell boundary and a single cell that is by itself too big to fit in the chunking window
    is divided by text-splitting.

    The returned `html` value is always a parseable HTML `<table>` subtree.
    """

    def __init__(self, table_element: HtmlTable, opts: ChunkingOptions, header_row_count: int = 0):
        self._table_element = table_element
        self._opts = opts
        self._header_row_count = max(0, header_row_count)

    @classmethod
    def iter_subtables(
        cls, table_element: HtmlTable, opts: ChunkingOptions, header_row_count: int = 0
    ) -> Iterator[TextAndHtml]:
        """Generate (text, html) pair for each split of this table pre-chunk.

        Each split is on an even row boundary whenever possible, falling back to even cell and even
        word boundaries when a row or cell is by itself oversized, respectively.
        """
        return cls(table_element, opts, header_row_count=header_row_count)._iter_subtables()

    def _iter_subtables(self) -> Iterator[TextAndHtml]:
        """Generate (text, html) pairs containing as many whole rows as will fit in window.

        Falls back to splitting rows into whole cells when a single row is by itself too big to
        fit in the chunking window.
        """
        is_first_chunk = True
        accum = _RowAccumulator(maxlen=self._maxlen(is_first_chunk), measure=self._opts.measure)

        for row in self._table_element.iter_rows():
            # -- if row won't fit, any WIP chunk is done, send it on its way --
            if not accum.will_fit(row):
                for text, html in accum.flush():
                    yield self._prepend_repeated_headers(text, html, is_first_chunk)
                    is_first_chunk = False
                accum = _RowAccumulator(
                    maxlen=self._maxlen(is_first_chunk), measure=self._opts.measure
                )
            # -- if row fits, add it to accumulator --
            if accum.will_fit(row):
                accum.add_row(row)
            else:  # -- otherwise, single row is bigger than chunking window --
                for text, html in self._iter_row_splits(row, maxlen=self._maxlen(is_first_chunk)):
                    yield self._prepend_repeated_headers(text, html, is_first_chunk)
                    is_first_chunk = False
                accum = _RowAccumulator(
                    maxlen=self._maxlen(is_first_chunk), measure=self._opts.measure
                )

        for text, html in accum.flush():
            yield self._prepend_repeated_headers(text, html, is_first_chunk)
            is_first_chunk = False

    def _iter_row_splits(self, row: HtmlRow, maxlen: int) -> Iterator[TextAndHtml]:
        """Split oversized row into (text, html) pairs containing as many cells as will fit."""
        accum = _CellAccumulator(maxlen=maxlen)

        for cell in row.iter_cells():
            # -- if cell won't fit, flush and check again --
            if not accum.will_fit(cell):
                yield from accum.flush()
            # -- if cell fits, add it to accumulator --
            if accum.will_fit(cell):
                accum.add_cell(cell)
            else:  # -- otherwise, single cell is bigger than chunking window --
                yield from self._iter_cell_splits(cell, maxlen=maxlen)

        yield from accum.flush()

    def _iter_cell_splits(self, cell: HtmlCell, maxlen: int) -> Iterator[TextAndHtml]:
        """Split a single oversized cell into sub-sub-sub-table HTML fragments."""
        # -- 33 is len("<table><tr><td></td></tr></table>"), HTML overhead beyond text content --
        # -- For token-based chunking, we subtract 33 chars worth of overhead but still use tokens
        # -- for the actual content limit. For character-based, we use the reduced character limit.
        if self._opts.use_token_counting:
            # -- In token mode, keep token limit but account for HTML overhead in char terms --
            # -- The HTML tags themselves are usually ~10-15 tokens, so we reduce by a small amount
            opts = ChunkingOptions(
                max_tokens=max(1, maxlen - 10),
                tokenizer=self._opts._kwargs.get("tokenizer"),
            )
        else:
            opts = ChunkingOptions(max_characters=max(1, maxlen - 33))
        split = _TextSplitter(opts)

        text, remainder = split(cell.text)
        yield text, f"<table><tr><td>{text}</td></tr></table>"

        # -- an oversized cell will have a remainder, split that up into additional chunks.
        while remainder:
            text, remainder = split(remainder)
            yield text, f"<table><tr><td>{text}</td></tr></table>"

    @cached_property
    def _header_text(self) -> str:
        """Concatenated text for leading header rows identified by caller."""
        return " ".join(text for row in self._header_rows for text in row.iter_cell_texts())

    @cached_property
    def _header_rows(self) -> tuple[HtmlRow, ...]:
        """Leading rows that should be repeated on continuation chunks, if any."""
        if self._header_row_count <= 0:
            return ()

        rows: list[HtmlRow] = []
        for idx, row in enumerate(self._table_element.iter_rows()):
            if idx >= self._header_row_count:
                break
            rows.append(row)
        return tuple(rows)

    @cached_property
    def _header_rows_html(self) -> str:
        """HTML for repeated header rows, preserving header semantics."""
        if not self._header_rows:
            return ""

        rows_html = "".join(self._as_header_row_html(row) for row in self._header_rows)
        return f"<thead>{rows_html}</thead>"

    @cached_property
    def carried_over_header_row_count(self) -> int:
        """Header-row count prepended to each continuation chunk, or 0 when disabled."""
        return len(self._header_rows) if self._should_repeat_headers else 0

    @cached_property
    def _should_repeat_headers(self) -> bool:
        """True when header repetition is enabled and not pathologically expensive."""
        if not self._header_rows:
            return False

        # -- guard against pathological headers where a single repeated header row would consume
        # -- more than half the chunking window.
        return self._max_header_row_len <= (self._opts.hard_max + 1) // 2

    @cached_property
    def _max_header_row_len(self) -> int:
        """Largest leading-header row text length."""
        if not self._header_rows:
            return 0
        return max(self._opts.measure(" ".join(row.iter_cell_texts())) for row in self._header_rows)

    @cached_property
    def _header_text_len(self) -> int:
        """Size of repeated header text in chunking units."""
        return self._opts.measure(self._header_text)

    def _maxlen(self, is_first_chunk: bool) -> int:
        """Available size for non-header row content of the next chunk."""
        if is_first_chunk or not self._should_repeat_headers:
            return self._opts.hard_max

        # -- reserve one separator between repeated header text and chunk body text --
        return max(1, self._opts.hard_max - self._header_text_len - 1)

    def _prepend_repeated_headers(self, text: str, html: str, is_first_chunk: bool) -> TextAndHtml:
        """Prepend repeated header rows to continuation chunk when enabled."""
        if is_first_chunk or not self._should_repeat_headers:
            return text, html

        header_text = self._header_text
        chunk_text = f"{header_text} {text}" if header_text and text else (header_text or text)

        html_inner = html.removeprefix("<table>").removesuffix("</table>")
        chunk_html = f"<table>{self._header_rows_html}{html_inner}</table>"
        return chunk_text, chunk_html

    @staticmethod
    def _as_header_row_html(row: HtmlRow) -> str:
        """Serialize `row` preserving source HTML while converting direct-child `<td>` to `<th>`."""
        row_html = row.source_html or row.html
        tr = _HtmlTableSplitter._parse_row_fragment(row_html)
        if tr is None and row.source_html:
            tr = _HtmlTableSplitter._parse_row_fragment(row.html)
        if tr is None:
            return row.html

        for cell in tr:
            if getattr(cell, "tag", None) == "td":
                cell.tag = "th"

        return tostring(tr, encoding=str)

    @staticmethod
    def _parse_row_fragment(row_html: str):
        """Parse `row_html` and return a `<tr>` element when recoverable."""
        try:
            parsed = fragment_fromstring(row_html)
        except (ParserError, ValueError):
            return None

        if parsed.tag == "tr":
            return parsed

        rows = parsed.xpath(".//tr")
        return rows[0] if rows else None


class _TextSplitter:
    """Provides a text-splitting function configured on construction.

    Text is split on the best-available separator, falling-back from the preferred separator
    through a sequence of alternate separators.

    - The separator is removed by splitting so only whitespace strings are suitable separators.
    - A "blank-line" ("\n\n") is unlikely to occur in an element as it would have been used as an
      element boundary during partitioning.

    This is a *callable* object. Constructing it essentially produces a function:

        split = _TextSplitter(opts)
        fragment, remainder = split(s)

    This allows it to be configured with length-options etc. on construction and used throughout a
    chunking operation on a given element-stream.
    """

    def __init__(self, opts: ChunkingOptions):
        self._opts = opts

    def __call__(self, s: str) -> tuple[str, str]:
        """Return pair of strings split from `s` on the best match of configured patterns.

        The first string is the split, the second is the remainder of the string. The split string
        will never be longer than `maxlen` (in characters or tokens depending on mode). The
        separators are tried in order until a match is found. The last separator is "" which matches
        between any two characters so there will always be a split.

        The separator is removed and does not appear in the split or remainder.

        An `s` that is already less than the maximum length is returned unchanged with no remainder.
        This allows this function to be called repeatedly with the remainder until it is consumed
        and returns a remainder of "".
        """
        maxlen = self._opts.hard_max

        # -- for token counting, use the measurement abstraction for size check --
        if self._opts.use_token_counting:
            if self._opts.measure(s) <= maxlen:
                return s, ""
            return self._split_by_tokens(s)

        # -- character-based splitting (original logic) --
        if len(s) <= maxlen:
            return s, ""

        for p, sep_len in self._patterns:
            # -- length of separator must be added to include that separator when it happens to be
            # -- located exactly at maxlen. Otherwise the search-from-end regex won't find it.
            fragment, remainder = self._split_from_maxlen(p, sep_len, s)
            if (
                # -- no available split with this separator --
                not fragment
                # -- split did not progress, consuming part of the string --
                or len(remainder) >= len(s)
            ):
                continue
            return fragment.rstrip(), remainder.lstrip()

        # -- the terminal "" pattern is not actually executed via regex since its implementation is
        # -- trivial and provides a hard back-stop here in this method. No separator is used between
        # -- tail and remainder on arb-char split.
        return s[:maxlen].rstrip(), s[maxlen - self._opts.overlap :].lstrip()

    def _split_by_tokens(self, s: str) -> tuple[str, str]:
        """Split text `s` on a separator boundary while respecting token limits.

        Tries each separator in order of preference, looking for the rightmost split position
        that keeps the fragment under the token limit. Falls back to splitting on whitespace
        boundaries if no separator works.
        """
        maxlen = self._opts.hard_max
        overlap = self._opts.overlap
        measure = self._opts.measure

        # -- try each separator in order of preference --
        for pattern, _ in self._patterns:
            # -- find all matches of this separator in the string --
            # -- note: (?r) flag makes finditer return matches right-to-left --
            matches = list(pattern.finditer(s))
            # -- iterate through matches (already right-to-left due to (?r) flag) --
            for match in matches:
                match_start, match_end = match.span()
                fragment = s[:match_start].rstrip()
                # -- check if fragment fits within token limit --
                if measure(fragment) <= maxlen:
                    # -- skip if fragment is too short (needs at least some content) --
                    if measure(fragment) == 0:
                        continue
                    raw_remainder = s[match_end:].lstrip()
                    # -- add overlap if configured --
                    if overlap > 0:
                        # -- token-based overlap: find tail with ~overlap tokens --
                        tail = self._get_token_overlap_tail(fragment, overlap)
                        overlapped_remainder = tail + " " + raw_remainder
                        return fragment, overlapped_remainder
                    return fragment, raw_remainder

        # -- fallback: split on whitespace boundary using binary search to find token limit --
        # -- find the approximate character position that corresponds to maxlen tokens --
        low, high = 0, len(s)
        best_pos = max(overlap + 1, 1)  # -- minimum viable position --

        while low <= high:
            mid = (low + high) // 2
            if measure(s[:mid]) <= maxlen:
                best_pos = mid
                low = mid + 1
            else:
                high = mid - 1

        # -- try to find a whitespace boundary near best_pos, searching backwards --
        split_pos = best_pos
        for i in range(best_pos, max(overlap, 0), -1):
            if i < len(s) and s[i].isspace():
                split_pos = i
                break

        # -- ensure the fragment still fits after whitespace adjustment --
        fragment = s[:split_pos].rstrip()
        if measure(fragment) > maxlen and split_pos > overlap + 1:
            # -- whitespace boundary pushed us over; use the binary search result directly --
            fragment = s[:best_pos].rstrip()
            split_pos = best_pos

        raw_remainder = s[split_pos:].lstrip()

        if overlap > 0 and fragment:
            tail = self._get_token_overlap_tail(fragment, overlap)
            overlapped_remainder = tail + " " + raw_remainder
            return fragment, overlapped_remainder

        return fragment, raw_remainder

    def _get_token_overlap_tail(self, text: str, target_tokens: int) -> str:
        """Extract tail of text containing approximately `target_tokens` tokens.

        Uses binary search to find the character position from which the tail contains
        approximately the specified number of tokens. Adjusts to word boundaries to avoid
        splitting words.
        """
        measure = self._opts.measure

        # -- if the entire text has fewer tokens than target, return all of it --
        if measure(text) <= target_tokens:
            return text.strip()

        # -- binary search to find the character position that yields ~target_tokens --
        low, high = 0, len(text)

        while low < high:
            mid = (low + high) // 2
            tail = text[mid:]
            token_count = measure(tail)
            if token_count > target_tokens:
                low = mid + 1
            else:
                high = mid

        # -- adjust to word boundary: search forward for whitespace then skip it --
        pos = low
        while pos < len(text) and not text[pos].isspace():
            pos += 1
        while pos < len(text) and text[pos].isspace():
            pos += 1

        # -- if we've moved too far, fall back to just stripping leading whitespace --
        if pos >= len(text):
            return text[low:].lstrip()

        return text[pos:]

    @cached_property
    def _patterns(self) -> tuple[tuple[regex.Pattern[str], int], ...]:
        """Sequence of (pattern, len) pairs to match against.

        Patterns appear in order of preference, those following are "fall-back" patterns to be used
        if no match of a prior pattern is found.

        NOTE these regexes search *from the end of the string*, which is what the "(?r)" bit
        specifies. This is much more efficient than starting at the beginning of the string which
        could result in hundreds of matches before the desired one.
        """
        separators = self._opts.text_splitting_separators
        return tuple((regex.compile(f"(?r){sep}"), len(sep)) for sep in separators)

    def _split_from_maxlen(
        self, pattern: regex.Pattern[str], sep_len: int, s: str
    ) -> tuple[str, str]:
        """Return (split, remainder) pair split from `s` on the right-most match before `maxlen`.

        Returns `"", s` if no suitable match was found. Also returns `"", s` if splitting on this
        separator produces a split shorter than the required overlap (which would produce an
        infinite loop).

        `split` will never be longer than `maxlen` and there is no longer split available using
        `pattern`.

        The separator is removed and does not appear in either the split or remainder.
        """
        maxlen, overlap = self._opts.hard_max, self._opts.overlap

        # -- A split not longer than overlap will not progress (infinite loop). On the right side,
        # -- need to extend search range to include a separator located exactly at maxlen.
        match = pattern.search(s, pos=overlap + 1, endpos=maxlen + sep_len)
        if match is None:
            return "", s

        # -- characterize match location
        match_start, match_end = match.span()
        # -- matched separator is replaced by single-space in overlap string --
        separator = " "

        # -- in multi-space situation, fragment may have trailing whitespace because match is from
        # -- right to left
        fragment = s[:match_start].rstrip()
        # -- remainder can have leading space when match is on "\n" followed by spaces --
        raw_remainder = s[match_end:].lstrip()

        if overlap <= len(separator):
            return fragment, raw_remainder

        # -- compute overlap --
        tail_len = overlap - len(separator)
        tail = fragment[-tail_len:].lstrip()
        overlapped_remainder = tail + separator + raw_remainder
        return fragment, overlapped_remainder


class _CellAccumulator:
    """Incrementally build `<table>` fragment cell-by-cell to maximally fill chunking window.

    Accumulate cells until chunking window is filled, then generate the text and HTML for the
    subtable composed of all those rows that fit in the window.
    """

    def __init__(self, maxlen: int):
        self._maxlen = maxlen
        self._cells: list[HtmlCell] = []

    def add_cell(self, cell: HtmlCell) -> None:
        """Add `cell` to this accumulation. Caller is responsible for ensuring it will fit."""
        self._cells.append(cell)

    def flush(self) -> Iterator[TextAndHtml]:
        """Generate zero-or-one (text, html) pairs for accumulated sub-sub-table."""
        if not self._cells:
            return
        text = " ".join(self._iter_cell_texts())
        tds_str = "".join(c.html for c in self._cells)
        html = f"<table><tr>{tds_str}</tr></table>"
        self._cells.clear()
        yield text, html

    def will_fit(self, cell: HtmlCell) -> bool:
        """True when `cell` will fit within remaining space left by accummulated cells."""
        return self._remaining_space >= len(cell.text)

    def _iter_cell_texts(self) -> Iterator[str]:
        """Generate contents of each accumulated cell as a separate string.

        A cell that is empty or contains only whitespace does not generate a string.
        """
        for cell in self._cells:
            if not (text := cell.text):
                continue
            yield text

    @property
    def _remaining_space(self) -> int:
        """Number of characters remaining when text of accumulated cells is joined."""
        # -- separators are one space (" ") at the end of each cell's text, including last one to
        # -- account for space before prospective next cell.
        separators_len = len(self._cells)
        return self._maxlen - separators_len - sum(len(c.text) for c in self._cells)


class _RowAccumulator:
    """Maybe `SubtableAccumulator`.

    Accumulate rows until chunking window is filled, then generate the text and HTML for the
    subtable composed of all those rows that fit in the window.
    """

    def __init__(self, maxlen: int, measure: Callable[[str], int] = len):
        self._maxlen = maxlen
        self._measure = measure
        self._rows: list[HtmlRow] = []
        self._row_text_len = 0

    def add_row(self, row: HtmlRow) -> None:
        """Add `row` to this accumulation. Caller is responsible for ensuring it will fit."""
        self._rows.append(row)
        self._row_text_len += self._measured_row_text_len(row)

    def flush(self) -> Iterator[TextAndHtml]:
        """Generate zero-or-one (text, html) pairs for accumulated sub-table."""
        if not self._rows:
            return
        text = " ".join(self._iter_cell_texts())
        trs_str = "".join(r.html for r in self._rows)
        html = f"<table>{trs_str}</table>"
        self._rows.clear()
        self._row_text_len = 0
        yield text, html

    def will_fit(self, row: HtmlRow) -> bool:
        """True when `row` will fit within remaining space left by accummulated rows."""
        return self._remaining_space >= self._measured_row_text_len(row)

    def _iter_cell_texts(self) -> Iterator[str]:
        """Generate contents of each row cell as a separate string.

        A cell that is empty or contains only whitespace does not generate a string.
        """
        for r in self._rows:
            yield from r.iter_cell_texts()

    @property
    def _remaining_space(self) -> int:
        """Number of chunk-size units remaining for accumulated row text."""
        # -- separators are one space (" ") at the end of each row's text, including last one to
        # -- account for space before prospective next row.
        separators_len = len(self._rows)
        return self._maxlen - separators_len - self._row_text_len

    def _measured_row_text_len(self, row: HtmlRow) -> int:
        """Length of `row` text in configured chunk-size units."""
        return self._measure(" ".join(row.iter_cell_texts()))


# ================================================================================================
# PRE-CHUNK COMBINER
# ================================================================================================


class PreChunkCombiner:
    """Filters pre-chunk stream to combine small pre-chunks where possible."""

    def __init__(self, pre_chunks: Iterable[PreChunk], opts: ChunkingOptions):
        self._pre_chunks = pre_chunks
        self._opts = opts

    def iter_combined_pre_chunks(self) -> Iterator[PreChunk]:
        """Generate pre-chunk objects, combining `PreChunk` objects when they'll fit in window."""
        accum = _PreChunkAccumulator(self._opts)

        for pre_chunk in self._pre_chunks:
            # -- finish accumulating pre-chunk when it's full --
            if not accum.will_fit(pre_chunk):
                yield from accum.flush()

            accum.add_pre_chunk(pre_chunk)

        yield from accum.flush()


class _PreChunkAccumulator:
    """Accumulates, measures, and combines pre-chunks.

    Used for combining pre-chunks for chunking strategies like "by-title" that can potentially
    produce undersized chunks and offer the `combine_text_under_n_chars` option.

    Provides `.add_pre_chunk()` allowing a pre-chunk to be added to the chunk and provides
    monitoring properties `.remaining_space` and `.text_length` suitable for deciding whether to add
    another pre-chunk.

    `.flush()` is used to combine the accumulated pre-chunks into a single `PreChunk` object.
    This method returns an interator that generates zero-or-one `PreChunk` objects and is used
    like so:

        yield from accum.flush()

    If no pre-chunks have been accumulated, no `PreChunk` is generated. Flushing the builder
    clears the pre-chunks it contains so it is ready to accept the next pre-chunk.
    """

    def __init__(self, opts: ChunkingOptions) -> None:
        self._opts = opts
        self._pre_chunk: PreChunk | None = None

    def add_pre_chunk(self, pre_chunk: PreChunk) -> None:
        """Add a pre-chunk to the accumulator for possible combination with next pre-chunk."""
        self._pre_chunk = (
            pre_chunk if self._pre_chunk is None else self._pre_chunk.combine(pre_chunk)
        )

    def flush(self) -> Iterator[PreChunk]:
        """Generate accumulated pre-chunk as a single combined pre-chunk.

        Does not generate a pre-chunk when none has been accumulated.
        """
        # -- nothing to do if no pre-chunk has been accumulated --
        if not self._pre_chunk:
            return
        # -- otherwise generate the combined pre-chunk --
        yield self._pre_chunk
        # -- and reset the accumulator (to empty) --
        self._pre_chunk = None

    def will_fit(self, pre_chunk: PreChunk) -> bool:
        """True when there is room for `pre_chunk` in accumulator.

        An empty accumulator always has room. Otherwise there is only room when `pre_chunk` can be
        combined with any other pre-chunks in the accumulator without exceeding the combination
        limits specified for the chunking run.
        """
        # -- an empty accumulator always has room --
        if self._pre_chunk is None:
            return True

        return self._pre_chunk.can_combine(pre_chunk)


# ================================================================================================
# CHUNK BOUNDARY PREDICATES
# ------------------------------------------------------------------------------------------------
# A *boundary predicate* is a function that takes an element and returns True when the element
# represents the start of a new semantic boundary (such as section or page) to be respected in
# chunking.
#
# Some of the functions below *are* a boundary predicate and others *construct* a boundary
# predicate.
#
# These can be mixed and matched to produce different chunking behaviors like "by_title" or left
# out altogether to produce "basic-chunking" behavior.
#
# The effective lifetime of the function that produce a predicate (rather than directly being one)
# is limited to a single element-stream because these retain state (e.g. current page number) to
# determine when a semantic boundary has been crossed.
# ================================================================================================


def is_on_next_page() -> BoundaryPredicate:
    """Not a predicate itself, calling this returns a predicate that triggers on each new page.

    The lifetime of the returned callable cannot extend beyond a single element-stream because it
    stores current state (current page-number) that is particular to that element stream.

    The returned predicate tracks the "current" page-number, starting at 1. An element with a
    greater page number returns True, indicating the element starts a new page boundary, and
    updates the enclosed page-number ready for the next transition.

    An element with `page_number == None` or a page-number lower than the stored value is ignored
    and returns False.
    """
    current_page_number: int = 1
    is_first: bool = True

    def page_number_incremented(element: Element) -> bool:
        nonlocal current_page_number, is_first

        page_number = element.metadata.page_number

        # -- The first element never reports a page break, it starts the first page of the
        # -- document. That page could be numbered (page_number is non-None) or not. If it is not
        # -- numbered we assign it page-number 1.
        if is_first:
            current_page_number = page_number or 1
            is_first = False
            return False

        # -- An element with a `None` page-number is assumed to continue the current page. It never
        # -- updates the current-page-number because once set, the current-page-number is "sticky"
        # -- until replaced by a different explicit page-number.
        if page_number is None:
            return False

        if page_number == current_page_number:
            return False

        # -- it's possible for a page-number to decrease. We don't expect that, but if it happens
        # -- we consider it a page-break.
        current_page_number = page_number
        return True

    return page_number_incremented


def is_title(element: Element) -> bool:
    """True when `element` is a `Title` element, False otherwise."""
    return isinstance(element, Title)
