# Define the input options for the API
from typing import Annotated, Any, Optional

from pydantic import AnyUrl, BaseModel, Field, PositiveInt, model_validator
from typing_extensions import Self

from docling.datamodel import vlm_model_specs
from docling.datamodel.base_models import InputFormat, OutputFormat
from docling.datamodel.pipeline_options import (
    PdfBackend,
    PictureDescriptionBaseOptions,
    ProcessingPipeline,
    TableFormerMode,
    TableStructureOptions,
)
from docling.datamodel.pipeline_options_vlm_model import (
    InferenceFramework,
    InlineVlmOptions,
    ResponseFormat,
    TransformersModelType,
)
from docling.datamodel.settings import (
    DEFAULT_PAGE_RANGE,
    PageRange,
)
from docling_core.types.doc import ImageRefMode


class PictureDescriptionLocal(BaseModel):
    repo_id: Annotated[
        str,
        Field(
            description="Repository id from the Hugging Face Hub.",
            examples=[
                "HuggingFaceTB/SmolVLM-256M-Instruct",
                "ibm-granite/granite-vision-3.3-2b",
            ],
        ),
    ]
    prompt: Annotated[
        str,
        Field(
            description="Prompt used when calling the vision-language model.",
            examples=[
                "Describe this image in a few sentences.",
                "This is a figure from a document. Provide a detailed description of it.",
            ],
        ),
    ] = "Describe this image in a few sentences."
    generation_config: Annotated[
        dict[str, Any],
        Field(
            description="Config from https://huggingface.co/docs/transformers/en/main_classes/text_generation#transformers.GenerationConfig",
            examples=[{"max_new_tokens": 200, "do_sample": False}],
        ),
    ] = {"max_new_tokens": 200, "do_sample": False}


class PictureDescriptionApi(BaseModel):
    url: Annotated[
        AnyUrl,
        Field(
            description="Endpoint which accepts openai-api compatible requests.",
            examples=[
                AnyUrl(
                    "http://localhost:8000/v1/chat/completions"
                ),  # example of a local vllm api
                AnyUrl(
                    "http://localhost:1234/v1/chat/completions"
                ),  # example of lm studio
                AnyUrl(
                    "http://localhost:11434/v1/chat/completions"
                ),  # example of ollama
            ],
        ),
    ]
    headers: Annotated[
        dict[str, str],
        Field(
            description="Headers used for calling the API endpoint. For example, it could include authentication headers."
        ),
    ] = {}
    params: Annotated[
        dict[str, Any],
        Field(
            description="Model parameters.",
            examples=[
                {  # on vllm
                    "model": "HuggingFaceTB/SmolVLM-256M-Instruct",
                    "max_completion_tokens": 200,
                },
                {  # on vllm
                    "model": "ibm-granite/granite-vision-3.3-2b",
                    "max_completion_tokens": 200,
                },
                {  # on ollama
                    "model": "granite3.2-vision:2b"
                },
            ],
        ),
    ] = {}
    timeout: Annotated[float, Field(description="Timeout for the API request.")] = 20
    concurrency: Annotated[
        PositiveInt,
        Field(
            description="Maximum number of concurrent requests to the API.",
            examples=[1],
        ),
    ] = 1
    prompt: Annotated[
        str,
        Field(
            description="Prompt used when calling the vision-language model.",
            examples=[
                "Describe this image in a few sentences.",
                "This is a figures from a document. Provide a detailed description of it.",
            ],
        ),
    ] = "Describe this image in a few sentences."


