# drive_paddy/detection/strategies/geometric.py
import cv2
import mediapipe as mp
import numpy as np
import math
from ..base_processor import BaseProcessor

# --- Helper Functions (No changes here) ---
def calculate_ear(eye_landmarks, frame_shape):
    coords = np.array([(lm.x * frame_shape[1], lm.y * frame_shape[0]) for lm in eye_landmarks])
    v1 = np.linalg.norm(coords[1] - coords[5]); v2 = np.linalg.norm(coords[2] - coords[4])
    h1 = np.linalg.norm(coords[0] - coords[3]); return (v1 + v2) / (2.0 * h1) if h1 > 0 else 0.0

def calculate_mar(mouth_landmarks, frame_shape):
    coords = np.array([(lm.x * frame_shape[1], lm.y * frame_shape[0]) for lm in mouth_landmarks])
    v1 = np.linalg.norm(coords[1] - coords[7]); v2 = np.linalg.norm(coords[2] - coords[6])
    v3 = np.linalg.norm(coords[3] - coords[5]); h1 = np.linalg.norm(coords[0] - coords[4])
    return (v1 + v2 + v3) / (2.0 * h1) if h1 > 0 else 0.0

class GeometricProcessor(BaseProcessor):
    def __init__(self, config):
        self.settings = config['geometric_settings']
        self.face_mesh = mp.solutions.face_mesh.FaceMesh(max_num_faces=1, refine_landmarks=True, min_detection_confidence=0.5, min_tracking_confidence=0.5)
        self.counters = { "eye_closure": 0, "yawning": 0, "head_nod": 0, "looking_away": 0 }
        self.L_EYE = [362, 385, 387, 263, 373, 380]; self.R_EYE = [33, 160, 158, 133, 153, 144]
        self.MOUTH = [61, 291, 39, 181, 0, 17, 84, 178]

    def process_frame(self, frame):
        h, w, _ = frame.shape
        gray = cv2.cvtColor(frame, cv2.COLOR_BGR2GRAY)
        brightness = np.mean(gray)
        is_low_light = brightness < self.settings['low_light_thresh']

        drowsiness_indicators = {
            "drowsiness_level": "Awake", "lighting": "Good", "details": {}
        }
        face_landmarks = None

        if not is_low_light:
            img_rgb = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
            results = self.face_mesh.process(img_rgb)
            face_landmarks = results.multi_face_landmarks

            if face_landmarks:
                landmarks = face_landmarks[0].landmark
                score = 0
                weights = self.settings['indicator_weights']

                # --- Draw Facial Landmarks (Logic Added Back) ---
                # This will draw the green dots for eyes and mouth.
                eye_mouth_landmarks = self.L_EYE + self.R_EYE + self.MOUTH
                for idx in eye_mouth_landmarks:
                    lm = landmarks[idx]
                    x, y = int(lm.x * w), int(lm.y * h)
                    cv2.circle(frame, (x, y), 1, (0, 255, 0), -1)

                # --- Drowsiness Calculations ---
                ear = (calculate_ear([landmarks[i] for i in self.L_EYE],(h,w)) + calculate_ear([landmarks[i] for i in self.R_EYE],(h,w)))/2.0
                if ear < self.settings['eye_ar_thresh']: self.counters['eye_closure']+=1
                else: self.counters['eye_closure']=0
                if self.counters['eye_closure'] >= self.settings['eye_ar_consec_frames']: score += weights['eye_closure']
                
                mar = calculate_mar([landmarks[i] for i in self.MOUTH], (h, w))
                if mar > self.settings['yawn_mar_thresh']: self.counters['yawning']+=1
                else: self.counters['yawning']=0
                if self.counters['yawning'] >= self.settings['yawn_consec_frames']: score += weights['yawning']

                face_3d = np.array([[0.0,0.0,0.0],[0.0,-330.0,-65.0],[-225.0,170.0,-135.0],[225.0,170.0,-135.0],[-150.0,-150.0,-125.0],[150.0,-150.0,-125.0]],dtype=np.float64)
                face_2d = np.array([(landmarks[1].x*w,landmarks[1].y*h),(landmarks[152].x*w,landmarks[152].y*h),(landmarks[263].x*w,landmarks[263].y*h),(landmarks[33].x*w,landmarks[33].y*h),(landmarks[287].x*w,landmarks[287].y*h),(landmarks[57].x*w,landmarks[57].y*h)],dtype=np.float64)
                cam_matrix = np.array([[w,0,w/2],[0,w,h/2],[0,0,1]],dtype=np.float64)
                _, rot_vec, _ = cv2.solvePnP(face_3d, face_2d, cam_matrix, np.zeros((4,1),dtype=np.float64))
                rmat, _ = cv2.Rodrigues(rot_vec); angles, _, _, _, _, _ = cv2.RQDecomp3x3(rmat)
                pitch, yaw = angles[0], angles[1]

                if pitch > self.settings['head_nod_thresh']: self.counters['head_nod']+=1
                else: self.counters['head_nod']=0
                if self.counters['head_nod'] >= self.settings['head_pose_consec_frames']: score += weights['head_nod']

                if abs(yaw) > self.settings['head_look_away_thresh']: self.counters['looking_away']+=1
                else: self.counters['looking_away']=0
                if self.counters['looking_away'] >= self.settings['head_pose_consec_frames']: score += weights['looking_away']
                
                levels = self.settings['drowsiness_levels']
                if score >= levels['very_drowsy_threshold']: drowsiness_indicators['drowsiness_level'] = "Very Drowsy"
                elif score >= levels['slightly_drowsy_threshold']: drowsiness_indicators['drowsiness_level'] = "Slightly Drowsy"
                
                drowsiness_indicators['details']['Score'] = score
        else: # is_low_light is True
            drowsiness_indicators["lighting"] = "Low"
        
        # --- Visualization on Video Frame ---
        level = drowsiness_indicators['drowsiness_level']
        score_val = drowsiness_indicators.get("details", {}).get("Score", 0)
        color = (0, 255, 0) # Green for Awake
        
        if drowsiness_indicators['lighting'] == "Low":
            color = (0, 165, 255) # Orange for low light
            cv2.putText(frame, "LOW LIGHT", (w // 2 - 120, h // 2), cv2.FONT_HERSHEY_SIMPLEX, 2, color, 3, cv2.LINE_AA)
        elif level == "Slightly Drowsy":
            color = (0, 255, 255) # Yellow
        elif level == "Very Drowsy":
            color = (0, 0, 255) # Red

        cv2.rectangle(frame, (0, 0), (w, h), color, 10)
        status_text = f"Status: {level} (Score: {score_val:.2f})"
        cv2.putText(frame, status_text, (20, 40), cv2.FONT_HERSHEY_SIMPLEX, 1, (255, 255, 255), 2, cv2.LINE_AA)

        return frame, drowsiness_indicators, face_landmarks