"""File parsers for converting raw files into documents."""
import re
from collections.abc import Iterable

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.
        :returns: 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()


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.
        :returns: 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
