from __future__ import annotations

import os
import logging
from dataclasses import dataclass
from typing import Any, Dict, List, Optional, Tuple


@dataclass
class LayoutBlock:
    block_id: str
    page: int
    order: int
    text: str
    block_type: str = "text"


_CONVERTER_CACHE: Dict[Tuple[Any, ...], Any] = {}


def _safe_env() -> None:
    os.environ["MKL_THREADING_LAYER"] = "GNU"
    os.environ["OMP_NUM_THREADS"] = "1"
    os.environ["MKL_SERVICE_FORCE_INTEL"] = "1"
    os.environ["KMP_DUPLICATE_LIB_OK"] = "TRUE"


def _configure_docling_logging() -> None:
    quiet = os.getenv("NAVEXA_DOCLING_QUIET", "1").strip().lower() in {"1", "true", "yes", "on"}
    if not quiet:
        return
    logging.getLogger("RapidOCR").setLevel(logging.WARNING)
    logging.getLogger("rapidocr").setLevel(logging.WARNING)


def _get_converter(
    *,
    do_ocr: bool,
    backend: str,
    force_full_page_ocr: bool,
    do_picture_description: bool,
    enable_remote_services: bool,
) -> Any:
    from docling.datamodel.base_models import InputFormat  # type: ignore
    from docling.datamodel.pipeline_options import PdfPipelineOptions, RapidOcrOptions  # type: ignore
    from docling.document_converter import DocumentConverter, PdfFormatOption  # type: ignore

    key = (do_ocr, backend, force_full_page_ocr, do_picture_description, enable_remote_services)
    if key in _CONVERTER_CACHE:
        return _CONVERTER_CACHE[key]

    opts = PdfPipelineOptions()
    opts.do_ocr = do_ocr
    if do_ocr:
        ocr_options = RapidOcrOptions(backend=backend)  # type: ignore[arg-type]
        if hasattr(ocr_options, "force_full_page_ocr"):
            ocr_options.force_full_page_ocr = force_full_page_ocr
        opts.ocr_options = ocr_options
    if hasattr(opts, "do_picture_description"):
        opts.do_picture_description = do_picture_description
    if hasattr(opts, "enable_remote_services"):
        opts.enable_remote_services = enable_remote_services

    converter = DocumentConverter(format_options={InputFormat.PDF: PdfFormatOption(pipeline_options=opts)})
    _CONVERTER_CACHE[key] = converter
    return converter


def _resolve_parser_model(parser_model: Optional[str]) -> str:
    resolved = (parser_model or os.getenv("NAVEXA_PARSER_MODEL", "docling")).strip().lower() or "docling"
    if resolved != "docling":
        raise ValueError(f"Unsupported parser_model '{resolved}'. Currently supported: docling")
    return resolved


def _resolve_output_format(output_format: Optional[str]) -> str:
    resolved = (output_format or os.getenv("NAVEXA_DOCLING_OUTPUT_FORMAT", "markdown")).strip().lower() or "markdown"
    if resolved not in {"markdown", "text"}:
        raise ValueError("output_format must be 'markdown' or 'text'")
    return resolved


def extract_docling_content(
    pdf_path: str,
    output_format: Optional[str] = None,
    parser_model: Optional[str] = None,
) -> Tuple[str, List[LayoutBlock], Dict[int, str], int]:
    _safe_env()
    _configure_docling_logging()

    from docling_core.types.doc import ImageRefMode  # type: ignore

    _resolve_parser_model(parser_model)
    resolved_output_format = _resolve_output_format(output_format)

    do_ocr = os.getenv("NAVEXA_DOCLING_OCR", "0").strip().lower() in {"1", "true", "yes", "on"}
    backend = os.getenv("NAVEXA_RAPIDOCR_BACKEND", "torch").strip().lower() or "torch"
    force_full_page_ocr = os.getenv("NAVEXA_DOCLING_FORCE_FULL_PAGE_OCR", "0").strip().lower() in {"1", "true", "yes", "on"}
    do_picture_description = os.getenv("NAVEXA_DOCLING_PICTURE_DESCRIPTION", "0").strip().lower() in {"1", "true", "yes", "on"}
    enable_remote_services = os.getenv("NAVEXA_DOCLING_REMOTE_SERVICES", "0").strip().lower() in {"1", "true", "yes", "on"}
    image_mode_raw = os.getenv("NAVEXA_DOCLING_IMAGE_MODE", "placeholder").strip().lower() or "placeholder"

    converter = _get_converter(
        do_ocr=do_ocr,
        backend=backend,
        force_full_page_ocr=force_full_page_ocr,
        do_picture_description=do_picture_description,
        enable_remote_services=enable_remote_services,
    )
    result = converter.convert(pdf_path)
    document = getattr(result, "document", result)

    blocks: List[LayoutBlock] = []
    page_text_map: Dict[int, List[str]] = {}
    page_order: Dict[int, int] = {}

    for item in list(getattr(document, "texts", []) or []):
        text = (getattr(item, "text", None) or getattr(item, "orig", None) or "").strip()
        if not text:
            continue

        prov = list(getattr(item, "prov", []) or [])
        page_no = int(getattr(prov[0], "page_no", 1) or 1) if prov else 1
        order = page_order.get(page_no, 0)
        page_order[page_no] = order + 1

        label = getattr(item, "label", "text")
        block_type = getattr(label, "value", None) or str(label)

        blocks.append(LayoutBlock(block_id=f"d{page_no}_{order}", page=page_no, order=order, text=text, block_type=block_type))
        page_text_map.setdefault(page_no, []).append(text)

    blocks.sort(key=lambda b: (b.page, b.order))
    if not blocks:
        raise RuntimeError("Docling returned no text blocks.")

    page_texts = {p: "\n".join(parts).strip() for p, parts in page_text_map.items()}
    num_pages = len(getattr(document, "pages", {}) or {})

    if resolved_output_format == "text":
        output_text = "\n\n".join(page_texts[p] for p in sorted(page_texts) if page_texts[p].strip())
        if not output_text.strip():
            raise RuntimeError("Docling returned empty text output.")
        return output_text, blocks, page_texts, num_pages

    markdown_kwargs: Dict[str, object] = {}
    if image_mode_raw == "none":
        markdown_kwargs["image_mode"] = ImageRefMode.PLACEHOLDER
        markdown_kwargs["image_placeholder"] = ""
    else:
        mode_map = {
            "placeholder": ImageRefMode.PLACEHOLDER,
            "embedded": ImageRefMode.EMBEDDED,
            "referenced": ImageRefMode.REFERENCED,
        }
        markdown_kwargs["image_mode"] = mode_map.get(image_mode_raw, ImageRefMode.PLACEHOLDER)

    markdown = document.export_to_markdown(**markdown_kwargs)
    if not markdown or not markdown.strip():
        raise RuntimeError("Docling returned empty markdown.")
    return markdown, blocks, page_texts, num_pages
