"""Handles dispatch of elements to a chunking-strategy by name.

Also provides the `@add_chunking_strategy` decorator which is the chief current user of "by-name"
chunking dispatch.
"""

from __future__ import annotations

import copy
import dataclasses as dc
import functools
import inspect
from functools import cached_property
from typing import Any, Callable, Iterable, Optional, Protocol

from lxml.etree import ParserError, tostring
from lxml.html import fragment_fromstring
from typing_extensions import ParamSpec

from unstructured.chunking.basic import chunk_elements
from unstructured.chunking.title import chunk_by_title
from unstructured.documents.elements import Element, Table, TableChunk
from unstructured.utils import get_call_args_applying_defaults

_P = ParamSpec("_P")


class Chunker(Protocol):
    """Abstract interface for chunking functions."""

    def __call__(
        self, elements: Iterable[Element], *, max_characters: Optional[int]
    ) -> list[Element]:
        """A chunking function must have this signature.

        In particular it must minimally have an `elements` parameter and all chunkers will have a
        `max_characters` parameter (doesn't need to follow `elements` directly). All others can
        vary by chunker.
        """
        ...


def add_chunking_strategy(func: Callable[_P, list[Element]]) -> Callable[_P, list[Element]]:
    """Decorator for chunking text.

    Chunks the element sequence produced by the partitioner it decorates when a `chunking_strategy`
    argument is present in the partitioner call and it names an available chunking strategy.
    """
    # -- Patch the docstring of the decorated function to add chunking strategy and
    # -- chunking-related argument documentation. This only applies when `chunking_strategy`
    # -- is an explicit argument of the decorated function and "chunking_strategy" is not
    # -- already mentioned in the docstring.
    if func.__doc__ and (
        "chunking_strategy" in func.__code__.co_varnames and "chunking_strategy" not in func.__doc__
    ):
        func.__doc__ += (
            "\nchunking_strategy"
            + "\n\tStrategy used for chunking text into larger or smaller elements."
            + "\n\tDefaults to `None` with optional arg of 'basic' or 'by_title'."
            + "\n\tAdditional Parameters:"
            + "\n\t\tmultipage_sections"
            + "\n\t\t\tIf True, sections can span multiple pages. Defaults to True."
            + "\n\t\tcombine_text_under_n_chars"
            + "\n\t\t\tCombines elements (for example a series of titles) until a section"
            + "\n\t\t\treaches a length of n characters. Only applies to 'by_title' strategy."
            + "\n\t\tnew_after_n_chars"
            + "\n\t\t\tCuts off chunks once they reach a length of n characters; a soft max."
            + "\n\t\tmax_characters"
            + "\n\t\t\tChunks elements text and text_as_html (if present) into chunks"
            + "\n\t\t\tof length n characters, a hard max."
            + "\n\t\trepeat_table_headers"
            + "\n\t\t\tDefault: True. Repeat detected table headers on continuation"
            + "\n\t\t\ttable chunks. Set to False to opt out."
            + "\n\t\tskip_table_chunking"
            + "\n\t\t\tDefault: False. When True, Table elements are passed through"
            + "\n\t\t\tunchanged without being split into TableChunk elements."
            + "\n\t\tisolate_table"
            + "\n\t\t\tDefault: True. When True, Table/TableChunk elements are always"
            + "\n\t\t\tstaged in their own pre-chunk. Set to False to allow tables to"
            + "\n\t\t\tshare pre-chunks with adjacent non-table elements."
        )

    @functools.wraps(func)
    def wrapper(*args: _P.args, **kwargs: _P.kwargs) -> list[Element]:
        """The decorated function is replaced with this one."""

        # -- call the partitioning function to get the elements --
        elements = func(*args, **kwargs)

        # -- look for a chunking-strategy argument --
        call_args = get_call_args_applying_defaults(func, *args, **kwargs)
        chunking_strategy = call_args.pop("chunking_strategy", None)

        # -- no chunking-strategy means no chunking --
        if chunking_strategy is None:
            return elements

        # -- otherwise, chunk away :) --
        return chunk(elements, chunking_strategy, **call_args)

    return wrapper


