"""
Pydantic models for Docling Serve API v1.14.0.

This module provides fully-typed Pydantic v2 models that mirror every schema
in the Docling Serve OpenAPI specification.  Models use ``extra="allow"`` on
response types so the SDK is forward-compatible with future server versions,
while request types use ``extra="forbid"`` to catch typos at construction time.

Security note
-------------
* ``S3SourceRequest``, ``S3Target``: credential fields (``access_key``,
  ``secret_key``) are typed as ``SecretStr`` and excluded from ``repr``
  so they are never accidentally logged or printed.
"""

from __future__ import annotations

from enum import Enum
from typing import Any, Dict, List, Literal, Optional, Tuple, Union

from pydantic import BaseModel, ConfigDict, Field, SecretStr


# ---------------------------------------------------------------------------
# Enums  (alphabetical, values match OpenAPI spec exactly)
# ---------------------------------------------------------------------------

class ConversionStatus(str, Enum):
    """Document conversion status."""
    PENDING = "pending"
    STARTED = "started"
    SUCCESS = "success"
    FAILURE = "failure"
    PARTIAL_SUCCESS = "partial_success"
    SKIPPED = "skipped"


class ImageRefMode(str, Enum):
    """Image reference mode for export."""
    PLACEHOLDER = "placeholder"
    EMBEDDED = "embedded"
    REFERENCED = "referenced"


class InferenceFramework(str, Enum):
    """VLM inference framework."""
    MLX = "mlx"
    TRANSFORMERS = "transformers"
    VLLM = "vllm"


class InputFormat(str, Enum):
    """Input document formats supported by Docling Serve."""
    DOCX = "docx"
    PPTX = "pptx"
    HTML = "html"
    IMAGE = "image"
    PDF = "pdf"
    ASCIIDOC = "asciidoc"
    MD = "md"
    CSV = "csv"
    XLSX = "xlsx"
    XML_USPTO = "xml_uspto"
    XML_JATS = "xml_jats"
    XML_XBRL = "xml_xbrl"
    METS_GBS = "mets_gbs"
    JSON_DOCLING = "json_docling"
    AUDIO = "audio"
    VTT = "vtt"
    LATEX = "latex"


class OCREngine(str, Enum):
    """OCR engine selection."""
    AUTO = "auto"
    EASYOCR = "easyocr"
    OCRMAC = "ocrmac"
    RAPIDOCR = "rapidocr"
    TESSEROCR = "tesserocr"
    TESSERACT = "tesseract"


class OutputFormat(str, Enum):
    """Output document formats."""
    MD = "md"
    JSON = "json"
    YAML = "yaml"
    HTML = "html"
    HTML_SPLIT_PAGE = "html_split_page"
    TEXT = "text"
    DOCTAGS = "doctags"


class PdfBackend(str, Enum):
    """PDF parsing backends."""
    PYPDFIUM2 = "pypdfium2"
    DOCLING_PARSE = "docling_parse"
    DLPARSE_V1 = "dlparse_v1"
    DLPARSE_V2 = "dlparse_v2"
    DLPARSE_V4 = "dlparse_v4"


class ProcessingPipeline(str, Enum):
    """Document processing pipeline type."""
    LEGACY = "legacy"
    STANDARD = "standard"
    VLM = "vlm"
    ASR = "asr"


class ResponseFormat(str, Enum):
    """VLM response format."""
    DOCTAGS = "doctags"
    MARKDOWN = "markdown"
    DEEPSEEKOCR_MARKDOWN = "deepseekocr_markdown"
    HTML = "html"
    OTSL = "otsl"
    PLAINTEXT = "plaintext"


class TableFormerMode(str, Enum):
    """Table structure extraction mode."""
    FAST = "fast"
    ACCURATE = "accurate"


class TargetName(str, Enum):
    """Target name for multipart/form-data endpoints."""
    INBODY = "inbody"
    ZIP = "zip"


class TaskType(str, Enum):
    """Async task type."""
    CONVERT = "convert"
    CHUNK = "chunk"


class TransformersModelType(str, Enum):
    """Transformers auto-model type."""
    AUTOMODEL = "automodel"
    AUTOMODEL_VISION2SEQ = "automodel-vision2seq"
    AUTOMODEL_CAUSALLM = "automodel-causallm"
    AUTOMODEL_IMAGETEXTTOTEXT = "automodel-imagetexttotext"


