# -*- coding: utf-8 -*-
"""YOLO 人脸检测 + 姿态估计（后端推理）

链路：前端摄像头抽帧 → JPEG base64 → POST /api/vision/frame → 本模块推理 → 返回人脸/姿态/举手特征。
模型：
  - backend/models/yolov8n-face.pt （lindevs/yolov8-face，WIDERFACE 训练，torch 6.3MB）
  - backend/models/yolov8n-pose.pt （ultralytics 官方，COCO 17 关键点，torch 6.8MB）
推理：ultralytics 加载（torch 原生，自动 letterbox + NMS），CPU 设备。

手势判定分工：
  - 后端只做"单帧几何特征"（举手 = 手腕高于肩），无状态
  - 前端做时序判定（举手持续 1.5s → 开始对话；举手后手腕摆动 → 挥手结束对话）
"""
import base64
import io
import logging
import time
from pathlib import Path

from PIL import Image

log = logging.getLogger("dpm.vision")

MODEL_DIR = Path(__file__).resolve().parent.parent / "models"
FACE_MODEL_PATH = MODEL_DIR / "yolov8n-face.pt"
POSE_MODEL_PATH = MODEL_DIR / "yolov8n-pose.pt"
CONF_THRESHOLD = 0.25   # 远距离小脸可调低（如 0.15）
IMGSZ = 640

# COCO 17 关键点索引
KP_LEFT_SHOULDER = 5
KP_RIGHT_SHOULDER = 6
KP_LEFT_ELBOW = 7
KP_RIGHT_ELBOW = 8
KP_LEFT_WRIST = 9
KP_RIGHT_WRIST = 10
RAISE_LIFT = 0.12   # 举手：手腕高于对应肩 12% 画面高
FIST_ELBOW_LIFT = 0.05   # 举拳：手腕高于肘 5% 画面高（前臂上举）
FIST_SHOULDER_GAP = 0.15 # 举拳：拳不高于肩 15% 画面高（收在胸前）
FIST_CHEST_DIST = 0.25   # 举拳：拳与肩水平距离 < 25% 画面宽（贴近躯干）
KP_CONF_MIN = 0.3   # 关键点置信度过滤

_face_model = None
_pose_model = None


def get_face_model():
    global _face_model
    if _face_model is None:
        from ultralytics import YOLO
        log.info("vision: 加载 YOLO 人脸模型 %s（首次约 2-5s）", FACE_MODEL_PATH)
        _face_model = YOLO(str(FACE_MODEL_PATH), task="detect")
        log.info("vision: YOLO 人脸模型就绪")
    return _face_model


def get_pose_model():
    global _pose_model
    if _pose_model is None:
        from ultralytics import YOLO
        log.info("vision: 加载 YOLO 姿态模型 %s（首次约 2-5s）", POSE_MODEL_PATH)
        _pose_model = YOLO(str(POSE_MODEL_PATH), task="pose")
        log.info("vision: YOLO 姿态模型就绪")
    return _pose_model


def _parse_pose(results):
    """解析姿态结果 → persons（17 关键点归一化 [x,y,conf]）+ raised（举手）+ fists（举拳）"""
    persons = []
    raised = []
    fists = []
    for r in results:
        if r.keypoints is None:
            continue
        kps = r.keypoints.data  # [N,17,3] 归一化坐标
        for i in range(kps.shape[0]):
            kp = kps[i]
            person = [
                [round(float(kp[j][0]), 3), round(float(kp[j][1]), 3), round(float(kp[j][2]), 3)]
                for j in range(17)
            ]
            persons.append(person)
            # 每只手臂：举手（腕明显高于肩）或 举拳（前臂上举、拳收胸前），两者互斥
            for side, wrist_i, elbow_i, shoulder_i in (
                ("left", KP_LEFT_WRIST, KP_LEFT_ELBOW, KP_LEFT_SHOULDER),
                ("right", KP_RIGHT_WRIST, KP_RIGHT_ELBOW, KP_RIGHT_SHOULDER),
            ):
                w = kp[wrist_i]
                e = kp[elbow_i]
                s = kp[shoulder_i]
                if float(w[2]) < KP_CONF_MIN or float(s[2]) < KP_CONF_MIN or float(e[2]) < KP_CONF_MIN:
                    continue
                wx, wy = float(w[0]), float(w[1])
                sx, sy = float(s[0]), float(s[1])
                ey = float(e[1])
                if wy < sy - RAISE_LIFT:
                    raised.append({"side": side, "x": round(wx, 3), "y": round(wy, 3)})
                elif (
                    wy < ey - FIST_ELBOW_LIFT      # 前臂上举（腕高于肘）
                    and wy >= sy - FIST_SHOULDER_GAP  # 拳不高过肩太多（收在胸前）
                    and abs(wx - sx) < FIST_CHEST_DIST  # 拳贴近躯干中线
                ):
                    fists.append({"side": side, "x": round(wx, 3), "y": round(wy, 3)})
    return persons, raised, fists


def predict_jpeg(jpeg_bytes: bytes):
    """人脸检测 + 姿态估计 → {faces, boxes, pose, raised, latency_ms}"""
    t0 = time.time()
    img = Image.open(io.BytesIO(jpeg_bytes)).convert("RGB")

    # 人脸
    face_res = get_face_model().predict(img, conf=CONF_THRESHOLD, imgsz=IMGSZ, verbose=False, device="cpu")
    boxes = []
    for r in face_res:
        if r.boxes is None:
            continue
        for b in r.boxes:
            xyxy = [round(float(v), 1) for v in b.xyxy[0].tolist()]
            conf = round(float(b.conf[0]), 3)
            boxes.append({"box": xyxy, "conf": conf})

    # 姿态
    pose_res = get_pose_model().predict(img, conf=CONF_THRESHOLD, imgsz=IMGSZ, verbose=False, device="cpu")
    persons, raised, fists = _parse_pose(pose_res)

    latency_ms = round((time.time() - t0) * 1000, 1)
    log.info(
        "vision: 人脸 %d 姿态 %d 举手 %d 举拳 %d · %.0fms",
        len(boxes), len(persons), len(raised), len(fists), latency_ms,
    )
    return {
        "faces": len(boxes),
        "boxes": boxes,
        "pose": persons,
        "raised": raised,
        "fists": fists,
        "latency_ms": latency_ms,
    }


def predict_base64(b64: str):
    """入口：base64 JPEG → 检测结果 dict（含 ok 标记，失败时带 error）"""
    try:
        jpeg = base64.b64decode(b64)
    except Exception as e:  # noqa: BLE001
        log.warning("vision: base64 解码失败 %s", e)
        return {"ok": False, "error": f"bad base64: {e}"}
    try:
        data = predict_jpeg(jpeg)
        data["ok"] = True
        return data
    except Exception as e:  # noqa: BLE001
        log.error("vision: 推理失败 %s", e, exc_info=True)
        return {"ok": False, "error": str(e)}