def chunk(elements: Iterable[Element], chunking_strategy: str, **kwargs: Any) -> list[Element]:
    """Dispatch chunking of `elements` to the chunking function for `chunking_strategy`."""
    chunker_spec = _chunker_registry.get(chunking_strategy)

    if chunker_spec is None:
        raise ValueError(f"unrecognized chunking strategy {repr(chunking_strategy)}")

    # -- `kwargs` will in general be an omnibus dict of all keyword arguments to the partitioner;
    # -- pick out and use only those supported by this chunker.
    chunking_kwargs = {k: v for k, v in kwargs.items() if k in chunker_spec.kw_arg_names}

    return chunker_spec.chunker(elements, **chunking_kwargs)


def register_chunking_strategy(name: str, chunker: Chunker) -> None:
    """Make chunker available by using `name` as `chunking_strategy` arg in partitioner call."""
    _chunker_registry[name] = _ChunkerSpec(chunker)


@dc.dataclass(frozen=True)
class _ChunkerSpec:
    """A registry entry for a chunker."""

    chunker: Chunker
    """The "chunk_by_{x}() function that implements this chunking strategy."""

    @cached_property
    def kw_arg_names(self) -> tuple[str, ...]:
        """Keyword arguments supported by this chunker.

        These are all arguments other than the required `elements: list[Element]` first parameter.
        """
        sig = inspect.signature(self.chunker)
        return tuple(key for key in sig.parameters if key != "elements")


_chunker_registry: dict[str, _ChunkerSpec] = {
    "basic": _ChunkerSpec(chunk_elements),
    "by_title": _ChunkerSpec(chunk_by_title),
}


def reconstruct_table_from_chunks(elements: Iterable[Element]) -> list[Table]:
    """Reconstruct original tables from a mixed list of chunked elements.

    Filters `TableChunk` elements, groups them by `table_id`, orders by `chunk_index`, and
    merges each group into a single `Table` with combined text and HTML. Non-`TableChunk`
    elements are ignored. Returns reconstructed tables in reading order (order of first chunk
    appearance).
    """
    # -- filter to only TableChunk instances, preserving input order --
    table_chunks = [e for e in elements if isinstance(e, TableChunk)]
    if not table_chunks:
        return []

    # -- group by table_id, preserving first-seen order --
    groups: dict[str, list[TableChunk]] = {}
    for chunk in table_chunks:
        tid = chunk.metadata.table_id
        if tid is None:
            continue
        if tid not in groups:
            groups[tid] = []
        groups[tid].append(chunk)

    # -- sort each group by chunk_index and merge --
    tables: list[Table] = []

    def _chunk_sort_key(chunk: TableChunk) -> tuple[bool, int]:
        chunk_index = chunk.metadata.chunk_index
        return (chunk_index is None, 0 if chunk_index is None else chunk_index)

    for group in groups.values():
        group.sort(key=_chunk_sort_key)
        tables.append(_merge_table_chunks(group))

    return tables


def _merge_table_chunks(chunks: list[TableChunk]) -> Table:
    """Merge an ordered list of TableChunks from the same table into a single Table."""
    # -- combine text --
    text = " ".join(
        chunk_text for chunk in chunks if (chunk_text := _strip_carried_over_header_text(chunk))
    )

    # -- build metadata from first chunk --
    metadata = copy.deepcopy(chunks[0].metadata)
    metadata.is_continuation = None
    metadata.table_id = None
    metadata.chunk_index = None
    metadata.num_carried_over_header_rows = None

    # -- combine HTML if all chunks have it --
    if all(c.metadata.text_as_html for c in chunks):
        combined = fragment_fromstring("<table></table>")
        canonical_header_row_count, canonical_header_rows = _first_carried_header_rows(chunks)
        if canonical_header_rows:
            thead = fragment_fromstring("<thead></thead>")
            for row in canonical_header_rows:
                thead.append(row)
            combined.append(thead)

        for c in chunks:
            parsed = fragment_fromstring(c.metadata.text_as_html)
            carried_over_header_rows = _num_carried_over_header_rows(c)
            rows = parsed.xpath("./tr | ./thead/tr | ./tbody/tr | ./tfoot/tr")
            skip_count = carried_over_header_rows
            if c is chunks[0] and canonical_header_row_count:
                skip_count = canonical_header_row_count
            for row in rows[skip_count:]:
                combined.append(row)
        metadata.text_as_html = tostring(combined, encoding=str)
    else:
        metadata.text_as_html = None

    return Table(text=text, metadata=metadata)


