"""Marker task for converting PDF documents to text."""
# ruff: noqa: E402

import warnings
from collections.abc import Callable, Iterable
from pathlib import Path
from typing import Any

# `marker` is using outdated Pydantic config approach, suppress these warnings.
warnings.filterwarnings(
    "ignore",
    message="Support for class-based `config` is deprecated.*ConfigDict instead",
    category=DeprecationWarning,
)

from marker.converters.pdf import PdfConverter
from marker.converters.table import TableConverter
from marker.models import create_model_dict
from marker.output import text_from_rendered

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

Converter = PdfConverter | TableConverter


class Marker(Task):
    """Marker task for converting PDF documents to text."""

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

        :param converter: Custom PdfConverter or TableConverter instance. If None, a default one will be created.
        :param export_format: Format to export the document in ("markdown", "html", or "json").
        :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 extract_images: Whether to extract images from the PDF.
        :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._export_format = export_format
        self._converter = self._setup_converter(converter, self._export_format)
        self._extract_images = extract_images

    def _setup_converter(
        self, converter: PdfConverter | TableConverter | None, export_format: str
    ) -> PdfConverter | TableConverter:
        """Set up the converter with the specified renderer.

        :param converter: Custom converter instance or None.
        :param export_format: Format to export the document in.
        :return: Configured converter instance.
        """
        renderer: str = self._get_renderer(export_format)
        if converter is None:
            return PdfConverter(artifact_dict=create_model_dict(), renderer=renderer)

        # If a converter is provided, use its type but update the renderer
        if isinstance(converter, TableConverter):
            return TableConverter(artifact_dict=create_model_dict(), renderer=renderer)
        elif isinstance(converter, PdfConverter):
            return PdfConverter(artifact_dict=create_model_dict(), renderer=renderer)
        else:
            raise ValueError(f"Invalid converter type: {type(converter)}")

    def _get_renderer(self, export_format: str) -> str:
        """Get the renderer string based on the export format.

        :param export_format: Format to export the document in.
        :return: The renderer string.
        :raises ValueError: If the export format is invalid.
        """
        if export_format == "markdown":
            return "marker.renderers.markdown.MarkdownRenderer"
        elif export_format == "html":
            return "marker.renderers.html.HTMLRenderer"
        elif export_format == "json":
            return "marker.renderers.json.JSONRenderer"
        else:
            raise ValueError(f"Invalid export format: {export_format}")

    def _call(self, docs: Iterable[Doc]) -> Iterable[Doc]:
        """Process documents using Marker.

        :param docs: Documents to process.
        :return: Processed documents.
        """
        docs = list(docs)

        for doc in docs:
            # Convert URI to string if it's a Path
            uri = str(doc.uri) if isinstance(doc.uri, Path) else doc.uri
            # Process the document
            rendered = self._converter(uri)

            # Extract text and optionally images
            text, _, images = text_from_rendered(rendered)
            if self._extract_images:
                doc.images = images

            # Update document text
            doc.text = text

        for doc in docs:
            yield doc

    @property
    def _state(self) -> dict[str, Any]:
        """Get state for serialization.

        :return: State dictionary.
        """
        return {
            **super()._state,
            "converter": self._converter,
            "export_format": self._export_format,
            "extract_images": self._extract_images,
        }
