import streamlit as st
import os
import cv2
import tempfile
import torch

# Import the main app class and config flags
from motion_app import MotionComparisonApp
from config import REALSENSE_AVAILABLE, PYGAME_AVAILABLE

def main():
    """Main function to run the Streamlit app."""
    st.set_page_config(page_title="Motion Comparison", page_icon="🏃", layout="wide")
    st.title("🏃 Motion Comparison & Pose Analysis System")
    st.markdown("---")

    # Initialize the app object in session state
    if 'app' not in st.session_state:
        st.session_state.app = MotionComparisonApp()
    app = st.session_state.app

    # --- Sidebar UI ---
    with st.sidebar:
        st.header("🎛️ Control Panel")
        
        # Display settings
        resolution_options = {"High (1280x800)": "high", "Medium (960x720)": "medium", "Standard (640x480)": "low"}
        selected_res = st.selectbox("Display Resolution", list(resolution_options.keys()), index=1)
        app.display_settings['resolution_mode'] = resolution_options[selected_res]

        st.markdown("---")
        
        # Video Source Selection
        video_source = st.radio("Video Source", ["Preset Video", "Upload Video"])
        video_path = None
        
        if video_source == "Preset Video":
            preset_path = "preset_videos/liuzi.mp4"
            if os.path.exists(preset_path):
                st.success("✅ '六字诀' video found.")
                video_path = preset_path
            else:
                st.error("❌ Preset video not found. Please place 'liuzi.mp4' in 'preset_videos' folder.")
        else:
            uploaded_file = st.file_uploader("Upload a video", type=['mp4', 'avi', 'mov', 'mkv'])
            if uploaded_file:
                with tempfile.NamedTemporaryFile(delete=False, suffix='.mp4') as tmp_file:
                    tmp_file.write(uploaded_file.read())
                    video_path = tmp_file.name

        st.markdown("---")
        
        # System Initialization
        st.subheader("⚙️ System Initialization")
        if st.button("🚀 Initialize System", use_container_width=True):
            with st.spinner("Initializing detectors and cameras..."):
                app.initialize_detector()
                app.initialize_camera()
        
        # System Status
        st.subheader("ℹ️ System Status")
        st.info(f"Computation: {'GPU (CUDA)' if torch.cuda.is_available() else 'CPU'}")
        st.info(f"Camera: {'RealSense' if REALSENSE_AVAILABLE else 'USB Webcam'}")
        st.info(f"Audio: {'Enabled' if PYGAME_AVAILABLE else 'Disabled'}")

    # --- Main Page UI ---
    if video_path:
        # Display video info and control buttons
        # This part is identical to your original `main` function's logic
        # It sets up the "Preview Camera" and "Start Comparison" buttons
        # And calls app.start_comparison(video_path) when clicked.
        
        # Example of how you would structure the main page:
        if st.button("🚀 Start Comparison", use_container_width=True):
             if not app.body_detector:
                 st.error("⚠️ Please initialize the system from the sidebar first!")
             else:
                 # The start_comparison method now contains the main display loop
                 app.start_comparison(video_path)
    else:
        st.info("👈 Please select or upload a standard video from the sidebar to begin.")
        with st.expander("📖 Usage Guide", expanded=True):
             st.markdown("""
             1. **Select Video**: Choose a preset or upload your own video in the sidebar.
             2. **Initialize**: Click 'Initialize System' to prepare the camera and AI model.
             3. **Start**: Click 'Start Comparison' to begin the analysis.
             """)

if __name__ == "__main__":
    # Set environment variables for performance
    os.environ['OMP_NUM_THREADS'] = '1'
    os.environ['MKL_NUM_THREADS'] = '1'
    try:
        import torch
        torch.set_num_threads(1)
    except ImportError:
        pass
        
    main()