from __future__ import annotations

from pathlib import Path

from docling.datamodel.accelerator_options import AcceleratorDevice, AcceleratorOptions
from docling.datamodel.base_models import InputFormat
from docling.datamodel.layout_model_specs import (
    DOCLING_LAYOUT_EGRET_LARGE,
    DOCLING_LAYOUT_EGRET_MEDIUM,
    DOCLING_LAYOUT_EGRET_XLARGE,
    DOCLING_LAYOUT_HERON,
    DOCLING_LAYOUT_HERON_101,
)
from docling.datamodel.pipeline_options import (
    LayoutOptions,
    PdfPipelineOptions,
    PictureClassificationLabel,
    PictureDescriptionApiOptions,
    RapidOcrOptions,
    TableFormerMode,
    TableStructureOptions,
)
from docling.document_converter import (
    DocumentConverter,
    PdfFormatOption,
    PowerpointFormatOption,
    WordFormatOption,
)

_FIGURE_PROMPT_DEFAULT = (
    "You are extracting figure content for a vector search index over Japanese "
    "and English business / regulatory documents.\n\n"
    "Respond in 2-4 sentences (max 120 words). State:\n"
    "1. Figure type (screenshot, flowchart, diagram, chart, table-as-image).\n"
    "2. Its purpose in the surrounding procedure.\n"
    "3. Any specific labels, button names, field names, values, dates, codes — "
    "quote verbatim.\n\n"
    "Do NOT describe colors, fonts, window chrome, OS or browser names, or "
    "spatial layout. If the figure is a flowchart, list steps in order using "
    "arrows (→). Reply in the dominant language of the visible text."
)

_LAYOUT_SPEC = {
    "heron": DOCLING_LAYOUT_HERON,
    "heron_101": DOCLING_LAYOUT_HERON_101,
    "egret_medium": DOCLING_LAYOUT_EGRET_MEDIUM,
    "egret_large": DOCLING_LAYOUT_EGRET_LARGE,
    "egret_xlarge": DOCLING_LAYOUT_EGRET_XLARGE,
}


def _build_rapid_ocr_options(settings) -> RapidOcrOptions:
    """RapidOCR (PaddleOCR-ONNX) is Docling's recommended Japanese-capable
    CPU OCR backend. When det/rec model paths are configured, we pin to
    pre-downloaded PP-OCRv5 weights for stable offline operation; otherwise
    we fall back to RapidOCR's bundled defaults.

    Note: RapidOcrOptions.bitmap_area_threshold defaults to 0.05 in Docling
    2.89; force_full_page_ocr is the right switch when the embedded text
    layer is known to be unreliable (e.g. OCR'd-from-paperwork PDFs).
    """
    kwargs: dict = {
        "lang": settings.ocr_languages,
        "force_full_page_ocr": settings.ocr_mode == "force",
        "bitmap_area_threshold": settings.ocr_bitmap_area_threshold,
    }
    if settings.rapidocr_det_model_path:
        kwargs["det_model_path"] = settings.rapidocr_det_model_path
    if settings.rapidocr_rec_model_path:
        kwargs["rec_model_path"] = settings.rapidocr_rec_model_path
    if settings.rapidocr_rec_keys_path:
        kwargs["rec_keys_path"] = settings.rapidocr_rec_keys_path
    if settings.rapidocr_cls_model_path:
        kwargs["cls_model_path"] = settings.rapidocr_cls_model_path
        kwargs["use_cls"] = True
    else:
        # RapidOCR's text-direction classifier corrects lines rotated 180°.
        # When no cls model path is set, RapidOCR otherwise tries to load
        # its bundled default (a PP-OCRv4 cls weight at
        # /opt/docling/models/RapidOcr/onnx/PP-OCRv4/cls/...) and fails hard
        # if that file is not present on disk. Disabling cls avoids a hard
        # build-time dependency on weights that contribute little to
        # right-side-up printed Japanese documents. Set
        # RAPIDOCR_CLS_MODEL_PATH to re-enable orientation classification.
        kwargs["use_cls"] = False
    return RapidOcrOptions(**kwargs)


_DEFAULT_FIGURE_CLASSES_ALLOW: tuple[PictureClassificationLabel, ...] = (
    PictureClassificationLabel.SCREENSHOT_FROM_COMPUTER,
    PictureClassificationLabel.SCREENSHOT_FROM_MANUAL,
    PictureClassificationLabel.FLOW_CHART,
    PictureClassificationLabel.BAR_CHART,
    PictureClassificationLabel.LINE_CHART,
    PictureClassificationLabel.PIE_CHART,
    PictureClassificationLabel.BOX_PLOT,
    PictureClassificationLabel.SCATTER_PLOT,
    PictureClassificationLabel.TABLE,
    PictureClassificationLabel.ENGINEERING_DRAWING,
    PictureClassificationLabel.GEOGRAPHICAL_MAP,
    PictureClassificationLabel.TOPOGRAPHICAL_MAP,
    PictureClassificationLabel.PHOTOGRAPH,
)


