import streamlit as st
from colab import AIDancePartner
import tempfile
import os
import time

def main():
    st.title("AI Dance Partner")
    st.write("Upload your dance video and get an AI dance partner!")
    
    # Video upload
    uploaded_file = st.file_uploader("Choose a video file", type=['mp4', 'avi', 'mov'])
    
    if uploaded_file is not None:
        # Create a temporary file
        with tempfile.NamedTemporaryFile(delete=False, suffix='.mp4') as tfile:
            tfile.write(uploaded_file.read())
            temp_input_path = tfile.name
        
        # Dance style selection
        style = st.selectbox(
            "Choose dance partner style",
            ["Sync Partner", "Creative Partner"]
        )
        
        if st.button("Generate Dance Partner"):
            try:
                with st.spinner("Processing your dance video..."):
                    # Process video
                    dance_partner = AIDancePartner()
                    output_path = dance_partner.process_video(temp_input_path, mode=style)
                    
                    # Display result
                    st.video(output_path)
                    
                    # Download button
                    with open(output_path, 'rb') as file:
                        st.download_button(
                            label="Download video",
                            data=file,
                            file_name="ai_dance_partner.mp4",
                            mime="video/mp4"
                        )
                    
                    # Cleanup temporary files
                    os.unlink(temp_input_path)
                    os.unlink(output_path)
                    
            except Exception as e:
                st.error(f"An error occurred: {str(e)}")
                if os.path.exists(temp_input_path):
                    os.unlink(temp_input_path)

if __name__ == "__main__":
    main()