"""The on-prem GPU backend: open weights served on vLLM, for hosts with a GPU and no hosted API (#83).

Two stages, two models, both reached through the provider seam so no call site changes:

- Page transcription. `reocr.py` sends each page image to a vision model. olmOCR-2-7B is the default,
  a 7B model trained for document transcription, served on its own vLLM endpoint.
- Extraction. `llm_extraction.py` forces one tool call whose schema `fields.py` derives. A larger
  instruct model answers it, and vLLM constrains generation to that schema with xgrammar, so the output
  cannot take a shape the checklist rejects.

Configured with four environment variables. `VLLM_URL` and `VLLM_MODEL` name the extractor.
`VLLM_REOCR_URL` and `VLLM_REOCR_MODEL` name the transcription endpoint, each falling back to the
extractor's value and the olmOCR default, so a single server hosting one model still works.

The backend reports unavailable until `VLLM_URL` and `VLLM_MODEL` are set, so the settings list disables
it on a host with no server rather than letting a switch fail at extraction (#94).
"""

from __future__ import annotations

import os
from typing import Any, Optional

from endopath import llm

GPU_VLLM = "gpu_vllm"

URL_ENV = "VLLM_URL"
MODEL_ENV = "VLLM_MODEL"
REOCR_URL_ENV = "VLLM_REOCR_URL"
REOCR_MODEL_ENV = "VLLM_REOCR_MODEL"

# The transcription model the ticket names. Overridable, so a host serving a different vision model
# points at it without a code change.
DEFAULT_REOCR_MODEL = "allenai/olmOCR-2-7B-1025"

# vLLM's constrained-decoding backend. The extraction request sends the checklist schema as `guided_json`
# alongside this, so the server admits only tokens the schema allows.
GUIDED_DECODING = "xgrammar"


def _env(name: str) -> str:
    return (os.environ.get(name) or "").strip()


def _require(name: str) -> str:
    value = _env(name)
    if not value:
        raise RuntimeError(
            f"{name} is not set. The GPU backend serves open weights on vLLM: set {URL_ENV} and "
            f"{MODEL_ENV} to the extractor endpoint, and optionally {REOCR_URL_ENV} and "
            f"{REOCR_MODEL_ENV} for a separate transcription endpoint."
        )
    return value


def build_extraction_client() -> Any:
    """The extractor: a larger instruct model on vLLM, with xgrammar holding its output to the checklist
    schema `llm_extraction.build_checklist_tool` derives from `fields.py`."""
    return llm.openai_compatible_client(
        _require(URL_ENV), _require(MODEL_ENV), guided_decoding=GUIDED_DECODING
    )


def build_reocr_client() -> Any:
    """The page-transcription model. Falls back to the extractor's endpoint when only one server runs, so
    a single-server host still transcribes."""
    base_url = _env(REOCR_URL_ENV) or _require(URL_ENV)
    model = _env(REOCR_MODEL_ENV) or DEFAULT_REOCR_MODEL
    return llm.openai_compatible_client(base_url, model)


def _extraction_model() -> Optional[str]:
    """The extractor model this backend serves, recorded on the activity a run writes (#151). None when
    VLLM_MODEL is unset, because no extraction can have run under it."""
    return _env(MODEL_ENV) or None


def _available() -> bool:
    """Whether a vLLM endpoint is configured. The settings list disables the backend until it is, rather
    than offering a switch that fails at the first model call."""
    return bool(_env(URL_ENV) and _env(MODEL_ENV))


llm.register_backend(
    llm.Backend(
        GPU_VLLM,
        "GPU on-prem (open weights on vLLM)",
        needs_key=False,
        provider=llm.OPENAI_COMPATIBLE,
        build=build_extraction_client,
        build_reocr=build_reocr_client,
        runs_llm_extraction=True,
        extraction_model=_extraction_model,
        # No classical page reader: the vision model transcribes the page through the seam, the same way
        # the cloud backend does.
        page_reader=None,
        available=_available,
    )
)