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

import itertools
import sys
import warnings
from collections.abc import Callable, Iterable
from typing import Any, Literal

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

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

Converter = docling.document_converter.DocumentConverter


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

    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,
    ):
        """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.
        :param condition: Optional callable that determines whether to process each document.
        """
        super().__init__(task_id=task_id, include_meta=include_meta, batch_size=batch_size, condition=condition)
        self._converter = converter if converter else docling.document_converter.DocumentConverter()
        self._export_format = export_format

        # Set batch size in Docling performance settings.
        docling.datamodel.settings.settings.perf.doc_batch_size = (
            self._batch_size if self._batch_size > 0 else sys.maxsize
        )

    def _validate_docs(self, docs: Iterable[Doc]) -> None:
        """Validate documents.

        Raises warning if docs already have text values, as they will be overwritten.

        :param docs: Docs to validate.
        :raises ValueError: If documents do not have a URI.
        """
        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.")

    def _call(self, docs: Iterable[Doc]) -> Iterable[Doc]:
        """Parse resources using docling.

        :param docs: Resources to process.
        :return: Parsed documents.
        :raises ValueError: If documents failed to parse.
        """
        docs_iters = itertools.tee(docs, 3)
        self._validate_docs(docs_iters[0])

        for doc, parsed_resource in zip(
            docs_iters[2], self._converter.convert_all([resource.uri for resource in docs_iters[1]])
        ):
            try:
                if self._include_meta:
                    doc.meta |= {self.id: parsed_resource}
                if self._export_format == "html":
                    doc.text = parsed_resource.document.export_to_html()
                else:
                    doc.text = parsed_resource.document.export_to_markdown()

            except Exception as e:
                logger.error(f"Failed to parse file {doc.uri}: {str(e)}")
                raise e

            yield doc

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