"""Hugging Face Inference Endpoints custom handler: lead sheet image -> MusicXML with chord symbols.

Model: JuanCarlosMartinezSevilla/jazzmus-model (MIT): a YOLO staff detector, then the SMT
image-to-**kern model on each staff. The detector runs as ONNX on onnxruntime (MIT) with the
letterbox, box decoding and NMS below in numpy; the ultralytics package is not used. The
staff-cropping functions are copied from ../predict.py (2026-10-18 fixes: STAFF_PAD,
select_staff_boxes) and work on an in-memory image.

Errors are raised as EndpointError with a fixed short code. The Inference Toolkit turns any
exception from a handler into HTTP 400 with body {"error": str(exception)}; a handler cannot
choose the status code (see README.md).
"""
import base64
import binascii
import io
import logging
import os
import threading
from pathlib import Path

os.environ.setdefault("HF_HUB_OFFLINE", "1")  # load nothing from the network

import numpy as np  # noqa: E402
import onnxruntime as ort  # noqa: E402
import torch  # noqa: E402
from PIL import Image  # noqa: E402
from safetensors.torch import load_file  # noqa: E402
from torch.nn import Conv1d  # noqa: E402

from jazzmus_smt.configuration_smt import SMTConfig  # noqa: E402
from jazzmus_smt.modeling_smt import SMTModelForCausalLM  # noqa: E402
from kern2musicxml import kern_to_musicxml  # noqa: E402

log = logging.getLogger(__name__)

