import mediapipe as mp
import cv2, numpy as np, time, fitz
from helper import compute_distance
from config import PINCH_THRESHOLD, RELEASE_THRESHOLD

# hand_state.py
from dataclasses import dataclass, field
from typing import Tuple, Dict, Optional

from helper import finger_straight

# ----- Initialize Hand Detector -----
mp_hands = mp.solutions.hands
mp_draw  = mp.solutions.drawing_utils
hands = mp_hands.Hands(
    static_image_mode=False,
    max_num_hands=8,
    min_detection_confidence=0.7,
    min_tracking_confidence=0.5
)

start_pt = None               # measurement start point
measuring = False             # measurement in progress

@dataclass
class HandState:
    id: int                              # a persistent identifier for this hand
    landmarks: Dict[str, Tuple[int,int]] = field(default_factory=dict)
    finger_straightness: Dict[str, float] = field(default_factory=dict)
    gesture: Optional[str] = None        # e.g. "pinch", "fist", "open"
    gesture_persistence: int = 0         # how many frames the current gesture has held

#handList = Dict[int, HandState] = {}
next_id = 0 

def update(mp_results, frame_vis):
    h, w, _ = frame_vis.shape
    detected_centroids = []
    landmarks_list     = []

    # 1) pull out centroids & raw landmarks
    if mp_results.multi_hand_landmarks:
        for hand in mp_results.multi_hand_landmarks:
            pts = []
            for lm in hand.landmark:
                pts.append((int(lm.x*w), int(lm.y*h)))
            centroid = np.mean(pts, axis=0)
            detected_centroids.append(tuple(centroid.astype(int)))
            landmarks_list.append((hand, pts))

    # 2) match to existing by nearest centroid
    new_hands = {}
    used_ids  = set()
    for (hand, pts), centroid in zip(landmarks_list, detected_centroids):
        # find best existing hand
        best_id, best_dist = None, 1e9
        for hid, state in handList.items():
            dx, dy = np.array(state.landmarks['centroid']) - centroid
            d = np.hypot(dx, dy)
            if d < best_dist and d < 100:  # 100px max match distance
                best_dist, best_id = d, hid

        if best_id is None:
            hid = next_id
            next_id += 1
            state = HandState(id=hid)
        else:
            hid   = best_id
            state = handList[hid]

        used_ids.add(hid)

        # 3) update state
        state.persistence += 1
        state.landmarks['centroid'] = centroid

        # compute fingertip positions
        idx_tip = pts[mp.solutions.hands.HandLandmark.INDEX_FINGER_TIP]
        mid_tip = pts[mp.solutions.hands.HandLandmark.MIDDLE_FINGER_TIP]
        state.landmarks['index_tip']  = idx_tip
        state.landmarks['middle_tip'] = mid_tip

        for name, tip_i, pip_i, mcp_i in [
            ('index', mp.solutions.hands.HandLandmark.INDEX_FINGER_TIP,
                        mp.solutions.hands.HandLandmark.INDEX_FINGER_PIP,
                        mp.solutions.hands.HandLandmark.INDEX_FINGER_MCP),
            ('middle', mp.solutions.hands.HandLandmark.MIDDLE_FINGER_TIP,
                        mp.solutions.hands.HandLandmark.MIDDLE_FINGER_PIP,
                        mp.solutions.hands.HandLandmark.MIDDLE_FINGER_MCP),
            # add ring, pinky, thumb similarly...
        ]:
            p_tip = pts[tip_i]
            p_pip = pts[pip_i]
            p_mcp = pts[mcp_i]
            state.finger_straightness[name] = finger_straight(p_tip, p_pip, p_mcp)

        # 5) simple gesture detection
        if state.finger_straightness['index'] < 0.2 and \
            state.finger_straightness['middle'] < 0.2:
            gesture = 'fist'
        elif state.finger_straightness['index'] > 0.8 and \
                state.finger_straightness['middle'] > 0.8:
            gesture = 'open'
        else:
            gesture = None

        if gesture == state.gesture:
            state.gesture_persistence += 1
        else:
            state.gesture = gesture
            state.gesture_persistence = 0

        new_hands[hid] = state

    # 6) drop hands not seen this frame
    handList = new_hands
    return list(hands.values())



def detect_hands(frame):
    rgb = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
    rgb.flags.writeable = False
    results = hands.process(rgb)    

    return results


def check_hand(hand, frame_vis, object_calibrating):
    h, w, _ = frame_vis.shape

    fingertip_idx = None
    fingertip_mid = None
    
    mp_draw.draw_landmarks(frame_vis, hand, mp_hands.HAND_CONNECTIONS)
    # get index and middle finger tips and PIP joints
    idx_tip = hand.landmark[mp_hands.HandLandmark.INDEX_FINGER_TIP]
    mid_tip = hand.landmark[mp_hands.HandLandmark.MIDDLE_FINGER_TIP]
    ix, iy = int(idx_tip.x * w), int(idx_tip.y * h)
    mx, my = int(mid_tip.x * w), int(mid_tip.y * h)
    fingertip_idx = (ix, iy)
    fingertip_mid = (mx, my)

    # draw fingertips
    cv2.circle(frame_vis, fingertip_idx, 8, (255,255,0), -1)
    cv2.circle(frame_vis, fingertip_mid,  8, (0,255,0), -1)

    measuring, start_pt = is_measuring(frame_vis, hand, object_calibrating)

    return fingertip_idx, measuring, start_pt



def is_measuring(frame_vis, hand, object_calibrating):
    global measuring, start_pt
    h, w, _ = frame_vis.shape

    idx_tip = hand.landmark[mp_hands.HandLandmark.INDEX_FINGER_TIP]
    idx_pip = hand.landmark[mp_hands.HandLandmark.INDEX_FINGER_PIP]
    mid_tip = hand.landmark[mp_hands.HandLandmark.MIDDLE_FINGER_TIP]
    mid_pip = hand.landmark[mp_hands.HandLandmark.MIDDLE_FINGER_PIP]

    ix, iy = int(idx_tip.x * w), int(idx_tip.y * h)
    mx, my = int(mid_tip.x * w), int(mid_tip.y * h)
    fingertip_idx = (ix, iy)
    fingertip_mid = (mx, my)

    # check extension
    index_ext = idx_tip.y < idx_pip.y
    middle_ext = mid_tip.y < mid_pip.y

    # if measurement in progress but fingers no longer both extended, stop measuring
    if measuring and not (index_ext and middle_ext):
        measuring = False

    # pinch distance
    pinch = compute_distance(fingertip_idx, fingertip_mid)
    cv2.putText(frame_vis, f"Pinch: {int(pinch)} px", (10,30), cv2.FONT_HERSHEY_SIMPLEX, 0.6, (0,255,0),2)

    # measurement gesture
    if not object_calibrating and index_ext and middle_ext:
        if pinch < PINCH_THRESHOLD and not measuring:
            measuring = True
            start_pt = fingertip_idx
        elif pinch > RELEASE_THRESHOLD and measuring:
            measuring = False

    return measuring, start_pt