"""Nuclio handler for Ultralytics YOLOv8 pose models exported to OpenVINO (CVAT skeleton output)."""

from __future__ import annotations

import base64
import io
import json
import os
from pathlib import Path
from typing import Any

import cv2
import numpy as np
from openvino import Core
from PIL import Image

MODELS_DIR = Path(os.environ.get("MODELS_DIR", "/models"))
CLASSES_PATH = Path(os.environ.get("CLASSES_PATH", "/opt/nuclio/classes.json"))
DEFAULT_IMGSZ = int(os.environ.get("MODEL_IMGSZ", "640"))
DEFAULT_NUM_KEYPOINTS = int(os.environ.get("NUM_KEYPOINTS", "9"))
DEFAULT_CONF_THRESHOLD = float(os.environ.get("CONF_THRESHOLD", "0.25"))
DEFAULT_IOU_THRESHOLD = float(os.environ.get("IOU_THRESHOLD", "0.45"))


def load_classes(path: Path) -> dict[int, str]:
    with path.open(encoding="utf-8") as handle:
        payload = json.load(handle)
    class_map = payload.get("class", payload)
    return {int(class_id): str(name) for class_id, name in class_map.items()}


def find_model_xml(models_dir: Path) -> Path:
    candidates = [
        models_dir / "best.xml",
        models_dir / "best_openvino_model" / "best.xml",
        *sorted(models_dir.glob("*.xml")),
        *sorted(models_dir.glob("**/*.xml")),
    ]
    for candidate in candidates:
        if candidate.is_file():
            return candidate
    raise FileNotFoundError(f"No OpenVINO XML model found under {models_dir}")


def letterbox(
    image: np.ndarray,
    new_shape: tuple[int, int] = (640, 640),
    color: tuple[int, int, int] = (114, 114, 114),
) -> tuple[np.ndarray, float, tuple[float, float]]:
    height, width = image.shape[:2]
    target_height, target_width = new_shape
    scale = min(target_height / height, target_width / width)
    new_unpad_width = int(round(width * scale))
    new_unpad_height = int(round(height * scale))
    resized = cv2.resize(image, (new_unpad_width, new_unpad_height), interpolation=cv2.INTER_LINEAR)

    pad_width = target_width - new_unpad_width
    pad_height = target_height - new_unpad_height
    pad_left = pad_width / 2
    pad_top = pad_height / 2
    padded = cv2.copyMakeBorder(
        resized,
        int(round(pad_top - 0.1)),
        int(round(pad_height - pad_top)),
        int(round(pad_left - 0.1)),
        int(round(pad_width - pad_left)),
        cv2.BORDER_CONSTANT,
        value=color,
    )
    return padded, scale, (pad_left, pad_top)


def preprocess_image(image_bgr: np.ndarray, imgsz: int) -> tuple[np.ndarray, float, tuple[float, float]]:
    letterboxed, scale, pad = letterbox(image_bgr, new_shape=(imgsz, imgsz))
    rgb = letterboxed[:, :, ::-1].transpose(2, 0, 1)
    tensor = np.expand_dims(rgb, axis=0).astype(np.float32) / 255.0
    return tensor, scale, pad


def xywh_to_xyxy(boxes: np.ndarray) -> np.ndarray:
    converted = np.empty_like(boxes)
    converted[:, 0] = boxes[:, 0] - boxes[:, 2] / 2
    converted[:, 1] = boxes[:, 1] - boxes[:, 3] / 2
    converted[:, 2] = boxes[:, 0] + boxes[:, 2] / 2
    converted[:, 3] = boxes[:, 1] + boxes[:, 3] / 2
    return converted


def box_iou(box: np.ndarray, boxes: np.ndarray) -> np.ndarray:
    inter_x1 = np.maximum(box[0], boxes[:, 0])
    inter_y1 = np.maximum(box[1], boxes[:, 1])
    inter_x2 = np.minimum(box[2], boxes[:, 2])
    inter_y2 = np.minimum(box[3], boxes[:, 3])
    inter_area = np.maximum(0.0, inter_x2 - inter_x1) * np.maximum(0.0, inter_y2 - inter_y1)
    box_area = (box[2] - box[0]) * (box[3] - box[1])
    boxes_area = (boxes[:, 2] - boxes[:, 0]) * (boxes[:, 3] - boxes[:, 1])
    return inter_area / (box_area + boxes_area - inter_area + 1e-6)