class VlmEngineType(str, Enum):
    """VLM inference engine type."""
    TRANSFORMERS = "transformers"
    MLX = "mlx"
    VLLM = "vllm"
    API = "api"
    API_OLLAMA = "api_ollama"
    API_LMSTUDIO = "api_lmstudio"
    API_OPENAI = "api_openai"
    AUTO_INLINE = "auto_inline"


class VlmModelType(str, Enum):
    """Pre-defined VLM model presets."""
    SMOLDOCLING = "smoldocling"
    SMOLDOCLING_VLLM = "smoldocling_vllm"
    GRANITE_VISION = "granite_vision"
    GRANITE_VISION_VLLM = "granite_vision_vllm"
    GRANITE_VISION_OLLAMA = "granite_vision_ollama"
    GOT_OCR_2 = "got_ocr_2"
    GRANITE_DOCLING = "granite_docling"
    GRANITE_DOCLING_VLLM = "granite_docling_vllm"
    DEEPSEEKOCR_OLLAMA = "deepseekocr_ollama"


# ---------------------------------------------------------------------------
# Request models  (extra="forbid" – catch typos immediately)
# ---------------------------------------------------------------------------

class _StrictRequest(BaseModel):
    """Base for all *request* models — forbids unexpected fields."""
    model_config = ConfigDict(extra="forbid", populate_by_name=True)


class _FlexibleResponse(BaseModel):
    """Base for all *response* models — allows extra fields from future API versions."""
    model_config = ConfigDict(extra="allow", populate_by_name=True)


# -- Source request models ---------------------------------------------------

class FileSourceRequest(_StrictRequest):
    """File source: a base64-encoded document."""
    base64_string: str = Field(..., description="Content of the file serialised in base64.")
    filename: str = Field(..., description="Filename of the uploaded document.")
    kind: Literal["file"] = Field(default="file")


class HttpSourceRequest(_StrictRequest):
    """HTTP source: a URL pointing to a document."""
    url: str = Field(..., description="HTTP url to process.")
    headers: Dict[str, Any] = Field(default_factory=dict, description="Additional headers.")
    kind: Literal["http"] = Field(default="http")


class S3SourceRequest(_StrictRequest):
    """S3 source: credentials + bucket/key locator.

    ``access_key`` and ``secret_key`` are ``SecretStr`` so they never leak
    into logs or ``repr()`` output.
    """
    endpoint: str = Field(..., description="S3 service endpoint, without protocol.")
    access_key: SecretStr = Field(..., description="S3 access key.")
    secret_key: SecretStr = Field(..., description="S3 secret key.")
    bucket: str = Field(..., description="S3 bucket name.")
    key_prefix: str = Field(default="", description="Prefix for object keys.")
    verify_ssl: bool = Field(default=True, description="Verify SSL certificates.")
    kind: Literal["s3"] = Field(default="s3")

    def model_dump(self, **kwargs: Any) -> Dict[str, Any]:  # noqa: D102
        """Override to expose secrets as plain strings for API serialisation."""
        data = super().model_dump(**kwargs)
        if "access_key" in data and isinstance(data["access_key"], SecretStr):
            data["access_key"] = data["access_key"].get_secret_value()
        if "secret_key" in data and isinstance(data["secret_key"], SecretStr):
            data["secret_key"] = data["secret_key"].get_secret_value()
        return data


# Union alias
SourceRequest = Union[FileSourceRequest, HttpSourceRequest, S3SourceRequest]


# -- Target models -----------------------------------------------------------

class InBodyTarget(_StrictRequest):
    """Return converted content inline in the JSON response body."""
    kind: Literal["inbody"] = Field(default="inbody")


class ZipTarget(_StrictRequest):
    """Return converted content as a ZIP archive."""
    kind: Literal["zip"] = Field(default="zip")


