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

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

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


class NaiveChunker(Task):
    """Chunks by sentence counts. Only for test purposes."""

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

        :param interval: Token count interval for chunks.
        :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._interval = interval

    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

            for doc in docs_batch:
                assert doc.text
                sentences = [sent for sent in re.split("[?!.]", doc.text) if len(sent.strip())]
                doc.chunks = [
                    ".".join(sentences[i : i + self._interval]) for i in range(0, len(sentences), self._interval)
                ]

                yield doc

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