from __future__ import annotations

import threading
import time
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:
    def __init__(
        self,
        *,
        model_path: str | None = None,
        prefer_gpu: bool = True,
        result_callback: Callable | None = None,
    ) -> None:
        self.model_path = model_path or DEFAULT_MODEL_PATH

        if prefer_gpu:
            try:
                self.delegate = BaseOptions.Delegate.GPU
                self._landmarker = self._create(PoseLandmarker.Delegate.GPU)
                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(PoseLandmarker.Delegate.CPU, result_callback)
        logger.info("MediaPipe PoseLandmarker initialized with CPU delegate")

    def _create(self, delegate, result_callback=None):
        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:
        self._landmarker.close()