"""Allows chunking of documents into segments."""

import itertools
import sys
from collections.abc import Callable, Iterable
from typing import Any

import chonkie

from sieves.data.doc import Doc
from sieves.tasks.core import Task


class Chonkie(Task):
    """Chunker wrapping the chonkie library."""

    def __init__(
        self,
        chunker: chonkie.BaseChunker,
        task_id: str | None = None,
        include_meta: bool = False,
        batch_size: int = -1,
        condition: Callable[[Doc], bool] | None = None,
    ):
        """Initialize chunker.

        :param task_id: Task ID.
        :param include_meta: Whether to include meta information generated by the task.
        :param batch_size: Batch size to use for processing. Use -1 to process all documents at once.
        :param condition: Optional callable that determines whether to process each document.
        """
        super().__init__(task_id=task_id, include_meta=include_meta, batch_size=batch_size, condition=condition)
        self._chunker = chunker

    def _call(self, docs: Iterable[Doc]) -> Iterable[Doc]:
        """Split documents into chunks.

        :param docs: Documents to split.
        :return: Split documents.
        """
        batch_size = self._batch_size if self._batch_size > 0 else sys.maxsize
        while docs_batch := [doc for doc in itertools.islice(docs, batch_size)]:
            if len(docs_batch) == 0:
                break

            # Signature varies in between versions, try both.
            try:
                chunks = self._chunker.chunk_batch([doc.text for doc in docs_batch], show_progress=False)
            except TypeError:
                chunks = self._chunker.chunk_batch([doc.text for doc in docs_batch], show_progress_bar=False)  # type: ignore[unknown-argument]
            assert len(chunks) == len(docs_batch)

            for doc, doc_chunks in zip(docs_batch, chunks):
                if self._include_meta:
                    doc.meta |= {self.id: {doc_chunks}}
                doc.chunks = [chunk.text for chunk in doc_chunks]

                yield doc

    @property
    def _state(self) -> dict[str, Any]:
        return {
            **super()._state,
            "chunker": self._chunker,
        }