class S3Target(_StrictRequest):
    """Upload converted content to an S3 bucket."""
    endpoint: str = Field(..., description="S3 service endpoint, without protocol.")
    access_key: SecretStr = Field(..., description="S3 access key.")
    secret_key: SecretStr = Field(..., description="S3 secret key.")
    bucket: str = Field(..., description="S3 bucket name.")
    key_prefix: str = Field(default="", description="Prefix for object keys.")
    verify_ssl: bool = Field(default=True, description="Verify SSL.")
    kind: Literal["s3"] = Field(default="s3")

    def model_dump(self, **kwargs: Any) -> Dict[str, Any]:  # noqa: D102
        data = super().model_dump(**kwargs)
        if "access_key" in data and isinstance(data["access_key"], SecretStr):
            data["access_key"] = data["access_key"].get_secret_value()
        if "secret_key" in data and isinstance(data["secret_key"], SecretStr):
            data["secret_key"] = data["secret_key"].get_secret_value()
        return data


class PutTarget(_StrictRequest):
    """PUT converted content to a presigned URL."""
    url: str = Field(..., description="Presigned URL to PUT the result to.")
    kind: Literal["put"] = Field(default="put")


# Union alias
TargetRequest = Union[InBodyTarget, ZipTarget, S3Target, PutTarget]


# -- Picture description models (deprecated, still in OpenAPI spec) -----------

class PictureDescriptionLocal(_StrictRequest):
    """DEPRECATED local VLM config for picture description.

    Matches the ``PictureDescriptionLocal`` schema in the OpenAPI spec.
    Used as a structured object in JSON-body endpoints; the SDK
    automatically serialises it to a JSON string for form-data endpoints.
    """
    repo_id: str = Field(..., description="Hugging Face repository ID.")
    prompt: str = Field(
        default="Describe this image in a few sentences.",
        description="Prompt for the vision-language model.",
    )
    generation_config: Dict[str, Any] = Field(
        default_factory=lambda: {"max_new_tokens": 200, "do_sample": False},
        description="HuggingFace GenerationConfig overrides.",
    )


class PictureDescriptionApi(_StrictRequest):
    """DEPRECATED remote API config for picture description.

    Matches the ``PictureDescriptionApi`` schema in the OpenAPI spec.
    """
    url: str = Field(..., description="OpenAI-compatible API endpoint URL.")
    headers: Dict[str, str] = Field(default_factory=dict, description="Extra HTTP headers.")
    params: Dict[str, Any] = Field(default_factory=dict, description="Model parameters.")
    timeout: float = Field(default=20.0, description="API request timeout in seconds.")
    concurrency: int = Field(default=1, gt=0, description="Max concurrent API requests.")
    prompt: str = Field(
        default="Describe this image in a few sentences.",
        description="Prompt for the vision-language model.",
    )


# -- Conversion options -------------------------------------------------------

