import cv2
import dlib
import numpy as np
import base64
from fastapi import FastAPI, WebSocket
from scipy.spatial import distance as dist
from imutils import face_utils
import uvicorn
import json

app = FastAPI()

import os

# --- CONFIGURATION FOR HUGGING FACE ---
# Try multiple paths for the model file
possible_paths = [
    "shape_predictor_68_face_landmarks.dat",
    "data/models/shape_predictor_68_face_landmarks.dat",
    "/app/shape_predictor_68_face_landmarks.dat",
    "/app/data/models/shape_predictor_68_face_landmarks.dat"
]

predictor_path = None
for path in possible_paths:
    if os.path.exists(path):
        predictor_path = path
        break

if not predictor_path:
    print("[ERROR] shape_predictor_68_face_landmarks.dat not found in any of the searched paths.")
    print(f"Current working directory: {os.getcwd()}")
    print(f"Files in current directory: {os.listdir('.')}")
else:
    print(f"[INFO] Using model at: {predictor_path}")

print("[INFO] Loading AI models...")
detector = dlib.get_frontal_face_detector()
try:
    if predictor_path:
        predictor = dlib.shape_predictor(predictor_path)
        print("[SUCCESS] Model loaded!")
    else:
        print("[ERROR] Predictor path is None, cannot load model.")
except Exception as e:
    print(f"[ERROR] Failed to load model: {e}")

# Landmark indexes
(lStart, lEnd) = face_utils.FACIAL_LANDMARKS_IDXS["left_eye"]
(rStart, rEnd) = face_utils.FACIAL_LANDMARKS_IDXS["right_eye"]
(mStart, mEnd) = face_utils.FACIAL_LANDMARKS_IDXS["mouth"]
(nStart, nEnd) = face_utils.FACIAL_LANDMARKS_IDXS["nose"]

def get_ear(eye):
    a = dist.euclidean(eye[1], eye[5])
    b = dist.euclidean(eye[2], eye[4])
    c = dist.euclidean(eye[0], eye[3])
    return (a + b) / (2.0 * c)

def get_mar(mouth):
    a = dist.euclidean(mouth[13], mouth[19])
    b = dist.euclidean(mouth[14], mouth[18])
    c = dist.euclidean(mouth[15], mouth[17])
    d = dist.euclidean(mouth[12], mouth[16])
    return (a + b + c) / (2.0 * d + 1e-6)

@app.websocket("/ws/analyze")
async def websocket_endpoint(websocket: WebSocket):
    await websocket.accept()
    counter = 0
    print("[INFO] Mobile App connected!")
    
    try:
        while True:
            data = await websocket.receive_text()
            header, encoded = data.split(",", 1)
            nparr = np.frombuffer(base64.b64decode(encoded), np.uint8)
            frame = cv2.imdecode(nparr, cv2.IMREAD_COLOR)

            if frame is None:
                continue

            gray = cv2.cvtColor(frame, cv2.COLOR_BGR2GRAY)
            faces = detector(gray, 0)
            
            response = {
                "drowsy": False,
                "distracted": False,
                "yawning": False,
                "ear": 0,
                "mar": 0
            }

            for rect in faces:
                shape = predictor(gray, rect)
                shape = face_utils.shape_to_np(shape)

                # Detection Logic
                ear = (get_ear(shape[lStart:lEnd]) + get_ear(shape[rStart:rEnd])) / 2.0
                mar = get_mar(shape[mStart:mEnd])
                
                nose_center = shape[nStart:nEnd].mean(axis=0)
                dist_left = dist.euclidean(shape[lStart:lEnd].mean(axis=0), nose_center)
                dist_right = dist.euclidean(shape[rStart:rEnd].mean(axis=0), nose_center)
                gaze_ratio = dist_left / (dist_right + 1e-6)

                response["ear"] = round(ear, 3)
                response["mar"] = round(mar, 3)
                
                if ear < 0.22: counter += 1
                else: counter = 0

                if counter >= 5: response["drowsy"] = True
                if gaze_ratio < 0.7 or gaze_ratio > 1.3: response["distracted"] = True
                if mar > 0.6: response["yawning"] = True

            await websocket.send_json(response)

    except Exception as e:
        print(f"[ERROR] Connection closed: {e}")

if __name__ == "__main__":
    # Port 7860 is required for Hugging Face Spaces
    uvicorn.run(app, host="0.0.0.0", port=7860)