"""Allows chunking of documents into segments."""
from collections.abc import Iterable
from typing import Any

import chonkie
from tqdm import tqdm

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,
        show_progress: bool = True,
        include_meta: bool = False,
    ):
        """Initialize chunker.
        :param task_id: Task ID.
        :param show_progress: Whether to show progress bar for processed documents.
        :param include_meta: Whether to include meta information generated by the task.
        """
        super().__init__(task_id=task_id, show_progress=show_progress, include_meta=include_meta)
        self._chunker = chunker

    def __call__(self, docs: Iterable[Doc]) -> Iterable[Doc]:
        """Split documents into chunks.
        :param docs: Documents to split.
        :return: Split documents.
        """
        docs = list(docs)

        pbar: tqdm | None = tqdm(total=len(docs)) if self._show_progress else None
        try:
            chunks = self._chunker.chunk_batch([doc.text for doc in docs])
            assert len(chunks) == len(docs)

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

            return docs
        finally:
            if pbar:
                pbar.close()

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