class ConvertDocumentsRequestOptions(_StrictRequest):
    """Full set of conversion options matching the OpenAPI spec."""

    from_formats: List[InputFormat] = Field(
        default=[
            InputFormat.DOCX, InputFormat.PPTX, InputFormat.HTML,
            InputFormat.IMAGE, InputFormat.PDF, InputFormat.ASCIIDOC,
            InputFormat.MD, InputFormat.CSV, InputFormat.XLSX,
            InputFormat.XML_USPTO, InputFormat.XML_JATS, InputFormat.XML_XBRL,
            InputFormat.METS_GBS, InputFormat.JSON_DOCLING,
            InputFormat.AUDIO, InputFormat.VTT, InputFormat.LATEX,
        ],
        description="Input format(s) to convert from.",
    )
    to_formats: List[OutputFormat] = Field(
        default=[OutputFormat.MD],
        description="Output format(s) to convert to.",
    )
    image_export_mode: ImageRefMode = Field(
        default=ImageRefMode.EMBEDDED,
        description="Image export mode (placeholder, embedded, referenced).",
    )
    do_ocr: bool = Field(default=True, description="Enable OCR processing.")
    force_ocr: bool = Field(default=False, description="Force OCR over existing text.")
    ocr_engine: OCREngine = Field(default=OCREngine.EASYOCR, description="OCR engine to use.")
    ocr_lang: Optional[List[str]] = Field(
        default=None,
        description="Languages for OCR engine (engine-specific names).",
    )
    pdf_backend: PdfBackend = Field(
        default=PdfBackend.DOCLING_PARSE,
        description="PDF parsing backend.",
    )
    table_mode: TableFormerMode = Field(
        default=TableFormerMode.ACCURATE,
        description="Table structure extraction mode.",
    )
    table_cell_matching: bool = Field(
        default=True,
        description="Match table cell predictions back to PDF cells.",
    )
    pipeline: ProcessingPipeline = Field(
        default=ProcessingPipeline.STANDARD,
        description="Processing pipeline type.",
    )
    page_range: Tuple[int, int] = Field(
        default=(1, 9_223_372_036_854_775_807),
        description="Page range to process (1-based, inclusive).",
    )
    document_timeout: float = Field(
        default=604800.0,
        gt=0,
        le=604800,
        description="Per-document timeout in seconds.",
    )
    abort_on_error: bool = Field(default=False, description="Abort batch on first error.")
    do_table_structure: bool = Field(default=True, description="Extract table structure.")
    include_images: bool = Field(default=True, description="Include images in output.")
    images_scale: float = Field(default=2.0, description="Image scale factor.")
    md_page_break_placeholder: str = Field(
        default="",
        description="Page break placeholder in markdown output.",
    )
    do_code_enrichment: bool = Field(default=False, description="Enable code enrichment via OCR.")
    do_formula_enrichment: bool = Field(default=False, description="Enable formula OCR (LaTeX).")
    do_picture_classification: bool = Field(default=False, description="Classify pictures.")
    do_chart_extraction: bool = Field(default=False, description="Extract chart data.")
    do_picture_description: bool = Field(default=False, description="Describe pictures via VLM.")
    picture_description_area_threshold: float = Field(
        default=0.05,
        description="Minimum picture area ratio to trigger description.",
    )

    # -- VLM pipeline presets & custom configs (new in v1.14) --
    vlm_pipeline_preset: Optional[str] = Field(
        default=None,
        description='Preset ID for VLM pipeline (e.g. "default", "granite_docling").',
    )
    picture_description_preset: Optional[str] = Field(
        default=None,
        description='Preset ID for picture description (e.g. "smolvlm").',
    )
    code_formula_preset: Optional[str] = Field(
        default=None,
        description='Preset ID for code/formula extraction.',
    )
    vlm_pipeline_custom_config: Optional[Dict[str, Any]] = Field(
        default=None,
        description="Custom VLM pipeline config (model_spec + engine_options).",
    )
    picture_description_custom_config: Optional[Dict[str, Any]] = Field(
        default=None,
        description="Custom picture description config.",
    )
    code_formula_custom_config: Optional[Dict[str, Any]] = Field(
        default=None,
        description="Custom code/formula extraction config.",
    )
    table_structure_custom_config: Optional[Dict[str, Any]] = Field(
        default=None,
        description="Custom table structure model config dict.",
    )
    layout_custom_config: Optional[Dict[str, Any]] = Field(
        default=None,
        description="Custom layout model config dict.",
    )

    # -- Deprecated VLM fields (kept for backward compat) --
    vlm_pipeline_model: Optional[VlmModelType] = Field(
        default=None,
        description="DEPRECATED: use vlm_pipeline_preset instead.",
    )
    picture_description_local: Optional[PictureDescriptionLocal] = Field(
        default=None,
        description="DEPRECATED: use picture_description_preset instead.",
    )
    picture_description_api: Optional[PictureDescriptionApi] = Field(
        default=None,
        description="DEPRECATED: use picture_description_preset instead.",
    )
    vlm_pipeline_model_local: Optional[str] = Field(
        default=None,
        description="DEPRECATED: use vlm_pipeline_preset instead.",
    )
    vlm_pipeline_model_api: Optional[str] = Field(
        default=None,
        description="DEPRECATED: use vlm_pipeline_preset instead.",
    )


# -- Convert request / response -----------------------------------------------

class ConvertDocumentsRequest(_StrictRequest):
    """Top-level conversion request (JSON body for ``/v1/convert/source``)."""
    sources: List[SourceRequest] = Field(..., description="Document sources.")
    options: ConvertDocumentsRequestOptions = Field(
        default_factory=ConvertDocumentsRequestOptions,
        description="Conversion options.",
    )
    target: TargetRequest = Field(
        default_factory=InBodyTarget,
        description="Output target.",
    )


