"""Ingestion task implementation."""

from __future__ import annotations

from collections.abc import Iterable
from typing import Any

import docling
import docling.document_converter
import marker
from marker.converters.pdf import PdfConverter
from marker.converters.table import TableConverter

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_

_ConverterType = docling.document_converter.DocumentConverter | PdfConverter | TableConverter


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: _ConverterType = docling.document_converter.DocumentConverter(),
        export_format: str = "markdown",
        task_id: str | None = None,
        include_meta: bool = False,
        batch_size: int = -1,
        **kwargs: Any,
    ):
        """Initialize the Ingestion task.

        :param converter: The Ingestion converter to use.
        :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 kwargs: Additional arguments for specific Ingestion implementations.
        """
        super().__init__(task_id=task_id, include_meta=include_meta, batch_size=batch_size)
        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.

        :return: Ingestion bridge implementation.
        """
        converter_type = type(self._converter)
        ingestion_task: Task
        match converter_type:
            case converter if issubclass(
                converter, (marker.converters.pdf.PdfConverter | marker.converters.table.TableConverter)
            ):
                ingestion_task = marker_.Marker(
                    converter=self._converter,
                    export_format=self._export_format,
                    task_id=self.id,
                    include_meta=self._include_meta,
                    batch_size=self._batch_size,
                    **self._kwargs,
                )
            case docling.document_converter.DocumentConverter:
                ingestion_task = docling_.Docling(
                    converter=self._converter,
                    export_format=self._export_format,
                    task_id=self.id,
                    include_meta=self._include_meta,
                    batch_size=self._batch_size,
                )
            case _:
                raise ValueError(
                    f"converter type {self._converter} is not supported. Please check the documentation "
                    f"and ensure you're providing a supported converter type."
                )
        assert isinstance(ingestion_task, Task)
        return ingestion_task

    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))
