"""
DoclingParser — docling pipeline machinery behind the `Parser` ABC.

`DoclingParser` is abstract: it owns the converter construction, the run
fingerprint, and the `parse` template, but defers the pipeline-option bundle
to a concrete subclass. The configuration is a *type*, not a runtime flag:

- `TunedFinancialParser` — the canonical configuration (TableFormerMode.ACCURATE,
  DoclingParse backend, OCR, picture classification, page_batch_size=32). The
  accelerator device defaults to AUTO and is resolved at runtime; RapidOCR is
  switched to CUDA only when the resolved device is CUDA. The default everywhere.
- `LegacyDoclingParser` — docling's default backend and PdfPipelineOptions,
  with do_ocr / do_table_structure / do_cell_matching set from the
  constructor. OCR is off unless a caller passes `do_ocr=True`, and no caller
  does, while docling's own default is on. Kept
  for regression comparison and as a fallback; select it with `--preset legacy`.

Use `parser_for_preset(preset)` to map a preset string to the right class.

The header-consolidation post-processor that builds `TableMetadata` lives in
`core.consolidation` and operates on the DoclingDocument. That keeps these
classes focused on a single job: source -> `ParseResult`, which carries the
DoclingDocument on `.document` beside the page scores, parsed-page cells and
per-table OCR/native provenance.
"""

import os
from abc import abstractmethod
from pathlib import Path
from typing import Any, ClassVar

import logfire
from docling.backend.docling_parse_backend import DoclingParseDocumentBackend
from docling.datamodel.accelerator_options import AcceleratorDevice, AcceleratorOptions
from docling.datamodel.base_models import InputFormat
from docling.datamodel.pipeline_options import (
    PdfPipelineOptions,
    RapidOcrOptions,
    TableFormerMode,
)
from docling.datamodel.settings import settings
from docling.document_converter import DocumentConverter, PdfFormatOption
from docling.utils.accelerator_utils import decide_device

from quber.core.parsers.base import Parser
from quber.core.parsers.result import ParseResult

CANONICAL_PAGE_BATCH_SIZE = 32


class DoclingParser(Parser):
    """Abstract docling parser. Subclasses define the pipeline-option bundle.

    Subclasses set the `preset` and `backend_label` class attributes and
    implement `build_pipeline_options`. Everything else (converter, perf
    settings, fingerprint, the `parse` template) is shared.
    """

    #: Stable name of the configuration, surfaced in the run fingerprint.
    preset: ClassVar[str]
    #: How the PDF backend is reported in the fingerprint.
    backend_label: ClassVar[str]

    def __init__(
        self,
        do_ocr: bool = False,
        do_table_structure: bool = True,
        do_cell_matching: bool = True,
        page_batch_size: int = CANONICAL_PAGE_BATCH_SIZE,
        accelerator_device: AcceleratorDevice = AcceleratorDevice.AUTO,
        num_threads: int | None = None,
    ) -> None:
        self.do_ocr = do_ocr
        self.do_table_structure = do_table_structure
        self.do_cell_matching = do_cell_matching
        self.page_batch_size = page_batch_size
        self.accelerator_device = accelerator_device
        self.num_threads = num_threads if num_threads is not None else (os.cpu_count() or 4)

    @abstractmethod
    def build_pipeline_options(self) -> PdfPipelineOptions:
        """Return the docling pipeline options for this configuration."""
        ...

    def build_converter(self, pipeline_options: PdfPipelineOptions) -> DocumentConverter:
        """Default converter (no explicit backend). Overridden by presets that pin one."""
        return DocumentConverter(
            format_options={InputFormat.PDF: PdfFormatOption(pipeline_options=pipeline_options)}
        )

    def apply_perf_settings(self) -> None:
        """Hook for process-wide perf tweaks. No-op by default."""

    def fingerprint(self, pipeline_options: PdfPipelineOptions) -> dict[str, Any]:
        device: Any = pipeline_options.accelerator_options.device
        device_name = device.name if hasattr(device, "name") else str(device)
        mode: Any = pipeline_options.table_structure_options.mode  # type: ignore[union-attr]
        mode_name = mode.name if hasattr(mode, "name") else str(mode)
        ocr_engine = pipeline_options.ocr_options.__class__.__name__
        ocr_cuda = isinstance(pipeline_options.ocr_options, RapidOcrOptions) and bool(
            (pipeline_options.ocr_options.rapidocr_params or {}).get(
                "EngineConfig.onnxruntime.use_cuda", False
            )
        )
        resolved_device = decide_device(self.accelerator_device.value)
        return {
            "preset": self.preset,
            "table_mode": mode_name,
            "do_cell_matching": pipeline_options.table_structure_options.do_cell_matching,  # type: ignore[union-attr]
            "do_ocr": pipeline_options.do_ocr,
            "do_picture_classification": pipeline_options.do_picture_classification,
            "generate_picture_images": pipeline_options.generate_picture_images,
            "backend": self.backend_label,
            "accelerator_device_requested": device_name,
            "accelerator_device_resolved": resolved_device,
            "num_threads": pipeline_options.accelerator_options.num_threads,
            "page_batch_size": self.page_batch_size,
            "ocr_engine": ocr_engine,
            "ocr_cuda": ocr_cuda,
        }

    def parse(self, source: Path) -> ParseResult:
        pipeline_options = self.build_pipeline_options()
        # Keep docling's parsed-page cells (text, box, from_ocr, confidence) alive
        # past document assembly. Off by default, so docling frees them; they are
        # the OCR-vs-native signal the reconciliation needs, so capture them for
        # every preset.
        pipeline_options.generate_parsed_pages = True
        converter = self.build_converter(pipeline_options)
        self.apply_perf_settings()

        fingerprint = self.fingerprint(pipeline_options)
        with logfire.span("docling_parser.parse", source=str(source), **fingerprint):
            result = converter.convert(str(source))
            return ParseResult.from_conversion(result)


