"""Wrapper for `Docling` for the conversion of complex files into markdown."""

import warnings
from collections.abc import Iterable
from typing import Any

import docling.datamodel.document
import docling.document_converter
from loguru import logger

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


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

    def __init__(
        self,
        converter: docling.document_converter.DocumentConverter | None = None,
        export_format: str = "markdown",
        task_id: str | None = None,
        include_meta: bool = False,
        batch_size: int = -1,
    ):
        """Initialize the docling parser.

        :param converter: Docling parser instance.
        :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.
        """
        super().__init__(task_id=task_id, include_meta=include_meta, batch_size=batch_size)
        self._converter = converter if converter else docling.document_converter.DocumentConverter()
        self._export_format = export_format

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

        parsed_resources: list[docling.datamodel.document.ConversionResult] = list(
            self._converter.convert_all([resource.uri for resource in docs])
        )
        assert len(parsed_resources) == len(docs)

        for doc, parsed_resource in zip(docs, parsed_resources):
            try:
                if self._include_meta:
                    doc.meta |= {self.id: parsed_resource}
                if self._export_format == "markdown":
                    doc.text = parsed_resource.document.export_to_markdown()
                elif self._export_format == "html":
                    doc.text = parsed_resource.document.export_to_html()
                elif self._export_format == "json":
                    doc.text = parsed_resource.document.export_to_dict()
            except Exception as e:
                logger.error(f"Failed to parse file {doc.uri}: {str(e)}")
                continue

        return docs

    @property
    def _state(self) -> dict[str, Any]:
        return {
            **super()._state,
            "converter": self._converter,
            "export_format": self._export_format,
        }