TORCH_THREADS = 2
MAX_BYTES = 760 * 1024
MAX_B64_CHARS = 4 * ((MAX_BYTES + 2) // 3)
MAX_SIDE = 6000  # refused above this
# Refused above this many pixels (6000 x 4000): a well-compressed page under MAX_BYTES can decode to
# hundreds of MB, so the bound is on pixels, not bytes or sides.
MAX_PIXELS = 24_000_000
# One decode at a time, whatever the toolkit's request concurrency: each holds up to a few hundred MB
# (four concurrent 6000 x 6000 decodes measured +1.5 GB on the endpoint's 4 GB).
DECODE_LOCK = threading.Lock()
WORK_SIDE = 3000  # larger images are scaled down before detection, as in app.py
MAX_STAVES = 12
CONTENT_TYPES = {"image/png": "PNG", "image/jpeg": "JPEG", "image/webp": "WEBP"}

# ---- YOLO staff detector on onnxruntime ------------------------------------------------------
# The same steps and defaults as ultralytics 8.4 predict on the .pt model: a rectangular letterbox
# (long side 640, padded to a multiple of 32 with grey 114), class-wise NMS at IoU 0.7 over boxes
# with confidence >= 0.25, at most 300, then boxes scaled back to the image.
DETECT_SIZE = 640
STRIDE = 32
DETECT_CONF = 0.25
NMS_IOU = 0.7
MAX_DET = 300
STAFF_CLASS = 12  # "staff" in the model's names (stored in the .onnx metadata)


def resize_linear(image: Image.Image, size: tuple[int, int]) -> Image.Image:
    """Bilinear resize without antialiasing, sampling at pixel centres as cv2.INTER_LINEAR does.

    The detector and the SMT model were used with cv2.resize; PIL's resize() antialiases when it
    shrinks, which moved staff boxes by up to 18 px. An affine transform samples like cv2.
    """
    if size == image.size:
        return image
    sx, sy = image.width / size[0], image.height / size[1]
    return image.transform(size, Image.AFFINE, (sx, 0, 0, 0, sy, 0), resample=Image.BILINEAR)


def letterbox(image: Image.Image) -> tuple[np.ndarray, tuple[float, float], tuple[int, int]]:
    """RGB image -> (1, 3, H, W) float32 input, (gain_x, gain_y), (pad_x, pad_y)."""
    w, h = image.size
    r = min(DETECT_SIZE / h, DETECT_SIZE / w)
    new_w, new_h = round(w * r), round(h * r)
    dw, dh = (DETECT_SIZE - new_w) % STRIDE / 2, (DETECT_SIZE - new_h) % STRIDE / 2
    top, bottom, left, right = round(dh - 0.1), round(dh + 0.1), round(dw - 0.1), round(dw + 0.1)
    canvas = Image.new("RGB", (new_w + left + right, new_h + top + bottom), (114, 114, 114))
    canvas.paste(resize_linear(image, (new_w, new_h)), (left, top))
    x = np.asarray(canvas, dtype=np.float32).transpose(2, 0, 1)[None] / 255.0
    return x, (new_w / w, new_h / h), (left, top)


def nms(boxes: np.ndarray, scores: np.ndarray, iou: float) -> list[int]:
    order, keep = scores.argsort()[::-1], []
    areas = (boxes[:, 2] - boxes[:, 0]) * (boxes[:, 3] - boxes[:, 1])
    while order.size:
        i, rest = order[0], order[1:]
        keep.append(int(i))
        xx1 = np.maximum(boxes[i, 0], boxes[rest, 0])
        yy1 = np.maximum(boxes[i, 1], boxes[rest, 1])
        xx2 = np.minimum(boxes[i, 2], boxes[rest, 2])
        yy2 = np.minimum(boxes[i, 3], boxes[rest, 3])
        inter = np.clip(xx2 - xx1, 0, None) * np.clip(yy2 - yy1, 0, None)
        order = rest[inter / (areas[i] + areas[rest] - inter + 1e-9) <= iou]
    return keep


def detect(session: ort.InferenceSession, image: Image.Image) -> list[tuple[float, int, tuple[float, float, float, float]]]:
    """All detections as (confidence, class, (x1, y1, x2, y2)) in image pixels."""
    x, (gx, gy), (px, py) = letterbox(image)
    pred = session.run(None, {"images": x})[0][0].T  # (anchors, 4 + classes): cx, cy, w, h, scores
    cls = pred[:, 4:].argmax(1)
    conf = pred[np.arange(len(pred)), 4 + cls]
    m = conf >= DETECT_CONF
    pred, cls, conf = pred[m], cls[m], conf[m]
    boxes = np.stack([pred[:, 0] - pred[:, 2] / 2, pred[:, 1] - pred[:, 3] / 2,
                      pred[:, 0] + pred[:, 2] / 2, pred[:, 1] + pred[:, 3] / 2], 1)
    keep = nms(boxes + cls[:, None] * 7680.0, conf, NMS_IOU)[:MAX_DET]  # offset by class: class-wise NMS
    w, h = image.size
    out = []
    for i in keep:
        x1, y1, x2, y2 = boxes[i]
        out.append((float(conf[i]), int(cls[i]), (
            float(np.clip((x1 - px) / gx, 0, w)), float(np.clip((y1 - py) / gy, 0, h)),
            float(np.clip((x2 - px) / gx, 0, w)), float(np.clip((y2 - py) / gy, 0, h)))))
    return out


# ---- copied from ../predict.py ------------------------------------------------------------
FIXED_IMG_HEIGHT = 128
MAX_IMG_WIDTH = 1000
STAFF_PAD = 0.30
MIN_CONFIDENCE = 0.5
MAX_OVERLAP = 0.5


def _vertical_overlap(a: tuple[int, int, int, int], b: tuple[int, int, int, int]) -> float:
    inter = min(a[3], b[3]) - max(a[1], b[1])
    return max(0, inter) / min(a[3] - a[1], b[3] - b[1])


def select_staff_boxes(detections: list[tuple[float, tuple[int, int, int, int]]]) -> list[tuple[int, int, int, int]]:
    """Keep confident boxes, best first, dropping any that overlap a kept one; return them top to bottom."""
    kept: list[tuple[int, int, int, int]] = []
    for conf, box in sorted(detections, key=lambda d: -d[0]):
        if conf < MIN_CONFIDENCE:
            continue
        if any(_vertical_overlap(box, k) > MAX_OVERLAP for k in kept):
            continue
        kept.append(box)
    return sorted(kept, key=lambda b: b[1])


def detect_staves(image: Image.Image, session: ort.InferenceSession) -> list[Image.Image]:
    """predict.detect_staves on an RGB PIL image, with the ONNX detector; returns greyscale crops."""
    gray = image.convert("L")
    detections = [(conf, tuple(map(int, box))) for conf, cls, box in detect(session, image) if cls == STAFF_CLASS]

    crops = []
    for x1, y1, x2, y2 in select_staff_boxes(detections)[:MAX_STAVES]:
        pad = int(round(STAFF_PAD * (y2 - y1)))
        crops.append(gray.crop((x1, max(0, y1 - pad), x2, min(gray.height, y2 + pad))))
    return crops


def preprocess_staff(img: Image.Image) -> torch.Tensor:
    width = int(np.ceil(img.width * FIXED_IMG_HEIGHT / img.height))
    width = min(width, MAX_IMG_WIDTH)
    img = resize_linear(img, (width, FIXED_IMG_HEIGHT))
    # jazzmus convert_img_to_tensor: greyscale, (1, H, W) float in [0, 1]
    return torch.from_numpy(np.asarray(img, dtype=np.float32) / 255.0)[None]


@torch.no_grad()
def transcribe_staff(model: SMTModelForCausalLM, image_tensor: torch.Tensor) -> str:
    predicted_sequence, _ = model.predict(input=image_tensor, convert_to_str=True)
    # jazzmus.dataset.tokenizer.untokenize
    return (
        "".join(predicted_sequence).replace("<t>", "\t").replace("<n>", "\n").replace("<s>", " ")
        .replace("<chord-pitch>", "").replace("<chord-extension>", "")
    )


def _sync_pe_device(model: SMTModelForCausalLM, device: torch.device):
    model.positional_2D.pe = model.positional_2D.pe.to(device)
    model.decoder.positional_1D.pe = model.decoder.positional_1D.pe.to(device)
# ---- end of copy ----------------------------------------------------------------------------


class EndpointError(ValueError):
    """A client-facing failure; str() is only the fixed code, never request content."""


MAGIC = {
    "image/png": lambda b: b.startswith(b"\x89PNG\r\n\x1a\n"),
    "image/jpeg": lambda b: b.startswith(b"\xff\xd8\xff"),
    "image/webp": lambda b: len(b) >= 12 and b[:4] == b"RIFF" and b[8:12] == b"WEBP",
}


def decode_image(data) -> Image.Image:
    """Validate the request body and return the image as RGB; raise EndpointError with a fixed code."""
    if not isinstance(data, dict):
        raise EndpointError("invalid_request")
    inputs = data.get("inputs")
    params = data.get("parameters") or {}
    if not isinstance(inputs, str) or not isinstance(params, dict):
        raise EndpointError("invalid_request")
    content_type = params.get("contentType")
    if content_type not in CONTENT_TYPES:
        raise EndpointError("unsupported_content_type")
    if len(inputs) > MAX_B64_CHARS:
        raise EndpointError("image_too_large")
    try:
        raw = base64.b64decode(inputs, validate=True)
    except (binascii.Error, ValueError):
        raise EndpointError("invalid_base64") from None
    if len(raw) > MAX_BYTES:
        raise EndpointError("image_too_large")
    if not MAGIC[content_type](raw):
        raise EndpointError("content_type_mismatch")
    with DECODE_LOCK:
        try:
            with Image.open(io.BytesIO(raw), formats=[CONTENT_TYPES[content_type]]) as im:
                # The header's size, checked before any pixel is decoded.
                if max(im.size) > MAX_SIDE or im.width * im.height > MAX_PIXELS:
                    raise EndpointError("image_dimensions_too_large")
                im.load()
                rgba = im.convert("RGBA")
        except EndpointError:
            raise
        except Exception:
            raise EndpointError("invalid_image") from None
        # Transparent areas become white paper, not black.
        image = Image.new("RGB", rgba.size, "white")
        image.paste(rgba, mask=rgba.getchannel("A"))
        del rgba
        image.thumbnail((WORK_SIDE, WORK_SIDE))
        return image


class EndpointHandler:
    def __init__(self, path: str = ""):
        torch.set_num_threads(TORCH_THREADS)
        model_dir = Path(path)
        self.device = torch.device("cpu")
        opts = ort.SessionOptions()
        opts.intra_op_num_threads = TORCH_THREADS
        opts.inter_op_num_threads = 1
        self.detector = ort.InferenceSession(
            str(model_dir / "yolo_staff_detector.onnx"), opts, providers=["CPUExecutionProvider"]
        )

        config = SMTConfig.from_pretrained(str(model_dir))
        sd = load_file(str(model_dir / "model.safetensors"))
        embed_size = sd["decoder.embedding.weight"].shape[0]
        out_size = sd["decoder.out_layer.weight"].shape[0]
        orig_out = config.out_categories
        config.out_categories = embed_size
        model = SMTModelForCausalLM(config)
        if embed_size != out_size:
            model.decoder.out_layer = Conv1d(config.d_model, out_size, kernel_size=1)
        model.load_state_dict(sd, strict=True)
        config.out_categories = orig_out
        model.to(self.device)
        _sync_pe_device(model, self.device)
        model.eval()
        self.smt = model
        self.lock = threading.Lock()  # one page at a time on the two cores

    def __call__(self, data: dict) -> dict:
        image = decode_image(data)
        try:
            with self.lock:
                staves = detect_staves(image, self.detector)
                if not staves:
                    raise EndpointError("no_staff_found")
                kerns = [transcribe_staff(self.smt, preprocess_staff(s).to(self.device)) for s in staves]
            musicxml, _ = kern_to_musicxml(kerns)
        except EndpointError:
            raise
        except Exception:
            log.exception("transcription failed")  # the traceback holds no request content
            raise EndpointError("transcription_failed") from None
        return {"musicxml": musicxml, "staves": len(staves)}