import cv2
import streamlit as st
import numpy as np
from drowsiness_detection import VideoFrameHandler
from streamlit_webrtc import webrtc_streamer, VideoProcessorBase, WebRtcMode

# Streamlit app UI
st.title("Drowsiness Detection System")

# Thresholds for EAR and drowsiness detection
thresholds = {
    "EAR_THRESH": 0.25,
    "WAIT_TIME": 2.0,  
}

class VideoProcessor(VideoProcessorBase):
    def __init__(self):
        self.frame_handler = VideoFrameHandler()
        self.play_alarm_flag = False

    def recv(self, frame):
        img = frame.to_ndarray(format="bgr24")
        processed_frame, self.play_alarm_flag = self.frame_handler.process(img, thresholds)
        return av.VideoFrame.from_ndarray(processed_frame, format="bgr24")

webrtc_ctx = webrtc_streamer(
    key="drowsiness-detection",
    mode=WebRtcMode.SENDRECV,
    video_processor_factory=VideoProcessor,
    media_stream_constraints={"video": True, "audio": False},
    async_processing=True,
)

if webrtc_ctx.video_processor:
    if webrtc_ctx.video_processor.play_alarm_flag:
        st.write(
            """
            <script>
            var audio = new Audio('audio/wake_up.wav');
            audio.play();
            </script>
            """,
            unsafe_allow_html=True,
        )
    webrtc_ctx.video_processor.thresholds = thresholds

st.write("Stopped Drowsiness Detection.")