class VlmModelLocal(BaseModel):
    repo_id: Annotated[
        str,
        Field(
            description="Repository id from the Hugging Face Hub.",
        ),
    ]
    prompt: Annotated[
        str,
        Field(
            description="Prompt used when calling the vision-language model.",
            examples=[
                "Convert this page to docling.",
                "Convert this page to markdown. Do not miss any text and only output the bare markdown!",
            ],
        ),
    ] = "Convert this page to docling."
    scale: Annotated[float, Field(description="Scale factor of the images used.")] = 2.0
    response_format: Annotated[
        ResponseFormat, Field(description="Type of response generated by the model.")
    ]
    inference_framework: Annotated[
        InferenceFramework, Field(description="Inference framework to use.")
    ]
    transformers_model_type: Annotated[
        TransformersModelType,
        Field(description="Type of transformers auto-model to use."),
    ] = TransformersModelType.AUTOMODEL
    extra_generation_config: Annotated[
        dict[str, Any],
        Field(
            description="Config from https://huggingface.co/docs/transformers/en/main_classes/text_generation#transformers.GenerationConfig",
            examples=[{"max_new_tokens": 800, "do_sample": False}],
        ),
    ] = {"max_new_tokens": 800, "do_sample": False}
    temperature: Annotated[
        float,
        Field(
            description="Temperature parameter controlling the reproducibility of the result.",
            examples=[0.0, 1.0],
        ),
    ] = 0.0

    @staticmethod
    def from_docling(options: InlineVlmOptions):
        return VlmModelLocal.model_validate(options.model_dump())


class VlmModelApi(BaseModel):
    url: Annotated[
        AnyUrl,
        Field(
            description="Endpoint which accepts openai-api compatible requests.",
            examples=[
                AnyUrl(
                    "http://localhost:8000/v1/chat/completions"
                ),  # example of a local vllm api
                AnyUrl(
                    "http://localhost:1234/v1/chat/completions"
                ),  # example of lm studio
            ],
        ),
    ]
    headers: Annotated[
        dict[str, str],
        Field(
            description="Headers used for calling the API endpoint. For example, it could include authentication headers."
        ),
    ] = {}
    params: Annotated[
        dict[str, Any],
        Field(
            description="Model parameters.",
            examples=[
                {  # on vllm
                    "model": "ibm-granite/granite-docling-258M",
                    "max_completion_tokens": 800,
                },
                {  # on vllm
                    "model": "ibm-granite/granite-vision-3.3-2b",
                    "max_completion_tokens": 800,
                },
            ],
        ),
    ] = {}
    timeout: Annotated[float, Field(description="Timeout for the API request.")] = 60
    concurrency: Annotated[
        PositiveInt,
        Field(
            description="Maximum number of concurrent requests to the API.",
            examples=[1],
        ),
    ] = 1
    prompt: Annotated[
        str,
        Field(
            description="Prompt used when calling the vision-language model.",
            examples=[
                "Convert this page to docling.",
                "Convert this page to markdown. Do not miss any text and only output the bare markdown!",
            ],
        ),
    ] = "Convert this page to docling."
    scale: Annotated[float, Field(description="Scale factor of the images used.")] = 2.0
    response_format: Annotated[
        ResponseFormat, Field(description="Type of response generated by the model.")
    ]
    temperature: Annotated[
        float,
        Field(
            description="Temperature parameter controlling the reproducibility of the result.",
            examples=[0.0, 1.0],
        ),
    ] = 0.0


