"""File preprocessing for converting raw files into documents."""
import warnings
from collections.abc import Callable, Iterable
from typing import Any

import nltk
import unstructured
import unstructured.documents.elements
import unstructured.partition.auto
from tqdm import tqdm

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

PartitionType = Callable[..., list[unstructured.documents.elements.Text]]
CleanerType = Callable[[str], str]


class Unstructured(Task):
    """Parser wrapping the unstructured library to convert files into documents."""

    def __init__(
        self,
        partition: PartitionType = unstructured.partition.auto.partition,
        cleaners: tuple[CleanerType, ...] = (),
        task_id: str | None = None,
        show_progress: bool = True,
        include_meta: bool = False,
        **kwargs: dict[str, Any],
    ):
        """Initialize the docling parser.
        :param partition: Function to use for partitioning.
        :param cleaners: Cleaning functions to apply.
        :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.
        :param kwargs: Kwargs to be supplied to partitioning call.
        """
        super().__init__(task_id=task_id, show_progress=show_progress, include_meta=include_meta)
        self._partition = partition
        self._partition_args = kwargs or {}
        self._cleaners = cleaners

        Unstructured._require()

    @staticmethod
    def _require() -> None:
        """Download all necessary resources that have to be installed from within Python."""
        # Some nltk resources seem necessary for basic functionality.
        for nltk_resource in ("punkt_tab", "averaged_perceptron_tagger_eng"):
            # Don't install if already available.
            try:
                nltk.data.find(nltk_resource)
            except LookupError:
                nltk.download(nltk_resource)

    def __call__(self, docs: Iterable[Doc]) -> Iterable[Doc]:
        """Parse resources using docling.
        :param docs: Resources to process.
        :return: Parsed documents.
        """
        docs = list(docs)

        # Validate docs.
        have_text = False
        for doc in docs:
            assert doc.uri, ValueError("Documents have to have a value for .uri.")
            if doc.text:
                have_text = True
        if have_text:
            warnings.warn(f"Task {self._task_id} is about to overwrite existing .text values.")

        # Wrap conversion in TQDM if progress should be shown.
        iterable = tqdm(docs, total=len(docs)) if self._show_progress else docs
        does_chunking = "chunking_strategy" in self._partition_args

        # import unstructured.documents.elements.Te
        for doc in iterable:
            try:
                # Parse and process document.
                parsed_resources: list[unstructured.documents.elements.Text] = self._partition(
                    doc.uri, **self._partition_args
                )

                # Apply specified cleaners.
                for cleaner in self._cleaners:
                    for pr in parsed_resources:
                        pr.apply(cleaner)

                # Integrate into Doc instances.
                if self._include_meta:
                    doc.meta |= {self.id: parsed_resources}

                # Use chunks.
                if does_chunking:
                    doc.chunks = [pr.text for pr in parsed_resources]

                # Merge texts from all elements into single string for the entire document.
                doc.text = "\n".join(resource.text for resource in parsed_resources)

            except FileNotFoundError as err:
                raise FileNotFoundError(
                    f"File at {doc.uri} not found. Ensure that this is a local file path - unstructured doesn't support"
                    f" loading files via network URIs."
                ) from err

        return docs

    @property
    def _state(self) -> dict[str, Any]:
        return {
            **super()._state,
            "partition": self._partition,
            "cleaners": self._cleaners,
            **self._partition_args,
        }