def non_max_suppression(
    boxes: np.ndarray,
    scores: np.ndarray,
    iou_threshold: float,
    max_detections: int = 300,
) -> list[int]:
    order = scores.argsort()[::-1]
    keep: list[int] = []
    while order.size > 0 and len(keep) < max_detections:
        current = int(order[0])
        keep.append(current)
        if order.size == 1:
            break
        remaining = order[1:]
        ious = box_iou(boxes[current], boxes[remaining])
        order = remaining[ious <= iou_threshold]
    return keep


def scale_boxes(
    input_shape: tuple[int, int],
    boxes: np.ndarray,
    image_shape: tuple[int, int],
    scale: float,
    pad: tuple[float, float],
) -> np.ndarray:
    boxes = boxes.copy()
    pad_x, pad_y = pad
    boxes[:, [0, 2]] -= pad_x
    boxes[:, [1, 3]] -= pad_y
    boxes[:, :4] /= scale
    boxes[:, [0, 2]] = boxes[:, [0, 2]].clip(0, image_shape[1])
    boxes[:, [1, 3]] = boxes[:, [1, 3]].clip(0, image_shape[0])
    return boxes


def scale_keypoints(
    keypoints: np.ndarray,
    scale: float,
    pad: tuple[float, float],
    image_shape: tuple[int, int],
) -> np.ndarray:
    scaled = keypoints.copy()
    pad_x, pad_y = pad
    scaled[..., 0] -= pad_x
    scaled[..., 1] -= pad_y
    scaled[..., :2] /= scale
    scaled[..., 0] = scaled[..., 0].clip(0, image_shape[1])
    scaled[..., 1] = scaled[..., 1].clip(0, image_shape[0])
    return scaled


def parse_pose_output(
    output: np.ndarray,
    num_classes: int,
    num_keypoints: int,
) -> np.ndarray:
    if output.ndim == 3:
        predictions = output[0].T
    elif output.ndim == 2:
        predictions = output
    else:
        raise ValueError(f"Unexpected model output rank: {output.ndim}")

    expected_channels = 4 + num_classes + num_keypoints * 3
    if predictions.shape[1] != expected_channels:
        raise ValueError(
            f"Expected {expected_channels} output channels for "
            f"{num_classes} classes and {num_keypoints} keypoints, "
            f"got {predictions.shape[1]}"
        )
    return predictions


def decode_detections(
    predictions: np.ndarray,
    class_names: dict[int, str],
    num_keypoints: int,
    input_shape: tuple[int, int],
    image_shape: tuple[int, int],
    scale: float,
    pad: tuple[float, float],
    conf_threshold: float,
    iou_threshold: float,
) -> list[dict[str, Any]]:
    boxes_xywh = predictions[:, :4]
    class_scores = predictions[:, 4 : 4 + len(class_names)]
    keypoints = predictions[:, 4 + len(class_names) :].reshape(-1, num_keypoints, 3)

    class_ids = np.argmax(class_scores, axis=1)
    confidences = class_scores[np.arange(class_scores.shape[0]), class_ids]
    mask = confidences >= conf_threshold
    if not np.any(mask):
        return []

    boxes_xywh = boxes_xywh[mask]
    boxes_xyxy = xywh_to_xyxy(boxes_xywh)
    keypoints = keypoints[mask]
    class_ids = class_ids[mask]
    confidences = confidences[mask]

    keep = non_max_suppression(boxes_xyxy, confidences, iou_threshold=iou_threshold)
    if not keep:
        return []

    boxes_xyxy = boxes_xyxy[keep]
    keypoints = keypoints[keep]
    class_ids = class_ids[keep]
    confidences = confidences[keep]

    image_hw = (image_shape[0], image_shape[1])
    boxes_xyxy = scale_boxes(input_shape, boxes_xyxy, image_hw, scale, pad)
    keypoints = scale_keypoints(keypoints, scale, pad, image_hw)

    detections: list[dict[str, Any]] = []
    for box, kpts, class_id, confidence in zip(boxes_xyxy, keypoints, class_ids, confidences):
        detections.append(
            {
                "class_id": int(class_id),
                "label": class_names[int(class_id)],
                "confidence": float(confidence),
                "box": box,
                "keypoints": kpts,
            }
        )
    return detections


