"""Thin wrapper around marker-pdf.

The model dict is built once at import time so subsequent requests reuse the
same weights on GPU. Don't construct PdfConverter at module level — its config
depends on the per-request `mode` (LLM enhancement on or off).
"""

from __future__ import annotations

import logging
import os
from pathlib import Path
from typing import Literal

import torch
from marker.converters.pdf import PdfConverter
from marker.models import create_model_dict
from marker.output import text_from_rendered

log = logging.getLogger(__name__)

Mode = Literal["fast", "quality"]

# Loaded once. ~5 GB on disk, lives in /data/hf in production.
_models = create_model_dict()
log.info("marker model dict loaded (cuda=%s)", torch.cuda.is_available())

CLAUDE_MODEL = os.getenv("CLAUDE_MODEL", "claude-sonnet-4-6")


def device() -> str:
    return "gpu" if torch.cuda.is_available() else "cpu"


def parse(pdf_path: str | Path, mode: Mode = "fast") -> tuple[str, int]:
    """Run marker on the PDF at `pdf_path`. Returns (markdown, page_count)."""
    config: dict[str, object] = {"output_format": "markdown"}
    llm_service: str | None = None
    if mode == "quality":
        api_key = os.getenv("ANTHROPIC_API_KEY")
        if not api_key:
            raise RuntimeError(
                "parse_mode=quality requires ANTHROPIC_API_KEY to be set."
            )
        config.update(
            {
                "use_llm": True,
                "claude_api_key": api_key,
                "claude_model_name": CLAUDE_MODEL,
            }
        )
        # Marker's PdfConverter takes a dotted class path STRING here and
        # resolves it via strings_to_classes(). It then constructs the
        # service using the matching keys in `config` (claude_api_key etc.).
        llm_service = "marker.services.claude.ClaudeService"

    converter = PdfConverter(
        artifact_dict=_models,
        config=config,
        llm_service=llm_service,
    )
    rendered = converter(str(pdf_path))
    text, _, _ = text_from_rendered(rendered)
    return text, _page_count(rendered, pdf_path)


def _page_count(rendered: object, pdf_path: str | Path) -> int:
    """Pull page count from marker's rendered metadata, fall back to pypdf.

    marker 1.x exposes per-page entries on `rendered.metadata` (key has shifted
    across versions — `page_stats` in recent releases). Probe a few likely
    keys, then fall back to reading the source PDF.
    """
    metadata = getattr(rendered, "metadata", None) or {}
    if isinstance(metadata, dict):
        for key in ("page_stats", "pages", "page_metadata"):
            value = metadata.get(key)
            if isinstance(value, list) and value:
                return len(value)
            if isinstance(value, int) and value > 0:
                return value

    try:
        from pypdf import PdfReader

        return len(PdfReader(str(pdf_path)).pages)
    except Exception:  # noqa: BLE001 — page count is best-effort metadata
        log.warning("could not determine page_count for %s", pdf_path)
        return 0