"""Hugging Face Inference Endpoints custom handler for monocular depth.

HF's managed task list does not expose "Depth Estimation", so we deploy
Depth-Anything-V2 through a custom handler instead. HF's default inference
container auto-detects an `EndpointHandler` class in this file and calls it per
request, bypassing the fixed task list.

Request  : raw image bytes (Content-Type: image/*) OR JSON {"inputs": "<b64>"}.
Response : JSON {"depth": "<base64 PNG>", "width": W, "height": H} — an 8-bit
           grayscale depth map (near = bright). This matches the AI engine's
           `hf_client.depth()` contract, which decodes {"depth": <base64 png>}.
"""
from __future__ import annotations

import base64
import io
from typing import Any, Dict

import numpy as np
from PIL import Image


class EndpointHandler:
    def __init__(self, model_dir: str = "", **kwargs: Any) -> None:
        # HF's inference toolkit calls this as EndpointHandler(model_dir=...),
        # so the first positional/keyword arg MUST be named `model_dir`.
        # Lazy heavy imports so container build/health checks stay light.
        import torch
        from transformers import pipeline

        self._torch = torch
        device = 0 if torch.cuda.is_available() else -1
        # This repo holds only the handler (no weights), so `model_dir` has no
        # model config. Always load the depth model from its public repo id.
        import os

        model = os.environ.get("TARGET_MODEL_ID", "depth-anything/Depth-Anything-V2-Large-hf")
        self.pipe = pipeline("depth-estimation", model=model, device=device)

    def __call__(self, data: Dict[str, Any]) -> Dict[str, Any]:
        image = self._read_image(data)
        with self._torch.inference_mode():
            out = self.pipe(image)

        depth = out["depth"] if isinstance(out, dict) else out
        arr = self._to_uint8(depth, image.size)

        buf = io.BytesIO()
        Image.fromarray(arr, mode="L").save(buf, format="PNG")
        b64 = base64.b64encode(buf.getvalue()).decode("ascii")
        h, w = arr.shape
        return {"depth": b64, "width": w, "height": h}

    # --- helpers -----------------------------------------------------------
    @staticmethod
    def _read_image(data: Dict[str, Any]) -> Image.Image:
        # HF passes raw request bytes under "inputs" (bytes) for binary payloads,
        # or a base64 string when sent as JSON.
        inputs = data.get("inputs", data)
        if isinstance(inputs, (bytes, bytearray)):
            raw = bytes(inputs)
        elif isinstance(inputs, str):
            raw = base64.b64decode(inputs)
        elif isinstance(inputs, Image.Image):
            return inputs.convert("RGB")
        else:
            raise ValueError("Unsupported input type for depth handler")
        return Image.open(io.BytesIO(raw)).convert("RGB")

    @staticmethod
    def _to_uint8(depth: Any, size) -> np.ndarray:
        """Normalize a depth output (PIL image or array/tensor) to uint8 HxW."""
        if isinstance(depth, Image.Image):
            arr = np.asarray(depth, dtype=np.float32)
        else:
            arr = np.asarray(depth, dtype=np.float32)
        if arr.ndim == 3:
            arr = arr[..., 0]
        lo, hi = float(arr.min()), float(arr.max())
        if hi - lo < 1e-6:
            arr = np.zeros_like(arr)
        else:
            arr = (arr - lo) / (hi - lo)
        return (arr * 255.0).clip(0, 255).astype(np.uint8)