# SPDX-FileCopyrightText: The Docling Contributors
# SPDX-License-Identifier: MIT

"""Options for the threaded layout+VLM pipeline."""

from typing import Union

from pydantic import Field, model_validator

from docling.datamodel.pipeline_options import (
    BaseLayoutOptions,
    LayoutObjectDetectionOptions,
    PaginatedPipelineOptions,
)
from docling.datamodel.pipeline_options_vlm_model import (
    ApiVlmOptions,
    InlineVlmOptions,
    ResponseFormat,
)
from docling.datamodel.vlm_model_specs import GRANITEDOCLING_2STAGE_TRANSFORMERS


class ThreadedLayoutVlmPipelineOptions(PaginatedPipelineOptions):
    """Pipeline options for the threaded layout+VLM pipeline."""

    images_scale: float = 2.0

    # VLM configuration (will be enhanced with layout awareness by the pipeline)
    vlm_options: Union[InlineVlmOptions, ApiVlmOptions] = (
        GRANITEDOCLING_2STAGE_TRANSFORMERS
    )

    # Layout model configuration
    layout_options: BaseLayoutOptions = Field(
        default_factory=lambda: LayoutObjectDetectionOptions(skip_cell_assignment=True),
    )

    # Threading and batching controls
    layout_batch_size: int = 4
    vlm_batch_size: int = 4
    batch_timeout_seconds: float = 2.0
    queue_max_size: int = 50

    @model_validator(mode="after")
    def validate_response_format(self):
        """Validate that VLM response format is DOCTAGS (required for this pipeline)."""
        if self.vlm_options.response_format != ResponseFormat.DOCTAGS:
            raise ValueError(
                f"ThreadedLayoutVlmPipeline only supports DOCTAGS response format, "
                f"but got {self.vlm_options.response_format}. "
                f"Please set vlm_options.response_format=ResponseFormat.DOCTAGS"
            )
        return self