class ExportDocumentResponse(_FlexibleResponse):
    """Single exported document content."""
    filename: str = Field(..., description="Source filename.")
    md_content: Optional[str] = Field(default=None, description="Markdown content.")
    json_content: Optional[Dict[str, Any]] = Field(default=None, description="JSON content.")
    html_content: Optional[str] = Field(default=None, description="HTML content.")
    text_content: Optional[str] = Field(default=None, description="Plain-text content.")
    doctags_content: Optional[str] = Field(default=None, description="DocTags content.")


class ErrorItem(_FlexibleResponse):
    """Conversion error detail."""
    component_type: str = Field(..., description="Component that raised the error.")
    module_name: str = Field(..., description="Module name.")
    error_message: str = Field(..., description="Error message.")


class ProfilingItem(_FlexibleResponse):
    """Profiling information for a conversion step."""
    scope: str = Field(..., description="Profiling scope (page or document).")
    count: int = Field(default=0)
    times: List[float] = Field(default_factory=lambda: list[float]())
    start_timestamps: List[str] = Field(default_factory=list)


class ExportResult(_FlexibleResponse):
    """Container of all exported content for a single document."""
    kind: Literal["ExportResult"] = Field(default="ExportResult")
    content: ExportDocumentResponse
    status: ConversionStatus
    errors: List[ErrorItem] = Field(default_factory=lambda: list[ErrorItem]())
    timings: Dict[str, ProfilingItem] = Field(default_factory=dict)


class ConvertDocumentResponse(_FlexibleResponse):
    """Response from ``/v1/convert/source`` and ``/v1/convert/file``."""
    document: ExportDocumentResponse = Field(..., description="Converted document data.")
    status: ConversionStatus = Field(..., description="Conversion status.")
    processing_time: float = Field(..., description="Processing time (s).")
    errors: List[ErrorItem] = Field(default_factory=lambda: list[ErrorItem](), description="Processing errors.")
    timings: Dict[str, ProfilingItem] = Field(default_factory=dict, description="Processing timings.")


class PresignedUrlConvertDocumentResponse(_FlexibleResponse):
    """Response when target is S3 or PUT (no inline document)."""
    processing_time: float = Field(..., description="Processing time (s).")
    num_converted: int = Field(..., description="Number of documents converted.")
    num_succeeded: int = Field(..., description="Number that succeeded.")
    num_failed: int = Field(..., description="Number that failed.")


# -- Health / Version --------------------------------------------------------

class HealthCheckResponse(_FlexibleResponse):
    """``GET /health`` response."""
    status: str = Field(default="ok", description="Health status.")


# -- Task status / result ----------------------------------------------------

class TaskProcessingMeta(_FlexibleResponse):
    """Progress metadata for an async task."""
    num_docs: int = Field(..., description="Total number of documents.")
    num_processed: int = Field(default=0, description="Documents processed so far.")
    num_succeeded: int = Field(default=0, description="Successful conversions.")
    num_failed: int = Field(default=0, description="Failed conversions.")


class TaskStatusResponse(_FlexibleResponse):
    """``GET /v1/status/poll/{task_id}`` response."""
    task_id: str = Field(..., description="Unique task identifier.")
    task_type: str = Field(..., description="Task type (convert or chunk).")
    task_status: str = Field(..., description="Current task status.")
    task_position: Optional[int] = Field(default=None, description="Queue position.")
    task_meta: Optional[TaskProcessingMeta] = Field(default=None, description="Processing meta.")
    error_message: Optional[str] = Field(default=None, description="Error message if failed.")


# -- Clear endpoints ---------------------------------------------------------

class ClearResponse(_FlexibleResponse):
    """``GET /v1/clear/converters`` and ``/v1/clear/results`` response."""
    status: str = Field(default="ok", description="Operation status.")


# -- Chunker options ---------------------------------------------------------

class HierarchicalChunkerOptions(_StrictRequest):
    """Configuration for the HierarchicalChunker."""
    chunker: Literal["hierarchical"] = Field(default="hierarchical")
    use_markdown_tables: bool = Field(
        default=False,
        description="Use markdown table format instead of triplets.",
    )
    include_raw_text: bool = Field(
        default=False,
        description="Include raw_text alongside contextualised text.",
    )


