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

import itertools
import sys
from collections.abc import 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,
    ):
        """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.
        """
        super().__init__(task_id=task_id, include_meta=include_meta, batch_size=batch_size)
        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

            chunks = self._chunker.chunk_batch([doc.text for doc in docs_batch], show_progress_bar=False)
            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,
        }
