"""
MediaPipe pose extraction wrapper.

MediaPipe provides 33 body keypoints, 21 hand keypoints per hand,
and 468 face landmarks. This wrapper provides a unified interface
for extracting these poses from videos.

Performance optimizations:
- Uses VIDEO running mode for temporal tracking (more efficient)
- Enables tracking by default for better consistency across frames
- Processes frames with timestamps for optimal VIDEO mode performance

Usage:
    from sltk.extraction.mediapipe import MediaPipeExtractor

    extractor = MediaPipeExtractor()
    extractor.load_model()
    result = extractor.extract_from_video("video.mp4")

    # Access poses
    body_poses = result.poses  # PoseSequence with body keypoints
"""

import hashlib
import logging
from dataclasses import dataclass
from pathlib import Path
from typing import Any

import numpy as np

logger = logging.getLogger(__name__)

# SHA256 checksums for MediaPipe model files
# These can be updated when new model versions are released
# To compute a checksum for a model file, use: compute_model_checksum(path)
MODEL_CHECKSUMS: dict[str, str | None] = {
    "pose_landmarker_heavy.task": None,  # Checksum not yet verified
    "pose_landmarker_full.task": None,
    "pose_landmarker_lite.task": None,
    "hand_landmarker.task": None,
    "holistic_landmarker.task": None,
}


def compute_model_checksum(file_path: Path) -> str:
    """
    Compute SHA256 checksum of a file.

    Useful for getting checksums of existing model files to add to MODEL_CHECKSUMS.

    Args:
        file_path: Path to the file to compute checksum for

    Returns:
        Hexadecimal SHA256 checksum string

    Example:
        >>> from sltk.extraction.mediapipe import compute_model_checksum
        >>> checksum = compute_model_checksum(Path.home() / ".cache/mediapipe/hand_landmarker.task")
        >>> print(f'"hand_landmarker.task": "{checksum}",')
    """
    sha256_hash = hashlib.sha256()
    with open(file_path, "rb") as f:
        for chunk in iter(lambda: f.read(8192), b""):
            sha256_hash.update(chunk)
    return sha256_hash.hexdigest()


def verify_checksum(file_path: Path, expected_hash: str) -> bool:
    """
    Verify SHA256 checksum of a file.

    Args:
        file_path: Path to the file to verify
        expected_hash: Expected SHA256 checksum (hexadecimal string)

    Returns:
        True if checksum matches, False otherwise
    """
    actual_hash = compute_model_checksum(file_path)
    return actual_hash == expected_hash


def download_with_verification(
    url: str,
    dest_path: Path,
    expected_hash: str | None = None,
) -> None:
    """
    Download a file and optionally verify its checksum.

    Downloads to a temporary file first, verifies checksum if provided,
    then moves to the final destination. This prevents corrupted/partial
    downloads from being used.

    Args:
        url: URL to download from
        dest_path: Destination path for the downloaded file
        expected_hash: Optional SHA256 checksum to verify against.
            If None, a warning is logged but download proceeds.

    Raises:
        ValueError: If checksum verification fails
        urllib.error.URLError: If download fails
    """
    import urllib.request

    # Download to temporary location first
    temp_path = dest_path.with_suffix(dest_path.suffix + ".tmp")

    try:
        urllib.request.urlretrieve(url, temp_path)

        # Verify checksum if provided
        if expected_hash is not None:
            if not verify_checksum(temp_path, expected_hash):
                actual_hash = compute_model_checksum(temp_path)
                temp_path.unlink()
                raise ValueError(
                    f"Checksum verification failed for {dest_path.name}. "
                    f"Expected: {expected_hash}, Got: {actual_hash}. "
                    "The downloaded file may be corrupted or tampered with."
                )
            logger.debug(f"Checksum verified for {dest_path.name}")
        else:
            # No checksum provided - log warning for security awareness
            logger.warning(
                f"No checksum available for {dest_path.name}. "
                "Download proceeding without verification. "
                "Consider computing and adding the checksum to MODEL_CHECKSUMS."
            )

        # Move to final location
        temp_path.rename(dest_path)

    except Exception:
        # Clean up temp file on any error
        if temp_path.exists():
            temp_path.unlink()
        raise


