import streamlit as st
import google.generativeai as genai
import os
import tempfile
import time
import mimetypes
import subprocess
from pathlib import Path

# --- Get API key from environment variable or user input ---
def get_api_key():
    GOOGLE_API_KEY = os.environ.get("GOOGLE_API_KEY", "")
    if not GOOGLE_API_KEY:
        GOOGLE_API_KEY = st.text_input("Enter your Gemini API key", type="password")
    return GOOGLE_API_KEY or "AIzaSyA8TTu9s6fJDG9RlMwOyHFxg270xLgpiyE"  # Warning: Hardcoded key


# VideoProcessor class
class VideoProcessor:
    def __init__(self, api_key):
        genai.configure(api_key=api_key)
        self.model = genai.GenerativeModel("gemini-2.0-flash")

    def reduce_resolution(self, input_path, output_path, target_height=480):
        """Reduce video resolution to speed up processing."""
        try:
            command = [
                'ffmpeg', '-i', input_path,
                '-vf', f'scale=-2:{target_height}',
                '-c:a', 'copy',
                '-y', output_path
            ]
            subprocess.run(command, check=True, capture_output=True)
            return output_path
        except subprocess.CalledProcessError:
            # If ffmpeg fails, return original path
            return input_path

    def upload_video(self, video_path, display_name="uploaded_video"):
        try:
            return genai.upload_file(path=video_path, display_name=display_name)
        except Exception as e:
            raise RuntimeError(f"Failed to upload video: {str(e)}")

    def wait_for_processing(self, video_file, status_placeholder):
        max_attempts = 60  # Maximum wait time: 2 minutes
        attempts = 0
        while video_file.state.name == "PROCESSING" and attempts < max_attempts:
            # Update status text with dots animation
            dots = "." * ((attempts % 3) + 1)
            status_placeholder.markdown(f"**Processing video{dots}**")
            time.sleep(2)
            video_file = genai.get_file(video_file.name)
            attempts += 1
        if video_file.state.name == "FAILED":
            raise RuntimeError("Video processing failed")
        if attempts >= max_attempts:
            raise RuntimeError("Video processing timeout")
        return video_file

    def chat_with_video(self, video_file, prompt):
        try:
            response = self.model.generate_content([video_file, prompt])
            return response.text
        except Exception as e:
            return f"Error generating response: {str(e)}"

# Initialize session state properly
if "messages" not in st.session_state:
    st.session_state.messages = []
if "video_processor" not in st.session_state:
    st.session_state.video_processor = None
if "video_file" not in st.session_state:
    st.session_state.video_file = None
if "video_name" not in st.session_state:
    st.session_state.video_name = None

# Buffering animation CSS
def show_buffering_animation():
    st.markdown("""
        <style>
        .buffering-container {
            display: flex;
            flex-direction: column;
            align-items: center;
            justify-content: center;
            padding: 20px;
        }
        .buffering-spinner {
            width: 50px;
            height: 50px;
            border: 5px solid #f3f3f3;
            border-top: 5px solid #3498db;
            border-radius: 50%;
            animation: spin 1s linear infinite;
        }
        @keyframes spin {
            0% { transform: rotate(0deg); }
            100% { transform: rotate(360deg); }
        }
        .buffering-text {
            margin-top: 10px;
            font-size: 16px;
            color: #666;
        }
        </style>
        <div class="buffering-container">
            <div class="buffering-spinner"></div>
            <div class="buffering-text">Processing video...</div>
        </div>
    """, unsafe_allow_html=True)

