# pyright: ignore
"""
Pipeline configurations for Docling converter.
"""

from docling.datamodel.base_models import InputFormat
from docling.datamodel.pipeline_options import PdfPipelineOptions, TableFormerMode, VlmPipelineOptions
from docling.datamodel.vlm_model_specs import GRANITEDOCLING_TRANSFORMERS
from docling.document_converter import DocumentConverter, PdfFormatOption
from docling.pipeline.vlm_pipeline import VlmPipeline


class PipelineConfig:
    """Configuration for Docling pipelines."""

    @staticmethod
    def get_regular_options() -> PdfPipelineOptions:
        """Get configuration for regular text extraction pipeline."""
        options = PdfPipelineOptions()
        options.do_table_structure = True
        options.table_structure_options.do_cell_matching = True
        options.table_structure_options.mode = TableFormerMode.ACCURATE
        options.do_ocr = True
        options.images_scale = 2.0
        options.generate_page_images = True
        return options

    @staticmethod
    def get_vlm_options() -> VlmPipelineOptions:
        """Get configuration for VLM pipeline with table extraction."""
        from docling.datamodel.accelerator_options import AcceleratorDevice

        vlm_options = VlmPipelineOptions(vlm_options=GRANITEDOCLING_TRANSFORMERS)
        vlm_options.generate_page_images = True  # Enable efficient page batching
        vlm_options.accelerator_options.device = AcceleratorDevice.CUDA
        return vlm_options


def init_converters():
    """
    Initialize converters.
    Called once per process for separate GPU contexts.
    """
    # Regular pipeline for table detection
    regular_options = PipelineConfig.get_regular_options()
    regular_format_options = PdfFormatOption(pipeline_options=regular_options)

    regular_converter = DocumentConverter(format_options={InputFormat.PDF: regular_format_options})

    # VLM pipeline for complex table extraction
    vlm_options = PipelineConfig.get_vlm_options()
    vlm_format_options = PdfFormatOption(pipeline_cls=VlmPipeline, pipeline_options=vlm_options)

    vlm_converter = DocumentConverter(format_options={InputFormat.PDF: vlm_format_options})

    return regular_converter, vlm_converter