from sltk.data.core import PoseSequence
from sltk.extraction.base import (
    ExtractionConfig,
    ExtractionResult,
    PoseExtractor,
    get_video_info,
    read_video_frames,
)


@dataclass
class MediaPipeConfig(ExtractionConfig):
    """MediaPipe-specific configuration."""

    # Detection components to enable
    enable_body: bool = True
    enable_hands: bool = True
    enable_face: bool = False

    # Model complexity (0, 1, 2)
    model_complexity: int = 1

    # Detection confidence
    min_detection_confidence: float = 0.5
    min_tracking_confidence: float = 0.5

    # Static image mode (disable tracking) - False enables temporal tracking
    static_image_mode: bool = False

    # Use VIDEO running mode for new API (better performance with tracking)
    use_video_mode: bool = True


class MediaPipeExtractor(PoseExtractor):
    """
    MediaPipe pose extractor.

    Extracts body pose (33 keypoints), hand landmarks (21 per hand),
    and optionally face landmarks (468 points).

    Performance improvements:
    - Uses VIDEO running mode for the new task-based API
    - Enables temporal tracking for consistent detections across frames
    - Processes frames with timestamps for optimal tracking
    """

    def __init__(self, config: MediaPipeConfig | None = None):
        super().__init__(config or MediaPipeConfig())
        self.config: MediaPipeConfig = self.config
        self._holistic = None
        self._hands = None
        self._pose: Any = None
        self._use_legacy_api = True  # Will be set in load_model
        self._mp: Any = None
        self._vision: Any = None  # Store vision module reference

    @property
    def name(self) -> str:
        return "mediapipe"

    @property
    def output_format(self) -> str:
        return "mediapipe"

    def load_model(self, **kwargs) -> None:
        """Load MediaPipe models."""
        try:
            import mediapipe as mp
        except ImportError:
            raise ImportError("mediapipe required: pip install mediapipe")

        self._mp = mp
        self._use_legacy_api = hasattr(mp, "solutions")

        if self._use_legacy_api:
            # Legacy API (MediaPipe < 0.10)
            # Use static_image_mode=False for tracking (better performance)
            if self.config.enable_body and self.config.enable_hands:
                self._holistic = mp.solutions.holistic.Holistic(
                    static_image_mode=self.config.static_image_mode,
                    model_complexity=self.config.model_complexity,
                    min_detection_confidence=self.config.min_detection_confidence,
                    min_tracking_confidence=self.config.min_tracking_confidence,
                )
            elif self.config.enable_body:
                self._pose = mp.solutions.pose.Pose(
                    static_image_mode=self.config.static_image_mode,
                    model_complexity=self.config.model_complexity,
                    min_detection_confidence=self.config.min_detection_confidence,
                    min_tracking_confidence=self.config.min_tracking_confidence,
                )
            elif self.config.enable_hands:
                self._hands = mp.solutions.hands.Hands(
                    static_image_mode=self.config.static_image_mode,
                    max_num_hands=2,
                    model_complexity=self.config.model_complexity,
                    min_detection_confidence=self.config.min_detection_confidence,
                    min_tracking_confidence=self.config.min_tracking_confidence,
                )
        else:
            # New task-based API (MediaPipe >= 0.10)
            from mediapipe.tasks import python
            from mediapipe.tasks.python import vision

            self._vision = vision

            model_dir = Path.home() / ".cache" / "mediapipe"
            model_dir.mkdir(parents=True, exist_ok=True)

            # Determine running mode - VIDEO for sequential processing with tracking
            running_mode = (
                vision.RunningMode.VIDEO if self.config.use_video_mode else vision.RunningMode.IMAGE
            )

            if self.config.enable_body:
                pose_model_path = model_dir / "pose_landmarker_heavy.task"
                if not pose_model_path.exists():
                    logger.info("Downloading pose landmarker model...")
                    download_with_verification(
                        "https://storage.googleapis.com/mediapipe-models/pose_landmarker/pose_landmarker_heavy/float16/latest/pose_landmarker_heavy.task",
                        pose_model_path,
                        MODEL_CHECKSUMS.get("pose_landmarker_heavy.task"),
                    )

                base_options = python.BaseOptions(model_asset_path=str(pose_model_path))
                options = vision.PoseLandmarkerOptions(
                    base_options=base_options,
                    running_mode=running_mode,
                    min_pose_detection_confidence=self.config.min_detection_confidence,
                    min_tracking_confidence=self.config.min_tracking_confidence,
                )
                self._pose = vision.PoseLandmarker.create_from_options(options)

            if self.config.enable_hands:
                hand_model_path = model_dir / "hand_landmarker.task"
                if not hand_model_path.exists():
                    logger.info("Downloading hand landmarker model...")
                    download_with_verification(
                        "https://storage.googleapis.com/mediapipe-models/hand_landmarker/hand_landmarker/float16/latest/hand_landmarker.task",
                        hand_model_path,
                        MODEL_CHECKSUMS.get("hand_landmarker.task"),
                    )

                base_options = python.BaseOptions(model_asset_path=str(hand_model_path))
                options = vision.HandLandmarkerOptions(
                    base_options=base_options,
                    running_mode=running_mode,
                    num_hands=2,
                    min_hand_detection_confidence=self.config.min_detection_confidence,
                    min_tracking_confidence=self.config.min_tracking_confidence,
                )
                self._hands = vision.HandLandmarker.create_from_options(options)

        self._initialized = True

    def _process_frame_legacy(self, frame: np.ndarray) -> dict[str, Any]:
        """Process a single frame using legacy API."""
        result: dict[str, Any] = {
            "body": None,
            "left_hand": None,
            "right_hand": None,
            "face": None,
        }

        if self._holistic is not None:
            detection = self._holistic.process(frame)

            if detection.pose_landmarks:
                result["body"] = self._landmarks_to_array(detection.pose_landmarks.landmark, 33)

            if detection.left_hand_landmarks:
                result["left_hand"] = self._landmarks_to_array(
                    detection.left_hand_landmarks.landmark, 21
                )

            if detection.right_hand_landmarks:
                result["right_hand"] = self._landmarks_to_array(
                    detection.right_hand_landmarks.landmark, 21
                )

            if self.config.enable_face and detection.face_landmarks:
                result["face"] = self._landmarks_to_array(detection.face_landmarks.landmark, 468)

        elif self._pose is not None:
            detection = self._pose.process(frame)
            if detection.pose_landmarks:
                result["body"] = self._landmarks_to_array(detection.pose_landmarks.landmark, 33)

        elif self._hands is not None:
            detection = self._hands.process(frame)
            if detection.multi_hand_landmarks:
                for hand_landmarks, handedness in zip(
                    detection.multi_hand_landmarks, detection.multi_handedness
                ):
                    hand_type = handedness.classification[0].label.lower()
                    key = f"{hand_type}_hand"
                    result[key] = self._landmarks_to_array(hand_landmarks.landmark, 21)

        return result

    def _process_frame_video_mode(self, frame: np.ndarray, timestamp_ms: int) -> dict[str, Any]:
        """Process a single frame using new task API with VIDEO mode."""
        result: dict[str, Any] = {
            "body": None,
            "left_hand": None,
            "right_hand": None,
            "face": None,
        }

        # Convert frame to MediaPipe Image
        mp_image = self._mp.Image(image_format=self._mp.ImageFormat.SRGB, data=frame)

        if self._pose is not None:
            detection = self._pose.detect_for_video(mp_image, timestamp_ms)
            if detection.pose_landmarks and len(detection.pose_landmarks) > 0:
                result["body"] = self._landmarks_to_array_new_api(detection.pose_landmarks[0], 33)

        if self._hands is not None:
            detection = self._hands.detect_for_video(mp_image, timestamp_ms)
            if detection.hand_landmarks:
                for i, hand_landmarks in enumerate(detection.hand_landmarks):
                    # Determine handedness
                    if detection.handedness and i < len(detection.handedness):
                        hand_type = detection.handedness[i][0].category_name.lower()
                    else:
                        hand_type = "right" if i == 0 else "left"
                    key = f"{hand_type}_hand"
                    result[key] = self._landmarks_to_array_new_api(hand_landmarks, 21)

        return result

    def _process_frame_image_mode(self, frame: np.ndarray) -> dict[str, Any]:
        """Process a single frame using new task API with IMAGE mode."""
        result: dict[str, Any] = {
            "body": None,
            "left_hand": None,
            "right_hand": None,
            "face": None,
        }

        # Convert frame to MediaPipe Image
        mp_image = self._mp.Image(image_format=self._mp.ImageFormat.SRGB, data=frame)

        if self._pose is not None:
            detection = self._pose.detect(mp_image)
            if detection.pose_landmarks and len(detection.pose_landmarks) > 0:
                result["body"] = self._landmarks_to_array_new_api(detection.pose_landmarks[0], 33)

        if self._hands is not None:
            detection = self._hands.detect(mp_image)
            if detection.hand_landmarks:
                for i, hand_landmarks in enumerate(detection.hand_landmarks):
                    # Determine handedness
                    if detection.handedness and i < len(detection.handedness):
                        hand_type = detection.handedness[i][0].category_name.lower()
                    else:
                        hand_type = "right" if i == 0 else "left"
                    key = f"{hand_type}_hand"
                    result[key] = self._landmarks_to_array_new_api(hand_landmarks, 21)

        return result

    def _process_frame(self, frame: np.ndarray, timestamp_ms: int = 0) -> dict[str, Any]:
        """Process a single frame."""
        if self._use_legacy_api:
            return self._process_frame_legacy(frame)
        elif self.config.use_video_mode:
            return self._process_frame_video_mode(frame, timestamp_ms)
        else:
            return self._process_frame_image_mode(frame)

    def _landmarks_to_array(self, landmarks, num_landmarks: int) -> np.ndarray:
        """Convert MediaPipe landmarks to numpy array (legacy API)."""
        arr = np.zeros((num_landmarks, 4), dtype=np.float32)

        for i, lm in enumerate(landmarks):
            if i >= num_landmarks:
                break
            arr[i] = [lm.x, lm.y, lm.z, lm.visibility if hasattr(lm, "visibility") else 1.0]

        return arr

    def _landmarks_to_array_new_api(self, landmarks, num_landmarks: int) -> np.ndarray:
        """Convert MediaPipe landmarks to numpy array (new task API)."""
        arr = np.zeros((num_landmarks, 4), dtype=np.float32)

        for i, lm in enumerate(landmarks):
            if i >= num_landmarks:
                break
            # New API uses .x, .y, .z attributes directly
            visibility = lm.visibility if hasattr(lm, "visibility") else 1.0
            if visibility is None:
                visibility = 1.0
            arr[i] = [lm.x, lm.y, lm.z, visibility]

        return arr

    def extract_from_video(
        self,
        video_path: str | Path,
        output_path: str | Path | None = None,
    ) -> ExtractionResult:
        """Extract poses from video."""
        if not self._initialized:
            self.load_model()

        video_path = Path(video_path)
        info = get_video_info(video_path)
        fps = info["fps"]

        # Read and process frames
        frames = read_video_frames(video_path)
        result = self.extract_from_frames(frames, fps=fps)
        result.metadata["source_video"] = str(video_path)
        result.metadata["video_info"] = info

        # Save if output path provided
        if output_path is not None:
            self.save_results(result, output_path)

        return result

    def extract_from_frames(
        self,
        frames: np.ndarray,
        fps: float = 25.0,
    ) -> ExtractionResult:
        """
        Extract poses from numpy frames.

        Uses VIDEO mode with timestamps for optimal performance when tracking
        is enabled (default).
        """
        if not self._initialized:
            self.load_model()

        num_frames = len(frames)

        # Initialize storage
        body_data = []
        left_hand_data = []
        right_hand_data = []
        body_conf = []
        left_hand_conf = []
        right_hand_conf = []

        # Calculate frame duration in milliseconds for timestamp computation
        frame_duration_ms = 1000.0 / fps

        for i, frame in enumerate(frames):
            if self.config.verbose and i % 100 == 0:
                logger.info(f"  Frame {i}/{num_frames}")

            # Compute timestamp for VIDEO mode
            timestamp_ms = int(i * frame_duration_ms)
            detection = self._process_frame(frame, timestamp_ms)

            # Body
            if detection["body"] is not None:
                body_data.append(detection["body"][:, :3])
                body_conf.append(detection["body"][:, 3])
            else:
                body_data.append(np.zeros((33, 3), dtype=np.float32))
                body_conf.append(np.zeros(33, dtype=np.float32))

            # Left hand
            if detection["left_hand"] is not None:
                left_hand_data.append(detection["left_hand"][:, :3])
                left_hand_conf.append(detection["left_hand"][:, 3])
            else:
                left_hand_data.append(np.zeros((21, 3), dtype=np.float32))
                left_hand_conf.append(np.zeros(21, dtype=np.float32))

            # Right hand
            if detection["right_hand"] is not None:
                right_hand_data.append(detection["right_hand"][:, :3])
                right_hand_conf.append(detection["right_hand"][:, 3])
            else:
                right_hand_data.append(np.zeros((21, 3), dtype=np.float32))
                right_hand_conf.append(np.zeros(21, dtype=np.float32))

        # Convert to arrays
        body_arr = np.array(body_data, dtype=np.float32)
        left_hand_arr = np.array(left_hand_data, dtype=np.float32)
        right_hand_arr = np.array(right_hand_data, dtype=np.float32)

        # Combine body and hands (total: 33 + 21 + 21 = 75 keypoints)
        combined_data = np.concatenate([body_arr, left_hand_arr, right_hand_arr], axis=1)
        combined_conf = np.concatenate(
            [np.array(body_conf), np.array(left_hand_conf), np.array(right_hand_conf)], axis=1
        )

        # Create PoseSequence
        poses = PoseSequence(
            data=combined_data,
            fps=fps,
            format="mediapipe",
            confidence=combined_conf,
            metadata={
                "num_body_keypoints": 33,
                "num_hand_keypoints": 21,
                "keypoint_order": ["body", "left_hand", "right_hand"],
                "extractor": self.name,
                "video_mode": self.config.use_video_mode and not self._use_legacy_api,
            },
        )

        return ExtractionResult(
            poses=poses,
            num_frames=num_frames,
            num_detections=num_frames,  # One detection per frame
            fps=fps,
            metadata=poses.metadata,
        )

    def close(self):
        """Release resources."""
        if self._holistic is not None:
            self._holistic.close()
            self._holistic = None
        if self._pose is not None:
            if hasattr(self._pose, "close"):
                self._pose.close()
            self._pose = None
        if self._hands is not None:
            if hasattr(self._hands, "close"):
                self._hands.close()
            self._hands = None
        self._initialized = False

    def __enter__(self):
        """Context manager entry."""
        return self

    def __exit__(self, exc_type, exc_val, exc_tb):
        """Context manager exit - release resources."""
        self.close()
        return False


# Convenience function
def extract_mediapipe(
    video_path: str | Path,
    output_path: str | Path | None = None,
    enable_hands: bool = True,
    enable_face: bool = False,
) -> ExtractionResult:
    """
    Quick extraction using MediaPipe.

    Args:
        video_path: Input video path
        output_path: Optional output H5 path
        enable_hands: Whether to extract hand landmarks
        enable_face: Whether to extract face landmarks

    Returns:
        ExtractionResult with poses
    """
    config = MediaPipeConfig(
        enable_hands=enable_hands,
        enable_face=enable_face,
    )
    extractor = MediaPipeExtractor(config)
    try:
        return extractor.extract_from_video(video_path, output_path)
    finally:
        extractor.close()
