import cv2
import streamlit as st
from face_detection import FaceDetector
from mark_detection import MarkDetector
from pose_estimation import PoseEstimator
from utils import refine
import tempfile

def main():
    # Streamlit Title and Sidebar for inputs
    st.title("Distraction Detection App on Hugging Face Spaces")
    st.sidebar.write("Please upload a video file for analysis.")
    
    # File uploader for video
    video_file = st.sidebar.file_uploader("Upload a Video File", type=["mp4", "avi", "mov"])
    
    if video_file is not None:
        # Create a temporary file to store the uploaded video
        with tempfile.NamedTemporaryFile(delete=False) as temp_file:
            temp_file.write(video_file.read())
            video_path = temp_file.name
        
        # Set up video capture and detectors
        cap = cv2.VideoCapture(video_path)
        frame_width = int(cap.get(cv2.CAP_PROP_FRAME_WIDTH))
        frame_height = int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT))
        
        face_detector = FaceDetector("assets/face_detector.onnx")
        mark_detector = MarkDetector("assets/face_landmarks.onnx")
        pose_estimator = PoseEstimator(frame_width, frame_height)

        # Streamlit placeholder for processed frames
        frame_placeholder = st.empty()

        while cap.isOpened():
            # Capture each frame
            frame_got, frame = cap.read()
            if not frame_got:
                break

            # Face detection and pose estimation
            faces, _ = face_detector.detect(frame, 0.7)
            if len(faces) > 0:
                face = refine(faces, frame_width, frame_height, 0.15)[0]
                x1, y1, x2, y2 = face[:4].astype(int)
                patch = frame[y1:y2, x1:x2]
                marks = mark_detector.detect([patch])[0].reshape([68, 2])
                marks *= (x2 - x1)
                marks[:, 0] += x1
                marks[:, 1] += y1
                
                distraction_status, pose_vectors = pose_estimator.detect_distraction(marks)
                status_text = "Distracted" if distraction_status else "Focused"
                
                # Overlay status text
                cv2.putText(frame, f"Status: {status_text}", (10, 50),
                            cv2.FONT_HERSHEY_SIMPLEX, 0.5, 
                            (0, 255, 0) if not distraction_status else (0, 0, 255))
                
                # Display the processed frame
                frame_placeholder.image(cv2.cvtColor(frame, cv2.COLOR_BGR2RGB), channels="RGB")
        
        cap.release()
    else:
        st.warning("Please upload a video file to proceed.")

if __name__ == "__main__":
    main()