def _build_picture_description_options(settings) -> PictureDescriptionApiOptions | None:
    if not settings.do_picture_description:
        return None
    if not settings.openai_api_key:
        raise RuntimeError(
            "DO_PICTURE_DESCRIPTION=true but OPENAI_API_KEY is empty. "
            "Either set the key or disable picture description."
        )
    return PictureDescriptionApiOptions(
        url="https://api.openai.com/v1/chat/completions",
        headers={"Authorization": f"Bearer {settings.openai_api_key}"},
        params={"model": settings.openai_model},
        prompt=settings.figure_prompt.strip() or _FIGURE_PROMPT_DEFAULT,
        timeout=60,
        # Only send the VLM figures we actually want captions for. Icons,
        # logos, barcodes, page thumbnails, signatures and stamps are
        # filtered out — they produce noisy captions and burn budget.
        classification_allow=list(_DEFAULT_FIGURE_CLASSES_ALLOW),
        classification_min_confidence=settings.picture_classification_min_confidence,
        picture_area_threshold=settings.picture_area_threshold,
        concurrency=settings.figure_concurrency,
    )


def _build_pipeline_options(settings) -> PdfPipelineOptions:
    opts = PdfPipelineOptions()

    # ---- OCR (Japanese-tuned PP-OCRv5 via RapidOCR/ONNX, CPU-friendly) ----
    opts.do_ocr = settings.ocr_mode != "off"
    if opts.do_ocr:
        opts.ocr_options = _build_rapid_ocr_options(settings)

    # ---- Layout (HERON by default; HERON-101 for highest accuracy) ----
    opts.layout_options = LayoutOptions(model_spec=_LAYOUT_SPEC[settings.layout_model])

    # ---- Table structure (TableFormer V1 + ACCURATE + cell_matching tuned) ----
    # V1 (TableStructureOptions) supports FAST/ACCURATE modes; V2 does not.
    # V1 ACCURATE is also Docling's own default and benchmarks competitively
    # with V2 on banded/borderless tables when do_cell_matching=False.
    opts.do_table_structure = True
    opts.table_structure_options = TableStructureOptions(
        mode=(TableFormerMode.ACCURATE if settings.table_mode == "accurate" else TableFormerMode.FAST),
        do_cell_matching=settings.table_do_cell_matching,
    )

    # ---- Figure enrichment (GPT-4.1-mini captions, area-thresholded) ----
    pd_opts = _build_picture_description_options(settings)
    opts.generate_picture_images = True
    # Required for region-level VLM table rescue (table_rescue.py crops from
    # the rendered page bitmap). Cheap on CPU since images are only
    # rasterized when this flag is true.
    opts.generate_page_images = settings.enable_vlm_table_rescue
    opts.images_scale = 2.0
    # `compile_model=True` (the Docling default) asks PyTorch Inductor to
    # JIT-compile the classifier, which requires a host C compiler — gcc on
    # Linux, MSVC `cl` on Windows. Disabling it avoids a hard dependency on
    # the build toolchain at runtime and costs ~5% throughput on CPU.
    opts.picture_classification_options.engine_options.compile_model = False
    opts.do_picture_classification = True  # required for classification_allow
    opts.do_picture_description = pd_opts is not None
    if pd_opts is not None:
        opts.picture_description_options = pd_opts
        opts.enable_remote_services = True

    # ---- Long-running document guard ----
    # Bound conversion wall-time so one slow doc cannot pin a worker.
    if settings.document_timeout_s:
        opts.document_timeout = settings.document_timeout_s

    # ---- CPU batch tuning ----
    opts.ocr_batch_size = settings.ocr_batch_size
    opts.layout_batch_size = settings.layout_batch_size
    opts.table_batch_size = settings.table_batch_size

    # ---- Accelerator (match AWS vCPUs) ----
    if settings.num_threads:
        opts.accelerator_options = AcceleratorOptions(
            num_threads=settings.num_threads,
            device=AcceleratorDevice.CPU,
        )

    return opts


def build_converter(settings) -> tuple[DocumentConverter, str]:
    """Returns (converter, backend_label). The backend_label is surfaced in
    every ParseResponse so operators can see at a glance which configuration
    produced the document.
    """
    opts = _build_pipeline_options(settings)
    label = (
        f"docling[layout={settings.layout_model} "
        f"ocr={settings.ocr_mode} "
        f"table={settings.table_mode}/cell_match={settings.table_do_cell_matching} "
        f"figures={'gpt-4.1-mini' if settings.do_picture_description else 'off'}"
        f"{' rescue=on' if settings.enable_vlm_table_rescue else ''}]"
    )
    converter = DocumentConverter(
        format_options={
            InputFormat.PDF: PdfFormatOption(pipeline_options=opts),
            InputFormat.DOCX: WordFormatOption(pipeline_options=opts),
            InputFormat.PPTX: PowerpointFormatOption(pipeline_options=opts),
        }
    )
    return converter, label