import cv2
import mediapipe as mp
import time

# TODO
# add pyautogui for mouse control
# import pyautogui

BaseOptions = mp.tasks.BaseOptions
FaceLandmarker = mp.tasks.vision.FaceLandmarker
FaceLandmarkerOptions = mp.tasks.vision.FaceLandmarkerOptions
VisionRunningMode = mp.tasks.vision.RunningMode

class HeadMouseController:
    def __init__(self, model_path='face_landmarker.task'):
        options = FaceLandmarkerOptions(
            base_options=BaseOptions(model_asset_path=model_path),
            running_mode=VisionRunningMode.VIDEO,
            output_face_blendshapes=True,
            num_faces=1
        )
        self.landmarker = FaceLandmarker.create_from_options(options)

    def process_frame(self, frame, timestamp):

        rgb_frame = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
        mp_image = mp.Image(image_format=mp.ImageFormat.SRGB, data=rgb_frame)
        
        # Detect landmarks
        result = self.landmarker.detect_for_video(mp_image, timestamp)

        if result.face_landmarks:
            h, w, _ = frame.shape
            # draw all landmarks as small green dots
            for landmark in result.face_landmarks[0]:
                cx = min(max(int(landmark.x * w), 0), w - 1)
                cy = min(max(int(landmark.y * h), 0), h - 1)
                cv2.circle(frame, (cx, cy), 1, (0, 255, 0), -1)

            # highlight nose tip in blue
            nose = result.face_landmarks[0][1]
            nx = min(max(int(nose.x * w), 0), w - 1)
            ny = min(max(int(nose.y * h), 0), h - 1)
            cv2.circle(frame, (nx, ny), 4, (255, 0, 0), -1)

            # debug output for nose position (normalized)
            print(f"nose: {nose.x:.3f}, {nose.y:.3f}")

            # TODO
            # design a way to map facial landmarks to mouse movements and clicks
            # you can use the raw landmark positions or "blendshape" for more complex gestures

        return frame

    def start(self):
        cap = cv2.VideoCapture(0)
        while cap.isOpened():
            success, frame = cap.read()
            if not success: break
            
            timestamp = int(time.time() * 1000)
            frame = self.process_frame(frame, timestamp)
            
            cv2.imshow('face cursor', cv2.flip(frame, 1))
            if cv2.waitKey(1) & 0xFF == ord('q'):
                break
        cap.release()
        cv2.destroyAllWindows()

if __name__ == "__main__":
    controller = HeadMouseController()
    controller.start()
