"""Allows chunking of documents into segments."""
import re
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,
        }


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

    def __init__(
        self,
        interval: int,
        task_id: str | None = None,
        show_progress: bool = True,
        include_meta: bool = False,
    ):
        """Initialize chunker.
        :param interval: Token count interval for chunks.
        :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._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)

        pbar: tqdm | None = tqdm(total=len(docs)) if self._show_progress else None
        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)]
            if pbar:
                pbar.update(1)
        if pbar:
            pbar.close()

        return docs

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