"""Implementation of baseline chunking.

This is the "plain-vanilla" chunking strategy. All the fundamental chunking behaviors are present in
this strategy and also in all other strategies. Those are:

- Maximally fill each chunk with sequential elements.
- Isolate oversized elements and divide (only) those chunks by text-splitting.
- Overlap when requested.

"Fancier" strategies add higher-level semantic-unit boundaries to be respected. For example, in the
by-title strategy, section boundaries are respected, meaning a chunk never contains text from two
different sections. When a new section is detected the current chunk is closed and a new one
started.
"""

from __future__ import annotations

from typing import Iterable, Iterator, Optional

from unstructured.chunking.base import ChunkingOptions, PreChunker
from unstructured.documents.elements import Element


def chunk_elements(
    elements: Iterable[Element],
    *,
    include_orig_elements: Optional[bool] = None,
    max_characters: Optional[int] = None,
    max_tokens: Optional[int] = None,
    new_after_n_chars: Optional[int] = None,
    new_after_n_tokens: Optional[int] = None,
    overlap: Optional[int] = None,
    overlap_all: Optional[bool] = None,
    tokenizer: Optional[str] = None,
    repeat_table_headers: Optional[bool] = None,
    skip_table_chunking: Optional[bool] = None,
    isolate_table: Optional[bool] = None,
) -> list[Element]:
    """Combine sequential `elements` into chunks, respecting specified text-length limits.

    Produces a sequence of `CompositeElement`, `Table`, and `TableChunk` elements (chunks).

    Parameters
    ----------
    elements
        A list of unstructured elements. Usually the output of a partition function.
    include_orig_elements
        When `True` (default), add elements from pre-chunk to the `.metadata.orig_elements` field
        of the chunk(s) formed from that pre-chunk. Among other things, this allows access to
        original-element metadata that cannot be consolidated and is dropped in the course of
        chunking.
    max_characters
        Hard maximum chunk length. No chunk will exceed this length. A single element that exceeds
        this length will be divided into two or more chunks using text-splitting. Mutually
        exclusive with `max_tokens`.
    max_tokens
        Hard maximum chunk token count. No chunk will exceed this token count. Requires `tokenizer`
        to be specified. Mutually exclusive with `max_characters`.
    new_after_n_chars
        A chunk that of this length or greater is not extended to include the next element, even if
        that element would fit without exceeding `max_characters`. A "soft max" length that can be
        used in conjunction with `max_characters` to limit most chunks to a preferred length while
        still allowing larger elements to be included in a single chunk without resorting to
        text-splitting. Defaults to `max_characters` when not specified, which effectively disables
        any soft window. 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`. A chunk with this token count or greater is
        not extended. Requires `max_tokens` and `tokenizer` to be specified.
    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 produces a certain
        level of "pollution" of otherwise clean semantic chunk boundaries.
    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`.
    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.
    skip_table_chunking
        Default: `False`. When `True`, `Table` elements are passed through unchanged without
        being split into `TableChunk` elements, regardless of their size.
    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).
    """
    # -- raises ValueError on invalid parameters --
    opts = _BasicChunkingOptions.new(
        include_orig_elements=include_orig_elements,
        max_characters=max_characters,
        max_tokens=max_tokens,
        new_after_n_chars=new_after_n_chars,
        new_after_n_tokens=new_after_n_tokens,
        overlap=overlap,
        overlap_all=overlap_all,
        tokenizer=tokenizer,
        repeat_table_headers=repeat_table_headers,
        skip_table_chunking=skip_table_chunking,
        isolate_table=isolate_table,
    )

    return _chunk_elements(elements, opts)


def iter_chunk_elements(
    elements: Iterable[Element],
    *,
    include_orig_elements: Optional[bool] = None,
    max_characters: Optional[int] = None,
    max_tokens: Optional[int] = None,
    new_after_n_chars: Optional[int] = None,
    new_after_n_tokens: Optional[int] = None,
    overlap: Optional[int] = None,
    overlap_all: Optional[bool] = None,
    tokenizer: Optional[str] = None,
    repeat_table_headers: Optional[bool] = None,
    skip_table_chunking: Optional[bool] = None,
    isolate_table: Optional[bool] = None,
) -> Iterator[Element]:
    """Lazy `chunk_elements`: yields each chunk as it is formed instead of returning a list.

    Accepts the same options and produces the same chunks in the same order; see
    `chunk_elements` for the parameter documentation. Prefer this form when `elements` is
    itself lazy and the chunks are consumed one at a time, e.g. written straight to a file.
    The formed chunks are then never accumulated in a list, and the elements held at once are
    bounded by the pre-chunk being formed (plus the one element that closes it) rather than by
    the whole document.

    That bound covers the elements this function holds, not what each chunk carries. Under the
    default `include_orig_elements=True` a chunk keeps the elements it was formed from in
    `metadata.orig_elements`, so their `metadata.image_base64` payloads live as long as the chunk
    does -- shared with the source element for text, duplicated for a table, which is deep-copied.
    Streaming alone therefore gives little relief when those payloads are the bottleneck; pass
    `include_orig_elements=False` as well.

    Chunking has always been lazy internally; this exposes that pipeline rather than adding
    a second one. Options are still validated eagerly, when this function is called, not
    when the returned iterator is first advanced, so an invalid combination -- or an unknown
    `tokenizer`, when chunking by `max_tokens` -- raises here.
    """
    # -- raises ValueError on invalid parameters --
    opts = _BasicChunkingOptions.new(
        include_orig_elements=include_orig_elements,
        max_characters=max_characters,
        max_tokens=max_tokens,
        new_after_n_chars=new_after_n_chars,
        new_after_n_tokens=new_after_n_tokens,
        overlap=overlap,
        overlap_all=overlap_all,
        tokenizer=tokenizer,
        repeat_table_headers=repeat_table_headers,
        skip_table_chunking=skip_table_chunking,
        isolate_table=isolate_table,
    )

    return _iter_chunk_elements(elements, opts)


def _chunk_elements(elements: Iterable[Element], opts: _BasicChunkingOptions) -> list[Element]:
    """Implementation of actual basic chunking."""
    # -- Note(scanny): it might seem like over-abstraction for this to be a separate function but
    # -- it eases overriding or adding individual chunking options when customizing a stock chunker.
    return list(_iter_chunk_elements(elements, opts))


def _iter_chunk_elements(
    elements: Iterable[Element], opts: _BasicChunkingOptions
) -> Iterator[Element]:
    """Lazy implementation of actual basic chunking, shared with `_chunk_elements`."""
    for pre_chunk in PreChunker.iter_pre_chunks(elements, opts):
        yield from pre_chunk.iter_chunks()
        # -- release the emitted pre-chunk before advancing. The loop variable would otherwise
        # -- keep its elements alive while the next pre-chunk fills, so a streaming caller would
        # -- hold two pre-chunks' worth of elements at each transition instead of one.
        del pre_chunk


class _BasicChunkingOptions(ChunkingOptions):
    """Options for `basic` chunking."""