def _num_carried_over_header_rows(chunk: TableChunk) -> int:
    """Header rows prepended synthetically to this chunk.

    Reconstruction can be called on user-provided/deserialized chunks, so treat missing values as
    "no carried header rows."
    """
    value = chunk.metadata.num_carried_over_header_rows
    return value or 0


def _first_carried_header_rows(chunks: list[TableChunk]) -> tuple[int, list[Any]]:
    """Header rows from first continuation chunk carrying repeated headers, if any."""
    first_chunk_rows = _top_level_table_rows(chunks[0].metadata.text_as_html)
    if first_chunk_rows is None:
        return 0, []

    for chunk in chunks:
        carried_row_count = _num_carried_over_header_rows(chunk)
        if carried_row_count <= 0:
            continue

        rows = _top_level_table_rows(chunk.metadata.text_as_html)
        if rows is None:
            continue

        if carried_row_count > len(rows):
            continue

        carried_rows = rows[:carried_row_count]
        if not _leading_row_texts_match(first_chunk_rows, carried_rows):
            continue

        return carried_row_count, [copy.deepcopy(row) for row in rows[:carried_row_count]]

    return 0, []


def _top_level_table_rows(text_as_html: str | None) -> list[Any] | None:
    """Top-level rows from a table fragment, preserving section ordering."""
    if not text_as_html:
        return None

    try:
        parsed = fragment_fromstring(text_as_html)
    except (ParserError, ValueError):
        return None

    return parsed.xpath("./tr | ./thead/tr | ./tbody/tr | ./tfoot/tr")


def _leading_row_texts_match(first_chunk_rows: list[Any], carried_rows: list[Any]) -> bool:
    """True when carried rows match first chunk's leading rows by normalized cell text."""
    if len(first_chunk_rows) < len(carried_rows):
        return False

    for first_row, carried_row in zip(first_chunk_rows, carried_rows):
        if _row_text_signature(first_row) != _row_text_signature(carried_row):
            return False

    return True


def _row_text_signature(row: Any) -> tuple[str, ...]:
    """Normalized cell text tuple for a row."""
    return tuple(" ".join(cell.text_content().split()) for cell in row.iter("td", "th"))


def _strip_carried_over_header_text(chunk: TableChunk) -> str:
    """Strip synthetic carried-over header text from continuation chunk text."""
    carried_row_count = _num_carried_over_header_rows(chunk)
    if carried_row_count == 0:
        return chunk.text

    text_as_html = chunk.metadata.text_as_html
    if not text_as_html:
        return chunk.text

    try:
        parsed = fragment_fromstring(text_as_html)
    except (ParserError, ValueError):
        return chunk.text

    rows = parsed.xpath("./tr | ./thead/tr | ./tbody/tr | ./tfoot/tr")
    if carried_row_count > len(rows):
        return chunk.text

    carried_header_text = " ".join(
        text
        for row in rows[:carried_row_count]
        for text in (" ".join(cell.text_content().split()) for cell in row.iter("td", "th"))
        if text
    )
    if not carried_header_text:
        return chunk.text

    chunk_text = chunk.text.lstrip()
    if chunk_text == carried_header_text:
        return ""
    if chunk_text.startswith(f"{carried_header_text} "):
        return chunk_text[len(carried_header_text) + 1 :]
    if chunk_text.startswith(carried_header_text):
        return chunk_text[len(carried_header_text) :].lstrip()
    return chunk.text