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

import re
from collections.abc import 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,
    ):
        """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.
        """
        super().__init__(task_id=task_id, include_meta=include_meta)
        self._interval = interval

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

        :param docs: Documents to split.
        :return: Split documents.
        """
        docs = list(docs)

        for doc in docs:
            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)]

        return docs

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