def build_skeleton_results(
    detections: list[dict[str, Any]],
    num_keypoints: int,
    conf_threshold: float,
) -> list[dict[str, Any]]:
    results: list[dict[str, Any]] = []
    sublabel_names = [str(index) for index in range(num_keypoints)]

    for detection in detections:
        elements = []
        for index, sublabel_name in enumerate(sublabel_names):
            x_coord, y_coord, kpt_conf = detection["keypoints"][index]
            elements.append(
                {
                    "label": sublabel_name,
                    "type": "points",
                    "outside": 0 if float(kpt_conf) >= conf_threshold else 1,
                    "points": [float(x_coord), float(y_coord)],
                    "confidence": str(float(kpt_conf)),
                }
            )

        if all(element["outside"] for element in elements):
            continue

        results.append(
            {
                "confidence": str(detection["confidence"]),
                "label": detection["label"],
                "type": "skeleton",
                "elements": elements,
            }
        )
    return results


def init_context(context) -> None:
    context.logger.info("Initializing OpenVINO YOLO pose handler")
    if not CLASSES_PATH.is_file():
        raise FileNotFoundError(f"Classes file not found: {CLASSES_PATH}")

    class_names = load_classes(CLASSES_PATH)
    model_xml = find_model_xml(MODELS_DIR)
    context.logger.info(f"Loading OpenVINO model from {model_xml}")

    core = Core()
    compiled_model = core.compile_model(model_xml, "CPU")
    input_layer = compiled_model.input(0)
    input_shape = tuple(input_layer.shape)
    if len(input_shape) == 4:
        imgsz = int(input_shape[2])
    else:
        imgsz = DEFAULT_IMGSZ

    context.user_data.compiled_model = compiled_model
    context.user_data.class_names = class_names
    context.user_data.num_keypoints = DEFAULT_NUM_KEYPOINTS
    context.user_data.imgsz = imgsz
    context.logger.info(
        f"Ready: classes={class_names}, imgsz={imgsz}, keypoints={DEFAULT_NUM_KEYPOINTS}"
    )


def handler(context, event):
    try:
        payload = event.body
        if isinstance(payload, (bytes, bytearray)):
            payload = json.loads(payload.decode("utf-8"))
        if isinstance(payload, str):
            payload = json.loads(payload)

        image_b64 = payload["image"]
        threshold = float(payload.get("threshold", DEFAULT_CONF_THRESHOLD))
        image_bytes = base64.b64decode(image_b64)
        image_rgb = np.array(Image.open(io.BytesIO(image_bytes)).convert("RGB"))
        image_bgr = image_rgb[:, :, ::-1]

        tensor, scale, pad = preprocess_image(image_bgr, context.user_data.imgsz)
        outputs = context.user_data.compiled_model([tensor])
        raw_output = next(iter(outputs.values()))
        predictions = parse_pose_output(
            np.array(raw_output),
            num_classes=len(context.user_data.class_names),
            num_keypoints=context.user_data.num_keypoints,
        )
        detections = decode_detections(
            predictions=predictions,
            class_names=context.user_data.class_names,
            num_keypoints=context.user_data.num_keypoints,
            input_shape=(context.user_data.imgsz, context.user_data.imgsz),
            image_shape=image_bgr.shape,
            scale=scale,
            pad=pad,
            conf_threshold=threshold,
            iou_threshold=DEFAULT_IOU_THRESHOLD,
        )
        results = build_skeleton_results(
            detections,
            num_keypoints=context.user_data.num_keypoints,
            conf_threshold=threshold,
        )
        context.logger.info(f"Returning {len(results)} skeleton detections")
        return context.Response(
            body=json.dumps(results),
            headers={},
            content_type="application/json",
            status_code=200,
        )
    except Exception as exc:
        context.logger.error(f"Pose handler failed: {exc}", exc_info=True)
        raise