# Main app function
def main():
    st.set_page_config(page_title="Video Retrieval-Augmented Generation", layout="wide")
    st.header("Video Retrieval-Augmented Generation - Gemini 2.0")
    st.markdown("---")

    # Step 1: API Key input
    api_key = get_api_key()

    if not api_key:
        st.error("Please enter your API key to proceed.")
        st.stop()

    # Step 2: Upload Video
    st.subheader("Step 1: Upload your video file")
    uploaded_file = st.file_uploader("Upload a video", type=['mp4', 'mov', 'avi', 'mkv', 'webm'])

    if uploaded_file:
        # Validate video file
        mime_type = mimetypes.guess_type(uploaded_file.name)[0]
        if mime_type and mime_type.startswith("video"):
            file_size = len(uploaded_file.getvalue()) / (1024**2)
            
            # Display file info
            col1, col2 = st.columns(2)
            with col1:
                st.info(f"File: {uploaded_file.name}")
            with col2:
                st.info(f"Size: {file_size:.2f} MB")
            
            # Check file size limit (200MB for Gemini)
            if file_size > 200:
                st.error("File too large! Maximum size is 200MB")
                st.stop()

            # Process new video if different from current
            if st.session_state.video_name != uploaded_file.name:
                tmp_path = None
                reduced_path = None
                try:
                    # Initialize processor if needed
                    if not st.session_state.video_processor:
                        st.session_state.video_processor = VideoProcessor(api_key)
                    
                    # Create temporary file
                    with tempfile.NamedTemporaryFile(delete=False, suffix=Path(uploaded_file.name).suffix) as tmp:
                        tmp.write(uploaded_file.getvalue())
                        tmp_path = tmp.name

                    # Show buffering animation container
                    buffering_container = st.empty()
                    status_text = st.empty()
                    
                    with buffering_container.container():
                        show_buffering_animation()
                    
                    # Reduce video resolution
                    status_text.markdown("**Reducing video resolution...**")
                    reduced_path = tmp_path.replace(Path(tmp_path).suffix, f"_reduced{Path(tmp_path).suffix}")
                    video_to_upload = st.session_state.video_processor.reduce_resolution(tmp_path, reduced_path)
                    
                    # Upload video
                    status_text.markdown("**Uploading video...**")
                    video_file = st.session_state.video_processor.upload_video(video_to_upload, uploaded_file.name)
                    
                    # Process video with status updates
                    processed_file = st.session_state.video_processor.wait_for_processing(video_file, status_text)
                    
                    # Clear buffering animation
                    buffering_container.empty()
                    status_text.empty()
                    
                    # Update session state
                    st.session_state.video_file = processed_file
                    st.session_state.video_name = uploaded_file.name
                    st.session_state.messages = []  # Clear previous conversation
                    
                    st.success("Video processed successfully!")
                    time.sleep(1)  # Show success message briefly

                except Exception as e:
                    st.error(f"Error processing video: {str(e)}")
                    st.session_state.video_file = None
                    st.session_state.video_name = None
                    
                finally:
                    # Clean up temporary files
                    if tmp_path and os.path.exists(tmp_path):
                        os.unlink(tmp_path)
                    if reduced_path and os.path.exists(reduced_path):
                        os.unlink(reduced_path)

            # Display video player
            st.video(uploaded_file.getvalue())
        else:
            st.error("Please upload a valid video file")

    # Control buttons
    col1, col2 = st.columns(2)
    with col1:
        if st.button("Reset Chat", disabled=not st.session_state.messages):
            st.session_state.messages = []
            st.rerun()
    
    with col2:
        if st.button("Reset All", disabled=not st.session_state.video_file):
            for key in list(st.session_state.keys()):
                del st.session_state[key]
            st.rerun()

    # Step 3: Chat about Video
    st.subheader("Step 2: Chat with your video")

    if st.session_state.video_file:
        # Display chat history
        for msg in st.session_state.messages:
            with st.chat_message(msg["role"]):
                st.markdown(msg["content"])

        # Chat input
        user_question = st.chat_input("Ask a question about the video...")
        if user_question:
            # Add user message
            st.session_state.messages.append({"role": "user", "content": user_question})
            with st.chat_message("user"):
                st.markdown(user_question)

            # Generate and display assistant response
            with st.chat_message("assistant"):
                placeholder = st.empty()
                with st.spinner("Thinking..."):
                    try:
                        response = st.session_state.video_processor.chat_with_video(
                            st.session_state.video_file, 
                            user_question
                        )
                    except Exception as e:
                        response = f"Error: {str(e)}"

                placeholder.markdown(response)
                st.session_state.messages.append({"role": "assistant", "content": response})
    else:
        st.info("Please upload a video in Step 1 to start chatting.")

if __name__ == "__main__":
    main()