from __future__ import annotations

import platform
from typing import Callable

import mediapipe as mp
from loguru import logger

from app.vision.pose_models import DEFAULT_MODEL_PATH

PoseLandmarker = mp.tasks.vision.PoseLandmarker
PoseLandmarkerOptions = mp.tasks.vision.PoseLandmarkerOptions
VisionRunningMode = mp.tasks.vision.RunningMode
BaseOptions = mp.tasks.BaseOptions

class PoseLandmarkerWrapper:
    """MediaPipe姿态关键点检测器封装"""

    def __init__(
        self,
        *,
        model_path: str | None = None,
        prefer_gpu: bool = True,
        result_callback: Callable | None = None,
    ) -> None:
        """初始化姿态检测器，优先尝试GPU委托，失败则回退到CPU"""
        self.model_path = model_path or DEFAULT_MODEL_PATH

        if prefer_gpu:
            try:
                if platform.system() == "Windows":
                    logger.warning(
                        "MediaPipe GPU delegate requested, but MediaPipe Tasks Python does not support GPU delegate on Windows; "
                        "Intel Iris Xe cannot be used by this backend and CPU fallback is expected"
                    )
                self.delegate = BaseOptions.Delegate.GPU
                self._landmarker = self._create(self.delegate, result_callback)
                logger.info("MediaPipe PoseLandmarker initialized with GPU delegate")
                return
            except Exception as exc:
                logger.warning("MediaPipe GPU delegate unavailable, falling back to CPU: {}", exc)

        self.delegate = BaseOptions.Delegate.CPU
        self._landmarker = self._create(self.delegate, result_callback)
        logger.info("MediaPipe PoseLandmarker initialized with CPU delegate")

    def _create(self, delegate, result_callback=None):
        """根据委托类型和回调创建PoseLandmarker实例"""
        options = PoseLandmarkerOptions(
            base_options=BaseOptions(model_asset_path=self.model_path, delegate=delegate),
            running_mode=VisionRunningMode.LIVE_STREAM,
            result_callback=result_callback,
            num_poses=1,
            min_pose_detection_confidence=0.5,
            min_pose_presence_confidence=0.5,
            min_tracking_confidence=0.5,
        )
        return PoseLandmarker.create_from_options(options)

    def detect_async(self, mp_image, timestamp_ms: int) -> None:
        """异步执行姿态检测"""
        return self._landmarker.detect_async(mp_image, timestamp_ms)

    def close(self) -> None:
        """释放MediaPipe资源"""
        self._landmarker.close()