class ConvertDocumentsOptions(BaseModel):
    from_formats: Annotated[
        list[InputFormat],
        Field(
            description=(
                "Input format(s) to convert from. String or list of strings. "
                f"Allowed values: {', '.join([v.value for v in InputFormat])}. "
                "Optional, defaults to all formats."
            ),
            examples=[[v.value for v in InputFormat]],
        ),
    ] = list(InputFormat)

    to_formats: Annotated[
        list[OutputFormat],
        Field(
            description=(
                "Output format(s) to convert to. String or list of strings. "
                f"Allowed values: {', '.join([v.value for v in OutputFormat])}. "
                "Optional, defaults to Markdown."
            ),
            examples=[
                [OutputFormat.MARKDOWN],
                [OutputFormat.MARKDOWN, OutputFormat.JSON],
                [v.value for v in OutputFormat],
            ],
        ),
    ] = [OutputFormat.MARKDOWN]

    image_export_mode: Annotated[
        ImageRefMode,
        Field(
            description=(
                "Image export mode for the document (in case of JSON,"
                " Markdown or HTML). "
                f"Allowed values: {', '.join([v.value for v in ImageRefMode])}. "
                "Optional, defaults to Embedded."
            ),
            examples=[ImageRefMode.EMBEDDED.value],
            # pattern="embedded|placeholder|referenced",
        ),
    ] = ImageRefMode.EMBEDDED

    do_ocr: Annotated[
        bool,
        Field(
            description=(
                "If enabled, the bitmap content will be processed using OCR. "
                "Boolean. Optional, defaults to true"
            ),
            # examples=[True],
        ),
    ] = True

    force_ocr: Annotated[
        bool,
        Field(
            description=(
                "If enabled, replace existing text with OCR-generated "
                "text over content. Boolean. Optional, defaults to false."
            ),
            # examples=[False],
        ),
    ] = False

    ocr_engine: Annotated[
        str,
        Field(
            description=("The OCR engine to use. String. "),
        ),
    ] = "auto"

    ocr_lang: Annotated[
        Optional[list[str]],
        Field(
            description=(
                "List of languages used by the OCR engine. "
                "Note that each OCR engine has "
                "different values for the language names. String or list of strings. "
                "Optional, defaults to empty."
            ),
            examples=[["fr", "de", "es", "en"]],
        ),
    ] = None

    pdf_backend: Annotated[
        PdfBackend,
        Field(
            description=(
                "The PDF backend to use. String. "
                f"Allowed values: {', '.join([v.value for v in PdfBackend])}. "
                f"Optional, defaults to {PdfBackend.DLPARSE_V4.value}."
            ),
            examples=[PdfBackend.DLPARSE_V4],
        ),
    ] = PdfBackend.DLPARSE_V4

    table_mode: Annotated[
        TableFormerMode,
        Field(
            description=(
                "Mode to use for table structure, String. "
                f"Allowed values: {', '.join([v.value for v in TableFormerMode])}. "
                "Optional, defaults to accurate."
            ),
            examples=[TableStructureOptions().mode],
            # pattern="fast|accurate",
        ),
    ] = TableStructureOptions().mode

    table_cell_matching: Annotated[
        bool,
        Field(
            description="If true, matches table cells predictions back to PDF cells. Can break table output if PDF cells are merged across table columns. If false, let table structure model define the text cells, ignore PDF cells.",
            examples=[TableStructureOptions().do_cell_matching],
        ),
    ] = TableStructureOptions().do_cell_matching

    pipeline: Annotated[
        ProcessingPipeline,
        Field(description="Choose the pipeline to process PDF or image files."),
    ] = ProcessingPipeline.STANDARD

    page_range: Annotated[
        PageRange,
        Field(
            description="Only convert a range of pages. The page number starts at 1.",
            examples=[DEFAULT_PAGE_RANGE, (1, 4)],
        ),
    ] = DEFAULT_PAGE_RANGE

    document_timeout: Annotated[
        Optional[float],
        Field(
            description="The timeout for processing each document, in seconds.",
        ),
    ] = None

    abort_on_error: Annotated[
        bool,
        Field(
            description=(
                "Abort on error if enabled. Boolean. Optional, defaults to false."
            ),
            # examples=[False],
        ),
    ] = False

    do_table_structure: Annotated[
        bool,
        Field(
            description=(
                "If enabled, the table structure will be extracted. "
                "Boolean. Optional, defaults to true."
            ),
            examples=[True],
        ),
    ] = True

    include_images: Annotated[
        bool,
        Field(
            description=(
                "If enabled, images will be extracted from the document. "
                "Boolean. Optional, defaults to true."
            ),
            examples=[True],
        ),
    ] = True

    images_scale: Annotated[
        float,
        Field(
            description="Scale factor for images. Float. Optional, defaults to 2.0.",
            examples=[2.0],
        ),
    ] = 2.0

    md_page_break_placeholder: Annotated[
        str,
        Field(
            description="Add this placeholder between pages in the markdown output.",
            examples=["<!-- page-break -->", ""],
        ),
    ] = ""

    do_code_enrichment: Annotated[
        bool,
        Field(
            description=(
                "If enabled, perform OCR code enrichment. "
                "Boolean. Optional, defaults to false."
            ),
            examples=[False],
        ),
    ] = False

    do_formula_enrichment: Annotated[
        bool,
        Field(
            description=(
                "If enabled, perform formula OCR, return LaTeX code. "
                "Boolean. Optional, defaults to false."
            ),
            examples=[False],
        ),
    ] = False

    do_picture_classification: Annotated[
        bool,
        Field(
            description=(
                "If enabled, classify pictures in documents. "
                "Boolean. Optional, defaults to false."
            ),
            examples=[False],
        ),
    ] = False

    do_picture_description: Annotated[
        bool,
        Field(
            description=(
                "If enabled, describe pictures in documents. "
                "Boolean. Optional, defaults to false."
            ),
            examples=[False],
        ),
    ] = False

    picture_description_area_threshold: Annotated[
        float,
        Field(
            description="Minimum percentage of the area for a picture to be processed with the models.",
            examples=[PictureDescriptionBaseOptions().picture_area_threshold],
        ),
    ] = PictureDescriptionBaseOptions().picture_area_threshold

    picture_description_local: Annotated[
        Optional[PictureDescriptionLocal],
        Field(
            description="Options for running a local vision-language model in the picture description. The parameters refer to a model hosted on Hugging Face. This parameter is mutually exclusive with picture_description_api.",
            examples=[
                PictureDescriptionLocal(repo_id="ibm-granite/granite-vision-3.2-2b"),
                PictureDescriptionLocal(repo_id="HuggingFaceTB/SmolVLM-256M-Instruct"),
            ],
        ),
    ] = None

    picture_description_api: Annotated[
        Optional[PictureDescriptionApi],
        Field(
            description="API details for using a vision-language model in the picture description. This parameter is mutually exclusive with picture_description_local.",
            examples=[
                PictureDescriptionApi(
                    url="http://localhost:1234/v1/chat/completions",
                    params={"model": "granite3.2-vision:2b"},
                ),
                PictureDescriptionApi(
                    url="http://localhost:11434/v1/chat/completions",
                    params={"model": "granite3.2-vision:2b"},
                ),
            ],
        ),
    ] = None

    vlm_pipeline_model: Annotated[
        Optional[vlm_model_specs.VlmModelType],
        Field(
            description="Preset of local and API models for the vlm pipeline. This parameter is mutually exclusive with vlm_pipeline_model_local and vlm_pipeline_model_api. Use the other options for more parameters.",
            examples=[vlm_model_specs.VlmModelType.GRANITEDOCLING],
        ),
    ] = None
    vlm_pipeline_model_local: Annotated[
        Optional[VlmModelLocal],
        Field(
            description="Options for running a local vision-language model for the vlm pipeline. The parameters refer to a model hosted on Hugging Face. This parameter is mutually exclusive with vlm_pipeline_model_api and vlm_pipeline_model.",
            examples=[
                VlmModelLocal.from_docling(vlm_model_specs.GRANITEDOCLING_TRANSFORMERS),
                VlmModelLocal.from_docling(vlm_model_specs.GRANITEDOCLING_MLX),
                VlmModelLocal.from_docling(vlm_model_specs.GRANITE_VISION_TRANSFORMERS),
            ],
        ),
    ] = None

    vlm_pipeline_model_api: Annotated[
        Optional[VlmModelApi],
        Field(
            description="API details for using a vision-language model for the vlm pipeline. This parameter is mutually exclusive with vlm_pipeline_model_local and vlm_pipeline_model.",
            examples=[
                VlmModelApi(
                    url="http://localhost:1234/v1/chat/completions",
                    params={"model": "ibm-granite/granite-docling-258M-mlx"},
                    response_format=ResponseFormat.DOCTAGS,
                    prompt="Convert this page to docling.",
                )
            ],
        ),
    ] = None

    @model_validator(mode="after")
    def picture_description_exclusivity(self) -> Self:
        # Validate picture description options
        if (
            self.picture_description_local is not None
            and self.picture_description_api is not None
        ):
            raise ValueError(
                "The parameters picture_description_local and picture_description_api are mutually exclusive, only one of them can be set."
            )

        return self

    @model_validator(mode="after")
    def vlm_model_exclusivity(self) -> Self:
        # Validate vlm model options
        num_not_nan = sum(
            (
                opt is not None
                for opt in (
                    self.vlm_pipeline_model,
                    self.vlm_pipeline_model_local,
                    self.vlm_pipeline_model_api,
                )
            )
        )
        if num_not_nan > 1:
            raise ValueError(
                "The parameters vlm_pipeline_model, vlm_pipeline_model_local and vlm_pipeline_model_api are mutually exclusive, only one of them can be set."
            )

        return self
