"""Hugging Face Inference Endpoints handler for the Kámárí CNN age model.

A cheap, platform-diverse fallback for the age model: host `cnn-age-v0` on a small CPU
Inference Endpoint (scale-to-zero) and point the gateway's third CNN fallback at it. Mirrors
the Modal serve logic: detect and crop the largest face, then estimate age, under-18
probability, uncertainty, and a sharpness-based quality.

Send the raw image bytes with `Content-Type: image/jpeg`. Returns:
    {estimated_age, p_under_18, uncertainty, face_quality, faces_detected, model_version}
"""
import base64
import io
import os

import numpy as np
import timm
import torch
import torch.nn as nn
from PIL import Image

MEAN = [0.485, 0.456, 0.406]
STD = [0.229, 0.224, 0.225]
MODEL_VERSION = "cnn_v0.1.0-hf"
BACKBONE = os.environ.get("CNN_BACKBONE", "tf_efficientnetv2_s")


class AgeNet(nn.Module):
    def __init__(self):
        super().__init__()
        self.backbone = timm.create_model(BACKBONE, pretrained=False, num_classes=0)
        d = self.backbone.num_features
        self.age = nn.Linear(d, 1)
        self.minor = nn.Linear(d, 1)
        self.logvar = nn.Linear(d, 1)

    def forward(self, x):
        f = self.backbone(x)
        return self.age(f).squeeze(1), self.minor(f).squeeze(1), self.logvar(f).squeeze(1)


def _to_bytes(data) -> bytes:
    """Accept the many shapes HF may hand a custom handler: raw bytes, a PIL image,
    a base64 / data-URL string, or a dict wrapping any of those under `inputs`."""
    if isinstance(data, dict):
        data = data.get("inputs", data)
    if isinstance(data, (bytes, bytearray)):
        return bytes(data)
    if isinstance(data, str):
        if data.startswith("data:"):
            data = data.split(",", 1)[1]
        return base64.b64decode(data)
    buf = io.BytesIO()          # already a PIL image
    data.convert("RGB").save(buf, format="JPEG")
    return buf.getvalue()


class EndpointHandler:
    def __init__(self, path: str = ""):
        import cv2  # noqa: F401 - imported lazily so the model loads even if cv2 is slow
        model = AgeNet()
        model.load_state_dict(
            torch.load(os.path.join(path, "best.pt"), map_location="cpu", weights_only=True))
        model.eval()
        torch.set_num_threads(max(1, os.cpu_count() or 2))
        self.model = model
        self.cv2 = cv2
        self.cascade = cv2.CascadeClassifier(
            cv2.data.haarcascades + "haarcascade_frontalface_default.xml")

    def __call__(self, data):
        try:
            return self._infer(_to_bytes(data))
        except Exception:
            return {"estimated_age": 0.0, "p_under_18": 0.5, "uncertainty": 1.0,
                    "face_quality": 0.0, "faces_detected": 0, "model_version": MODEL_VERSION}

    def _infer(self, image_bytes: bytes) -> dict:
        cv2 = self.cv2
        try:
            img = Image.open(io.BytesIO(image_bytes)).convert("RGB")
        except Exception:
            return {"estimated_age": 0.0, "p_under_18": 0.5, "uncertainty": 1.0,
                    "face_quality": 0.0, "faces_detected": 0, "model_version": MODEL_VERSION}
        w, h = img.size
        gray = cv2.cvtColor(np.asarray(img), cv2.COLOR_RGB2GRAY)
        faces = self.cascade.detectMultiScale(gray, scaleFactor=1.1, minNeighbors=5,
                                              minSize=(48, 48))
        if len(faces) == 0:
            return {"estimated_age": 0.0, "p_under_18": 0.5, "uncertainty": 1.0,
                    "face_quality": 0.0, "faces_detected": 0, "model_version": MODEL_VERSION}

        fx, fy, fw, fh = (int(v) for v in max(faces, key=lambda f: f[2] * f[3]))
        mx, my = int(fw * 0.3), int(fh * 0.3)
        x0, y0 = max(0, fx - mx), max(0, fy - my)
        x1, y1 = min(w, fx + fw + mx), min(h, fy + fh + my)
        face = img.crop((x0, y0, x1, y1)).resize((224, 224))

        arr = np.asarray(face, np.float32) / 255.0
        sharp = float(cv2.Laplacian(cv2.cvtColor(np.asarray(face), cv2.COLOR_RGB2GRAY),
                                    cv2.CV_64F).var())
        quality = float(np.clip(0.55 + sharp / 800.0, 0.0, 1.0))

        x = ((arr - MEAN) / STD).transpose(2, 0, 1)[None].astype(np.float32)
        with torch.no_grad():
            age, minor_logit, logvar = self.model(torch.from_numpy(x))
        p_under = 1.0 / (1.0 + float(np.exp(-float(minor_logit[0]))))
        uncertainty = float(np.clip(float(np.exp(0.5 * float(logvar[0]))) / 10.0, 0, 1))
        return {
            "estimated_age": round(float(age[0]), 1),
            "p_under_18": round(p_under, 3),
            "uncertainty": round(uncertainty, 3),
            "face_quality": round(quality, 3),
            "faces_detected": int(len(faces)),
            "model_version": MODEL_VERSION,
        }