import av
import os
import sys
import streamlit as st
import cv2
import tempfile

# Define the base directory path
BASE_DIR = os.path.abspath(os.path.join(__file__, '../../'))
sys.path.append(BASE_DIR)

# Import functions from other python files
from utils import get_mediapipe_pose
from process_frame import ProcessFrame
from thresholds import get_thresholds_beginner, get_thresholds_pro

st.title('Mediapipe Exercise Sample')

# Create a button to select mode (Beginner or Pro)
mode = st.radio('Select Mode', ['Beginner', 'Pro'], horizontal=True)

# Initialize thresholds based on the selected mode
thresholds = None
if mode == 'Beginner':
    thresholds = get_thresholds_beginner()

elif mode == 'Pro':
    thresholds = get_thresholds_pro()

# Create a ProcessFrame object with the loaded thresholds
upload_process_frame = ProcessFrame(thresholds=thresholds)

# Initialize Mediapipe pose solution
pose = get_mediapipe_pose()

# Initialize download flag
download = None

# Set initial state for download in Streamlit session
if 'download' not in st.session_state:
    st.session_state['download'] = False

# Define the output video file name
output_video_file = f'output_recorded.mp4'

# Remove the output video file if it exists
if os.path.exists(output_video_file):
    os.remove(output_video_file)

# Create a form for uploading a video
with st.form('Upload', clear_on_submit=True):
    up_file = st.file_uploader("Upload a Video", ['mp4','mov', 'avi'])
    uploaded = st.form_submit_button("Upload")

# Create an empty element to display the video frame
stframe = st.empty()

# Define HTML strings for displaying input video and warning message
ip_vid_str = '<p style="font-family:Helvetica; font-weight: bold; font-size: 16px;">Input Video</p>'
warning_str = '<p style="font-family:Helvetica; font-weight: bold; color: Red; font-size: 17px;">Please Upload a Video first!!!</p>'

# Create an empty element to display the warning message
warn = st.empty()

# Create an empty element for the download button
download_button = st.empty()

# Process uploaded video if a video is uploaded and submit button is clicked
if up_file and uploaded:
    # Clear previous download button and warning message
    download_button.empty()
    
    # Create a temporary file to store the uploaded video
    tfile = tempfile.NamedTemporaryFile(delete=False)
    
    try:
        warn.empty()
        # Write the uploaded video content to the temporary file
        tfile.write(up_file.read())
        
        # Open the temporary file using OpenCV VideoCapture
        vf = cv2.VideoCapture(tfile.name)
        
        # Get video properties (FPS, width, height)
        fps = int(vf.get(cv2.CAP_PROP_FPS))
        width = int(vf.get(cv2.CAP_PROP_FRAME_WIDTH))
        height = int(vf.get(cv2.CAP_PROP_FRAME_HEIGHT))
        frame_size = (width, height)
        
        # Define video writer fourcc code for mp4 format
        fourcc = cv2.VideoWriter_fourcc(*'mp4v')
        
        # Create video writer object for the output video
        video_output = cv2.VideoWriter(output_video_file, fourcc, fps, frame_size)
        
        # Display the uploaded video on the sidebar
        txt = st.sidebar.markdown(ip_vid_str, unsafe_allow_html=True)
        ip_video = st.sidebar.video(tfile.name)
        
        # Process each frame of the video
        while vf.isOpened():
            ret, frame = vf.read()
            if not ret:
                break
          
            # Convert frame from BGR to RGB for Mediapipe processing
            frame = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
          
            # Process the frame using the ProcessFrame object and
            # (continued from previous code block)
            out_frame, _ = upload_process_frame.process(frame, pose)
            stframe.image(out_frame)
            video_output.write(out_frame[...,::-1])  # Write the processed frame to output video
        
        # Release video capture and writer resources
        vf.release()
        video_output.release()
        
        # Clear elements that displayed video and temporary file
        stframe.empty()
        ip_video.empty()
        txt.empty()
        tfile.close()
    
    except AttributeError:
        # Handle errors during processing (e.g., invalid video format)
        warn.markdown(warning_str, unsafe_allow_html=True)

# Check if output video exists and offer download button
if os.path.exists(output_video_file):
    with open(output_video_file, 'rb') as op_vid:
        download = download_button.download_button('Download Video', data=op_vid, file_name='output_recorded.mp4')

    # If video is downloaded, update download flag in Streamlit session state
    if download:
        st.session_state['download'] = True

# Remove output video and reset download flag if video downloaded 
if os.path.exists(output_video_file) and st.session_state['download']:
    os.remove(output_video_file)
    st.session_state['download'] = False
    download_button.empty()