"""Hand tracking module using MediaPipe and a trained ML model."""

import os
import pickle
import cv2
import mediapipe as mp
import numpy as np


class HandTracker:
    """Recognizes hand gestures using MediaPipe and a pre-trained model."""

    # ASL letter labels
    LABELS_DICT = {
        0: 'Hi', 1: 'A', 2: 'B', 3: 'D', 4: 'E', 5: 'F',
        6: 'G', 7: 'H', 8: 'I', 9: 'J', 10: 'K', 11: 'L',
        12: 'M', 13: 'N', 14: 'O', 15: 'P', 16: 'Q', 17: 'R',
        18: 'S', 19: 'T', 20: 'U', 21: 'V', 22: 'W', 23: 'X', 24: 'Y'
    }

    def __init__(self, model_path=None):
        """
        Initialize the HandTracker.

        Args:
            model_path: Path to the pre-trained model pickle file.
                       If None, looks for 'model.p' in the package directory.
        """
        if model_path is None:
            # Look for model.p in the package directory
            base_dir = os.path.dirname(os.path.abspath(__file__))
            model_path = os.path.join(base_dir, 'model.p')

        if not os.path.exists(model_path):
            raise FileNotFoundError(f"Model file not found at {model_path}")

        model_dict = pickle.load(open(model_path, 'rb'))
        self.model = model_dict['model']

        # Initialize MediaPipe
        self.mp_hands = mp.solutions.hands
        self.mp_drawing = mp.solutions.drawing_utils
        self.mp_drawing_styles = mp.solutions.drawing_styles
        self.hands = self.mp_hands.Hands(
            static_image_mode=True,
            min_detection_confidence=0.3
        )

    def predict_gesture(self, frame):
        """
        Predict hand gesture from a video frame.

        Args:
            frame: A video frame (numpy array).

        Returns:
            dict: Contains:
                - 'character': Predicted character/gesture (str or None)
                - 'landmarks': Hand landmarks for the frame
                - 'bbox': Bounding box coordinates (x1, y1, x2, y2) or None
        """
        data_aux = []
        x_ = []
        y_ = []

        H, W, _ = frame.shape
        frame_rgb = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
        results = self.hands.process(frame_rgb)

        result = {
            'character': None,
            'landmarks': results.multi_hand_landmarks,
            'bbox': None
        }

        if results.multi_hand_landmarks:
            for hand_landmarks in results.multi_hand_landmarks:
                for i in range(len(hand_landmarks.landmark)):
                    x = hand_landmarks.landmark[i].x
                    y = hand_landmarks.landmark[i].y
                    x_.append(x)
                    y_.append(y)

                for i in range(len(hand_landmarks.landmark)):
                    x = hand_landmarks.landmark[i].x
                    y = hand_landmarks.landmark[i].y
                    data_aux.append(x - min(x_))
                    data_aux.append(y - min(y_))

            x1 = int(min(x_) * W) - 10
            y1 = int(min(y_) * H) - 10
            x2 = int(max(x_) * W) + 10
            y2 = int(max(y_) * H) + 10

            prediction = self.model.predict([np.asarray(data_aux)])
            predicted_character = self.LABELS_DICT[int(prediction[0])]

            result['character'] = predicted_character
            result['bbox'] = (x1, y1, x2, y2)

        return result

    def draw_predictions(self, frame, prediction_result):
        """
        Draw hand landmarks and prediction on the frame.

        Args:
            frame: Video frame to draw on.
            prediction_result: Dict from predict_gesture().

        Returns:
            frame: Annotated frame.
        """
        if prediction_result['landmarks']:
            for hand_landmarks in prediction_result['landmarks']:
                self.mp_drawing.draw_landmarks(
                    frame,
                    hand_landmarks,
                    self.mp_hands.HAND_CONNECTIONS,
                    self.mp_drawing_styles.get_default_hand_landmarks_style(),
                    self.mp_drawing_styles.get_default_hand_connections_style()
                )

        if prediction_result['bbox']:
            x1, y1, x2, y2 = prediction_result['bbox']
            cv2.rectangle(frame, (x1, y1), (x2, y2), (0, 0, 0), 4)
            if prediction_result['character']:
                cv2.putText(
                    frame,
                    prediction_result['character'],
                    (x1, y1 - 10),
                    cv2.FONT_HERSHEY_SIMPLEX,
                    1.3,
                    (255, 255, 255),
                    3,
                    cv2.LINE_AA
                )

        return frame


def run_hand_tracking(video_source=0, model_path=None):
    """
    Run real-time hand gesture recognition.

    Args:
        video_source: Camera index (default: 0 for default camera).
        model_path: Path to the pre-trained model pickle file.
    """
    tracker = HandTracker(model_path=model_path)
    cap = cv2.VideoCapture(video_source)

    if not cap.isOpened():
        raise RuntimeError("Failed to open video source")

    while True:
        ret, frame = cap.read()

        if not ret:
            print("Failed to grab frame")
            break

        prediction = tracker.predict_gesture(frame)
        frame = tracker.draw_predictions(frame, prediction)

        cv2.imshow('Hand Tracking', frame)

        key = cv2.waitKey(1) & 0xFF
        if key == ord('q'):
            break

    cap.release()
    cv2.destroyAllWindows()
