import base64
import io
from typing import Any

import numpy as np
import torch
import torch.nn.functional as F
from PIL import Image, ImageFilter
from transformers import AutoImageProcessor, AutoModelForDepthEstimation


MODEL_ID = "depth-anything/Depth-Anything-V2-Small-hf"
MAX_PIXELS = 12_000_000


class EndpointHandler:
    def __init__(self, path: str = "") -> None:
        self.device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
        model_path = path or MODEL_ID
        self.processor = AutoImageProcessor.from_pretrained(model_path)
        self.model = AutoModelForDepthEstimation.from_pretrained(model_path).to(self.device)
        self.model.eval()

    def __call__(self, data: dict[str, Any]) -> dict[str, Any]:
        encoded = data.get("inputs")
        if not isinstance(encoded, str) or not encoded:
            raise ValueError("inputs must be a base64-encoded image")
        try:
            image_bytes = base64.b64decode(encoded, validate=True)
            image = Image.open(io.BytesIO(image_bytes)).convert("RGB")
        except Exception as error:
            raise ValueError("inputs must contain a valid base64-encoded image") from error

        width, height = image.size
        if width < 1 or height < 1 or width * height > MAX_PIXELS:
            raise ValueError("image dimensions exceed the endpoint limit")

        inputs = self.processor(images=image, return_tensors="pt")
        inputs = {key: value.to(self.device) for key, value in inputs.items()}
        with torch.inference_mode():
            predicted = self.model(**inputs).predicted_depth
            predicted = F.interpolate(
                predicted.unsqueeze(1),
                size=(height, width),
                mode="bicubic",
                align_corners=False,
            )[0, 0]

        values = predicted.float().cpu().numpy()
        values = np.nan_to_num(values, nan=0.0, posinf=0.0, neginf=0.0)
        low = float(values.min())
        high = float(values.max())
        if high > low:
            normalized = np.clip((values - low) / (high - low), 0.0, 1.0)
        else:
            normalized = np.zeros_like(values)
        depth = Image.fromarray(np.round(normalized * 255).astype(np.uint8), mode="L")
        depth = depth.filter(ImageFilter.GaussianBlur(radius=0.8))

        output = io.BytesIO()
        depth.save(output, format="PNG", optimize=True)
        return {
            "depth_map_png": base64.b64encode(output.getvalue()).decode("ascii"),
            "width": width,
            "height": height,
        }