# Copyright 2026 AutomatosX
# SPDX-License-Identifier: AGPL-3.0-only
# Modified 2026-10-17 for v1.2; prior component notices are in NOTICE.
"""Digest-verified PyTorch inference over a local nine-file bundle.

This path loads the staged safetensors into the in-repo RT-DETRv2-M
architecture. It is not ``from_pretrained`` and it never marks a bundle
publishable.
"""

from __future__ import annotations

from pathlib import Path

import torch

from ax_doclayout._torch_rtdetrv2 import build_rtdetrv2_r50vd_m
from ax_doclayout.artifact import (
    MODEL_ID,
    STAGED_SAFETENSORS_SHA256,
    _reject_json_constant,
    _unique_object,
    verify_artifact_manifest,
)
from ax_doclayout.contract import Element, PageImage
from ax_doclayout.errors import ArtifactError, BackendUnavailableError
from ax_doclayout.postprocess import (
    DEFAULT_SCORE_THRESHOLD,
    elements_from_letterbox_detections,
)
from ax_doclayout.preprocess import letterbox_nchw, letterbox_page, orig_target_sizes

ARCHITECTURE_ID = "rtdetrv2_r50vd_m"


class LocalPytorchBundle:
    """Run the staged safetensors after the manifest binds every payload digest."""

    def __init__(
        self,
        directory: Path,
        *,
        expected_safetensors_sha256: str | None = STAGED_SAFETENSORS_SHA256,
    ) -> None:
        try:
            from safetensors import SafetensorError
            from safetensors.torch import load_file
        except ImportError as error:
            raise BackendUnavailableError("safetensors is not installed") from error
        manifest = verify_artifact_manifest(directory)
        if manifest.model_id != MODEL_ID:
            raise ArtifactError("artifact_model_id_invalid")
        weights_digest = ""
        for item in manifest.artifact_files:
            if item.name == "model.safetensors":
                weights_digest = item.sha256
                break
        if not weights_digest:
            raise ArtifactError("artifact_file_unavailable")
        if (
            expected_safetensors_sha256 is not None
            and weights_digest != expected_safetensors_sha256
        ):
            raise ArtifactError("artifact_safetensors_digest_mismatch")
        weights_path = directory / "model.safetensors"
        if not weights_path.is_file():
            raise ArtifactError("artifact_file_unavailable")
        try:
            state = load_file(str(weights_path))
        except (SafetensorError, OSError, ValueError) as error:
            raise ArtifactError("artifact_safetensors_unreadable") from error
        model = build_rtdetrv2_r50vd_m()
        try:
            model.load_state_dict(state, strict=True)
        except RuntimeError as error:
            raise ArtifactError("artifact_state_dict_mismatch") from error
        model.eval()
        self._model = model
        self._directory = directory

    def predict_elements(
        self,
        page: PageImage,
        *,
        threshold: float = DEFAULT_SCORE_THRESHOLD,
    ) -> tuple[Element, ...]:
        """Detect on one RGB page and return original-raster elements."""

        letterbox = letterbox_page(page)
        images = torch.from_numpy(letterbox_nchw(letterbox))
        sizes = torch.from_numpy(orig_target_sizes(letterbox))
        with torch.no_grad():
            labels, boxes, scores = self._model(images, sizes)
        return elements_from_letterbox_detections(
            labels=labels.cpu().numpy(),
            boxes=boxes.cpu().numpy(),
            scores=scores.cpu().numpy(),
            page=page,
            letterbox=letterbox,
            threshold=threshold,
        )


def read_publishable(path: Path) -> bool:
    """Read the staging receipt flag. Missing or malformed receipts fail closed."""

    import json

    try:
        payload = json.loads(
            path.read_text(encoding="utf-8"),
            object_pairs_hook=_unique_object,
            parse_constant=_reject_json_constant,
        )
    except (OSError, ValueError) as error:
        raise ArtifactError("staging_receipt_unreadable") from error
    if not isinstance(payload, dict) or "publishable" not in payload:
        raise ArtifactError("staging_publishable_missing")
    flag = payload["publishable"]
    if not isinstance(flag, bool):
        raise ArtifactError("staging_publishable_invalid")
    return flag