# -*- 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

from .config import settings

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

_frame_seq = 0

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.06   # 举手：手腕高于对应肩 6% 画面高（降低阈值，更易稳定激活）
FIST_ELBOW_LIFT = 0.05   # 举拳：手腕高于肘 5% 画面高（前臂上举）
FIST_SHOULDER_GAP = 0.15 # 举拳：拳不高于肩 15% 画面高（收在胸前）
FIST_CHEST_DIST = 0.25   # 举拳：拳与肩水平距离 < 25% 画面宽（贴近躯干）
KP_CONF_MIN = 0.3   # 关键点置信度过滤

# ── 扩展手势几何阈值（基于 COCO 17 关键点） ──
POINT_REACH_X = 0.30      # 指向：手腕水平伸出距肩 ≥30% 画面宽（手臂向前/侧伸）
POINT_Y_RANGE = 0.28      # 指向：手腕与肩同高 ±28% 画面高（排除高举/下垂）
HANDS_CLOSE_DIST = 0.18   # 双手合十/靠近：双腕欧氏距离 <18% 画面宽

_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, W, H):
    """解析姿态结果 → persons + raised(举手) + fists(举拳)
       + both_up(双臂举起) + pointing(指向) + hands_close(双手合十)
       + hands(每手几何摘要，供前端挥手/手势跟随时序判定)
       注意：ultralytics keypoints.data 为【原图像素坐标】，
       此处统一归一化为 [0,1]（x/W, y/H）后再判定阈值与输出
    """
    persons = []
    raised = []
    fists = []
    pointing = []
    hands = []
    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)
            # 每只手臂：举手（腕明显高于肩）或 举拳（前臂上举、拳收胸前），两者互斥
            arm = {}
            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
                # 像素 → 归一化 [0,1]
                wx, wy = float(w[0]) / W, float(w[1]) / H
                sx, sy = float(s[0]) / W, float(s[1]) / H
                ey = float(e[1]) / H
                is_raised = wy < sy - RAISE_LIFT
                is_fist = (
                    wy < ey - FIST_ELBOW_LIFT
                    and wy >= sy - FIST_SHOULDER_GAP
                    and abs(wx - sx) < FIST_CHEST_DIST
                )
                if is_raised:
                    raised.append({"side": side, "x": round(wx, 4), "y": round(wy, 4)})
                elif is_fist:
                    fists.append({"side": side, "x": round(wx, 4), "y": round(wy, 4)})
                # 指向：手腕水平伸出距肩较远、且与肩同高区间（前伸/侧伸，排除高举与下垂）
                if (
                    abs(wx - sx) > POINT_REACH_X
                    and abs(wy - sy) < POINT_Y_RANGE
                ):
                    pointing.append({"side": side, "x": round(wx, 4), "y": round(wy, 4)})
                arm[side] = {
                    "side": side,
                    "wx": round(wx, 4), "wy": round(wy, 4),
                    "sx": round(sx, 4), "sy": round(sy, 4),
                    "raised": is_raised, "fist": is_fist,
                }
            if arm:
                hands.append(arm)
    # 双臂举起：左、右腕都高于各自肩
    both_up = False
    if len(hands) >= 1:
        h0 = hands[0]
        if "left" in h0 and "right" in h0:
            both_up = h0["left"]["raised"] and h0["right"]["raised"]
    # 双手合十/靠近：同一人的双腕欧氏距离 < 阈值
    hands_close = False
    if len(hands) >= 1:
        h0 = hands[0]
        if "left" in h0 and "right" in h0:
            dx = h0["left"]["wx"] - h0["right"]["wx"]
            dy = h0["left"]["wy"] - h0["right"]["wy"]
            hands_close = (dx * dx + dy * dy) ** 0.5 < HANDS_CLOSE_DIST
    return persons, raised, fists, both_up, pointing, hands_close, hands


def predict_jpeg(jpeg_bytes: bytes):
    """人脸检测 + 姿态估计 → {faces, boxes, pose, raised, fists, both_up, pointing, hands_close, hands, 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})

    # 姿态
    W, H = img.size   # 原图尺寸（关键点像素坐标 → 归一化基准）
    pose_res = get_pose_model().predict(img, conf=CONF_THRESHOLD, imgsz=IMGSZ, verbose=False, device="cpu")
    persons, raised, fists, both_up, pointing, hands_close, hands = _parse_pose(pose_res, W, H)

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


def _should_save(data: dict) -> bool:
    """仅当画面检测到人脸/手势时才值得落盘保存。"""
    return (
        int(data.get("faces") or 0) > 0
        or bool(data.get("raised"))
        or bool(data.get("fists"))
        or bool(data.get("both_up"))
        or bool(data.get("pointing"))
        or bool(data.get("hands_close"))
    )


def maybe_save_frame(jpeg: bytes, data: dict):
    """仅将命中（有人/有动作）的 JPEG 帧落盘到 MEDIA_DIR/video/<小时目录>/，便于排查识别链路。
       目录按小时分割（YYYYMMDD_HH），不自动清理。
       可通过 DPM_SAVE_VISION_FRAMES=0 关闭，或 DPM_VISION_SAVE_DIR 改目录。"""
    global _frame_seq
    if not settings.SAVE_VISION_FRAMES:
        return
    if not _should_save(data):
        return
    _frame_seq += 1
    ts = time.strftime("%Y%m%d_%H%M%S")
    hour = time.strftime("%Y%m%d_%H")
    try:
        save_dir = settings.VISION_FRAME_SAVE_DIR / hour
        save_dir.mkdir(parents=True, exist_ok=True)
        path = save_dir / f"{ts}_{_frame_seq:05d}.jpg"
        with open(path, "wb") as f:
            f.write(jpeg)
    except Exception as e:  # noqa: BLE001
        log.warning("vision: 保存帧失败 %s", e)


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
        maybe_save_frame(jpeg, data)
        return data
    except Exception as e:  # noqa: BLE001
        log.error("vision: 推理失败 %s", e, exc_info=True)
        return {"ok": False, "error": str(e)}