"""Ingestion task implementation."""

from __future__ import annotations

from collections.abc import Callable, Iterable
from types import NoneType
from typing import Any, Literal

from sieves.data.doc import Doc
from sieves.serialization import Config
from sieves.tasks.core import Task
from sieves.tasks.preprocessing.ingestion import docling_, marker_

Converter = docling_.Converter | marker_.Converter


class Ingestion(Task):
    """Base class for Ingestion tasks that extract text from documents.

    This unified interface allows different Ingestion converters to be used interchangeably.
    """

    def __init__(
        self,
        converter: Converter | None = None,
        export_format: Literal["markdown", "html"] = "markdown",
        task_id: str | None = None,
        include_meta: bool = False,
        batch_size: int = -1,
        condition: Callable[[Doc], bool] | None = None,
        **kwargs: Any,
    ):
        """Initialize the Ingestion task.

        :param converter: The Ingestion converter to use. If None, tries to initialize in this order:
            - A `docling` converter with default values
            - A `marker` converter with default values
            If neither is possible due to a lack of installed requirements, raises a `ModuleNotFoundError`.
        :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.
        :param kwargs: Additional arguments for specific `Ingestion` implementations.
        """
        super().__init__(task_id=task_id, include_meta=include_meta, batch_size=batch_size, condition=condition)
        self._export_format = export_format
        self._converter = converter
        self._kwargs = kwargs
        self._task = self._init_ingestion_task()

    def _init_ingestion_task(self) -> Task:
        """Initialize the bridge for the specific Ingestion implementation.

        :raises ValueError: On unsupported converter type.
        :return: Ingestion bridge implementation.
        """
        converter_type = type(self._converter)

        # Set converter type manually in predefined order, depending on which dependencies are installed.
        if self._converter is None:
            if docling_ is not None:
                converter_type = docling_.Converter
            elif marker_ is not None:
                converter_type = marker_.Converter.__args__[0]
            else:
                raise ModuleNotFoundError(
                    "None of the supported ingestion libraries (`docling`, `marker`) are installed. Install at least "
                    "one of them to run ingestion tasks."
                )
        assert converter_type is not NoneType

        # Identify the ingestion task that maps to the specified converter's type.
        converter_module_map = {
            docling_: getattr(docling_, "Docling", None),
            marker_: getattr(marker_, "Marker", None),
        }

        for module, ingestion_task_type in converter_module_map.items():
            if ingestion_task_type is None:
                continue

            assert hasattr(module, "Converter")
            try:
                module_converter_types = module.Converter.__args__
            except AttributeError:
                module_converter_types = (module.Converter,)

            if any(issubclass(converter_type, module_model_type) for module_model_type in module_converter_types):
                ingestion_task = ingestion_task_type(
                    converter=self._converter,
                    export_format=self._export_format,
                    task_id=self.id,
                    include_meta=self._include_meta,
                    batch_size=self._batch_size,
                    **self._kwargs,
                )
                assert isinstance(ingestion_task, Task)

                return ingestion_task

        raise ValueError(
            f"converter type {self._converter} is not supported. Please check the documentation and ensure that (1) "
            f"you're providing a supported converter type and that (2) the corresponding library is installed in your "
            f"environment."
        )

    def _call(self, docs: Iterable[Doc]) -> Iterable[Doc]:
        """Process documents with Ingestion to extract text.

        :param docs: Documents to process.
        :return: Processed documents with extracted text.
        """
        docs = list(docs)
        assert all(doc.uri for doc in docs), ValueError("Documents have to have a value for .uri.")
        result = self._task(docs)

        yield from result

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

        :return: Dict of attributes to serialize.
        """
        return {
            **super()._state,
            "converter": self._converter,
            "export_format": self._export_format,
            **self._kwargs,
        }

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

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