class HybridChunkerOptions(_StrictRequest):
    """Configuration for the HybridChunker."""
    chunker: Literal["hybrid"] = Field(default="hybrid")
    use_markdown_tables: bool = Field(
        default=False,
        description="Use markdown table format instead of triplets.",
    )
    include_raw_text: bool = Field(
        default=False,
        description="Include raw_text alongside contextualised text.",
    )
    max_tokens: Optional[int] = Field(
        default=None,
        description="Maximum tokens per chunk (None = auto from tokenizer).",
    )
    tokenizer: str = Field(
        default="sentence-transformers/all-MiniLM-L6-v2",
        description="HuggingFace tokenizer model name.",
    )
    merge_peers: bool = Field(
        default=True,
        description="Merge undersized successive chunks with same headings.",
    )


# -- Chunking request models (JSON body for /v1/chunk/*/source) ---------------

class HierarchicalChunkerOptionsDocumentsRequest(_StrictRequest):
    """Request body for ``/v1/chunk/hierarchical/source``."""
    convert_options: ConvertDocumentsRequestOptions = Field(
        default_factory=ConvertDocumentsRequestOptions,
        description="Conversion options.",
    )
    sources: List[SourceRequest] = Field(..., description="Document sources.")
    include_converted_doc: bool = Field(
        default=False,
        description="Include converted document alongside chunks.",
    )
    target: TargetRequest = Field(
        default_factory=InBodyTarget,
        description="Output target.",
    )
    chunking_options: Optional[HierarchicalChunkerOptions] = Field(
        default=None,
        description="Chunker-specific options.",
    )


class HybridChunkerOptionsDocumentsRequest(_StrictRequest):
    """Request body for ``/v1/chunk/hybrid/source``."""
    convert_options: ConvertDocumentsRequestOptions = Field(
        default_factory=ConvertDocumentsRequestOptions,
        description="Conversion options.",
    )
    sources: List[SourceRequest] = Field(..., description="Document sources.")
    include_converted_doc: bool = Field(
        default=False,
        description="Include converted document alongside chunks.",
    )
    target: TargetRequest = Field(
        default_factory=InBodyTarget,
        description="Output target.",
    )
    chunking_options: Optional[HybridChunkerOptions] = Field(
        default=None,
        description="Chunker-specific options.",
    )


# -- Chunk response -----------------------------------------------------------

class ChunkedDocumentResultItem(_FlexibleResponse):
    """A single chunk of a document with its metadata and content.

    Matches the ``ChunkedDocumentResultItem`` schema in the OpenAPI spec.
    """
    filename: str = Field(..., description="Source filename.")
    chunk_index: int = Field(..., description="Zero-based chunk index.")
    text: str = Field(..., description="Chunk text with structural context.")
    doc_items: List[str] = Field(..., description="Doc item references.")
    raw_text: Optional[str] = Field(default=None, description="Raw text without formatting.")
    num_tokens: Optional[int] = Field(default=None, description="Token count (if available).")
    headings: Optional[List[str]] = Field(default=None, description="Section headings.")
    captions: Optional[List[str]] = Field(default=None, description="Table/figure captions.")
    page_numbers: Optional[List[int]] = Field(default=None, description="Source page numbers.")
    metadata: Optional[Dict[str, Any]] = Field(default=None, description="Additional metadata.")


class ChunkDocumentResponse(_FlexibleResponse):
    """Response from chunking endpoints."""
    chunks: List[ChunkedDocumentResultItem] = Field(..., description="List of chunks.")
    documents: List[ExportResult] = Field(..., description="Converted documents.")
    processing_time: float = Field(..., description="Processing time (s).")


# -- HTTP validation error (mirrors FastAPI) -----------------------------------

class ValidationErrorDetail(_FlexibleResponse):
    """Single validation error from the server."""
    loc: List[Union[str, int]] = Field(default_factory=lambda: list[Union[str, int]](), description="Error location.")
    msg: str = Field(default="", description="Error message.")
    type: str = Field(default="", description="Error type.")


class HTTPValidationError(_FlexibleResponse):
    """422 response body returned by FastAPI when validation fails."""
    detail: List[ValidationErrorDetail] = Field(default_factory=lambda: list[ValidationErrorDetail]())
