"""Chunking task."""

from __future__ import annotations

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

import chonkie

from sieves.data.doc import Doc
from sieves.serialization import Config
from sieves.tasks.core import Task
from sieves.tasks.preprocessing import chunking
from sieves.tasks.preprocessing.chunking import chonkie_, naive

_ChunkerArgType = chonkie.BaseChunker | int
_ChunkerType = chonkie_.Chonkie | naive.NaiveChunker


class Chunking(Task):
    """Task for chunking documents using different strategies.

    This task acts as a wrapper around specific chunker implementations,
    allowing for flexible configuration based on the provided chunker object or interval.
    """

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

        :param chunker: The chunker instance (chonkie.BaseChunker) or the interval (int) for NaiveChunker.
        :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_arg = chunker
        self._task = self._init_chunker_task()

    def _init_chunker_task(self) -> _ChunkerType:
        """Initialize the specific chunker task based on the type of _chunker_arg.

        :return: Initialized chunker task instance.
        :raises TypeError: If the type of _chunker_arg is not supported.
        """
        chunker_task: _ChunkerType

        match self._chunker_arg:
            case chunker if isinstance(chunker, chonkie.BaseChunker):
                chunker_task = chunking.chonkie_.Chonkie(
                    chunker=chunker,
                    task_id=self.id,
                    include_meta=self._include_meta,
                    batch_size=self._batch_size,
                )
            case interval if isinstance(interval, int):
                chunker_task = chunking.naive.NaiveChunker(
                    interval=interval,
                    task_id=self.id,
                    include_meta=self._include_meta,
                    batch_size=self._batch_size,
                )
            case _:
                raise TypeError(
                    f"Unsupported type for 'chunker' argument: {type(self._chunker_arg)}. "
                    f"Expected chonkie.BaseChunker or int."
                )

        return chunker_task

    def _call(self, docs: Iterable[Doc]) -> Iterable[Doc]:
        """Process documents by chunking their text.

        :param docs: Documents to process.
        :return: Processed documents with chunks added.
        """
        docs_iters = itertools.tee(docs, 2)
        assert all(doc.text for doc in docs_iters[0]), ValueError("Documents have to have a value for .text.")
        yield from self._task(docs_iters[1])

    @property
    def _state(self) -> dict[str, Any]:
        """Return attributes to serialize.

        :return: Dict of attributes to serialize.
        """
        return {
            **super()._state,
            "chunker": self._chunker_arg,
        }

    @classmethod
    def deserialize(cls, config: Config, **kwargs: dict[str, Any]) -> Chunking:
        """Generate Chunker instance from config.

        :param config: Config to generate instance from.
        :param kwargs: Values to inject into loaded config.
        :return: Deserialized Chunker instance.
        """
        return cls(**config.to_init_dict(cls, **kwargs))
