import streamlit as st
from PIL import Image
import os
from utils.yolo_processor import YOLOProcessor
import tempfile
import numpy as np 
import base64


processed_image = None
processed_video_path = None

def detect_fall(image, model_path):
    model = YOLOProcessor(model_path)
    result_image = model.detect_fall(image)
    return result_image

def main():
    global processed_image, processed_video_path

    st.title("Fall Detection with YOLO")
    st.markdown("---")
    option = st.sidebar.selectbox("Choose an option", ["Image", "Video"])

    if option == "Image":
        st.subheader("Upload Image")
        uploaded_file = st.file_uploader("Choose an image", type=["jpg", "jpeg", "png"])

        if uploaded_file is not None:
            image = Image.open(uploaded_file)
            st.image(image, caption='Uploaded Image', use_column_width=True)
            st.markdown("---")
            st.subheader("Detecting Fall...")

            if processed_image is None:  # Process the image only if it hasn't been processed before
                with st.spinner('Detecting fall...'):
                    processed_image = detect_fall(image, "assets/best.pt")
            st.image(processed_image, caption='Result', use_column_width=True)

            # Download button for the result image
            if st.button('Download Result Image'):
                download_image(processed_image, filename='result_image.png')

    elif option == "Video":
        st.subheader("Upload Video")
        uploaded_file = st.file_uploader("Choose a video", type=["mp4"])

        if uploaded_file is not None:
            st.markdown("---")
            st.subheader("Processing and Detecting Fall...")

            temp_dir = tempfile.TemporaryDirectory()
            temp_file_path = os.path.join(temp_dir.name, "uploaded_video.mp4")
            with open(temp_file_path, "wb") as f:
                f.write(uploaded_file.read())

            output_path = os.path.join(temp_dir.name, "processed_video.mp4")

            if processed_video_path is None:  
                with st.spinner('Processing and detecting fall...'):
                    yolo_processor = YOLOProcessor("assets/best.pt")
                    yolo_processor.process_video(temp_file_path, output_path)
                processed_video_path = output_path

            st.subheader("Result Video")
            st.video(processed_video_path)
            
            if st.button('Download Result Video'):
                download_file(processed_video_path, filename='processed_video.mp4')

            temp_dir.cleanup()

def download_image(image, filename):
    if isinstance(image, np.ndarray):
        image = Image.fromarray(image)
    image.save(filename)
    with open(filename, "rb") as f:
        image_bytes = f.read()
    b64 = base64.b64encode(image_bytes).decode()
    href = f'<a href="data:image/png;base64,{b64}" download="{filename}">Click here to download {filename}</a>'
    st.markdown(href, unsafe_allow_html=True)

def download_file(file_path, filename):
    with open(file_path, 'rb') as f:
        data = f.read()
    b64 = base64.b64encode(data).decode()
    href = f'<a href="data:file/mp4;base64,{b64}" download="{filename}">Click here to download {filename}</a>'
    st.markdown(href, unsafe_allow_html=True)

if __name__ == "__main__":
    main()