class LegacyDoclingParser(DoclingParser):
    """Docling's default backend and PdfPipelineOptions, with OCR off unless asked for."""

    preset = "legacy"
    backend_label = "default"

    def build_pipeline_options(self) -> PdfPipelineOptions:
        pipeline_options = PdfPipelineOptions()
        pipeline_options.do_table_structure = self.do_table_structure
        pipeline_options.table_structure_options.do_cell_matching = self.do_cell_matching  # type: ignore[misc]
        pipeline_options.do_ocr = self.do_ocr
        return pipeline_options


class TunedFinancialParser(DoclingParser):
    """Canonical (tuned-financial) configuration for dense SEC financial documents."""

    preset = "tuned-financial"
    backend_label = DoclingParseDocumentBackend.__name__

    def build_pipeline_options(self) -> PdfPipelineOptions:
        pipeline_options = PdfPipelineOptions()
        pipeline_options.do_table_structure = self.do_table_structure
        pipeline_options.table_structure_options.do_cell_matching = self.do_cell_matching  # type: ignore[misc]
        pipeline_options.do_ocr = True
        pipeline_options.table_structure_options.mode = TableFormerMode.ACCURATE  # type: ignore[misc]
        pipeline_options.do_picture_classification = True
        pipeline_options.generate_picture_images = True
        pipeline_options.accelerator_options = AcceleratorOptions(
            device=self.accelerator_device,
            num_threads=self.num_threads,
        )
        # docling propagates accelerator_options.device=CUDA to RapidOCR's
        # paddle/torch engines but not the onnxruntime engine (which is the default).
        # Flip EngineConfig.onnxruntime.use_cuda explicitly so the OCR pass actually
        # runs on GPU. CUDA 12 runtime libs preloaded in quber/__init__.py.
        # Gate on the *resolved* device (decide_device handles AUTO -> cuda:0/cpu/mps)
        # so non-GPU hosts (e.g., Intel MacBook) skip the override and run OCR on CPU.
        resolved_device = decide_device(self.accelerator_device.value)
        if resolved_device.startswith("cuda"):
            pipeline_options.ocr_options = RapidOcrOptions(
                rapidocr_params={
                    "EngineConfig.onnxruntime.use_cuda": True,
                    "EngineConfig.onnxruntime.cuda_ep_cfg.device_id": 0,
                },
            )
        return pipeline_options

    def build_converter(self, pipeline_options: PdfPipelineOptions) -> DocumentConverter:
        return DocumentConverter(
            format_options={
                InputFormat.PDF: PdfFormatOption(
                    pipeline_options=pipeline_options,
                    backend=DoclingParseDocumentBackend,
                )
            }
        )

    def apply_perf_settings(self) -> None:
        settings.perf.page_batch_size = self.page_batch_size


#: Preset name -> concrete parser class. The CLI `--preset` choices mirror these keys.
PARSERS_BY_PRESET: dict[str, type[DoclingParser]] = {
    "tuned-financial": TunedFinancialParser,
    "legacy": LegacyDoclingParser,
}


def parser_for_preset(preset: str, **kwargs: Any) -> DoclingParser:
    """Instantiate the parser for a preset name (forwards kwargs to the class)."""
    try:
        cls = PARSERS_BY_PRESET[preset]
    except KeyError:
        valid = ", ".join(sorted(PARSERS_BY_PRESET))
        raise ValueError(f"Unknown parser preset {preset!r}; valid presets: {valid}") from None
    return cls(**kwargs)