import streamlit as st
import base64
from datetime import datetime
import plotly.graph_objects as go
import cv2
import os
import pytz
import random
import re
import requests
from moviepy.editor import VideoFileClip
from PIL import Image
import glob
from audio_recorder_streamlit import audio_recorder
import json
from openai import OpenAI
from dotenv import load_dotenv

# Page config
st.set_page_config(
    page_title="Bike Cinematic Universe 🎬",
    page_icon="🚲",
    layout="wide"
)

# Custom CSS with expanded styling
st.markdown("""
<style>
    .main {
        background: linear-gradient(to right, #1a1a1a, #2d2d2d);
        color: #ffffff;
    }
    .stMarkdown {
        font-family: 'Helvetica Neue', sans-serif;
    }
    .category-header {
        background: linear-gradient(45deg, #2b5876, #4e4376);
        padding: 20px;
        border-radius: 10px;
        margin: 10px 0;
    }
    .scene-card {
        background: rgba(0,0,0,0.3);
        padding: 15px;
        border-radius: 8px;
        margin: 10px 0;
        border: 1px solid rgba(255,255,255,0.1);
    }
    .media-gallery {
        display: grid;
        gap: 1rem;
        padding: 1rem;
    }
    .bike-card {
        background: rgba(255,255,255,0.05);
        border-radius: 10px;
        padding: 15px;
        transition: transform 0.3s;
    }
    .bike-card:hover {
        transform: scale(1.02);
    }
</style>
""", unsafe_allow_html=True)

# Load environment variables
load_dotenv()
client = OpenAI(api_key=os.getenv('OPENAI_API_KEY'))

# Bike Collections
bike_collections = {
    "Celestial Collection 🌌": {
        "Eclipse Vaulter": {
            "prompt": """Cinematic shot of a sleek black mountain bike silhouetted against a total solar eclipse. 
                     The corona creates an ethereal halo effect, with lens flares accentuating key points of the frame.
                     Dynamic composition shows the bike mid-leap, with stardust particles trailing behind.
                     Camera angle: Low angle, wide shot
                     Lighting: Dramatic rim lighting from eclipse
                     Color palette: Deep purples, cosmic blues, corona gold""",
            "emoji": "🌑"
        },
        "Starlight Leaper": {
            "prompt": """A black bike performing an epic leap under a vast Milky Way galaxy.
                     Shimmering stars blanket the sky while the bike's wheels leave a trail of stardust.
                     Camera angle: Wide-angle upward shot
                     Lighting: Natural starlight with subtle rim lighting
                     Color palette: Deep blues, silver highlights, cosmic purples""",
            "emoji": "✨"
        },
        "Moonlit Hopper": {
            "prompt": """A sleek black bike mid-hop over a moonlit meadow.
                     Full moon illuminating misty surroundings with fireflies dancing around.
                     Camera angle: Side profile with slight low angle
                     Lighting: Soft moonlight with atmospheric fog
                     Color palette: Silver blues, soft whites, deep shadows""",
            "emoji": "🌙"
        }
    },
    "Nature-Inspired Collection 🌲": {
        "Shadow Grasshopper": {
            "prompt": """A black bike jumping between forest paths.
                     Dappled sunlight streams through the canopy, creating dynamic shadows.
                     Camera angle: Through-the-trees tracking shot
                     Lighting: Natural forest lighting with sun rays
                     Color palette: Forest greens, golden sunlight, deep shadows""",
            "emoji": "🦗"
        },
        "Onyx Leapfrog": {
            "prompt": """A bike with obsidian-black finish jumping over a sparkling creek.
                     Water reflection creates mirror effect with ripples from the leap.
                     Camera angle: Low angle from water level
                     Lighting: Golden hour side lighting
                     Color palette: Deep blacks, water blues, forest greens""",
            "emoji": "🐸"
        }
    }
}

# File handling functions
def generate_filename(prompt, file_type):
    """Generate a safe filename from prompt and timestamp"""
    central = pytz.timezone('US/Central')
    safe_date_time = datetime.now(central).strftime("%m%d_%H%M")
    replaced_prompt = re.sub(r'[<>:"/\\|?*\n]', ' ', prompt)
    safe_prompt = re.sub(r'\s+', ' ', replaced_prompt).strip()[:240]
    return f"{safe_date_time}_{safe_prompt}.{file_type}"

def save_file(content, filename, is_binary=False):
    """Save content to file with proper mode"""
    mode = 'wb' if is_binary else 'w'
    with open(filename, mode) as f:
        f.write(content)
    return filename

def process_video(video_path, seconds_per_frame=1):
    """Extract frames and audio from video"""
    base64Frames = []
    video = cv2.VideoCapture(video_path)
    total_frames = int(video.get(cv2.CAP_PROP_FRAME_COUNT))
    fps = video.get(cv2.CAP_PROP_FPS)
    frames_to_skip = int(fps * seconds_per_frame)
    
    for frame_idx in range(0, total_frames, frames_to_skip):
        video.set(cv2.CAP_PROP_POS_FRAMES, frame_idx)
        success, frame = video.read()
        if not success:
            break
        _, buffer = cv2.imencode(".jpg", frame)
        base64Frames.append(base64.b64encode(buffer).decode("utf-8"))
    
    video.release()
    
    # Extract audio
    base_video_path = os.path.splitext(video_path)[0]
    audio_path = f"{base_video_path}.mp3"
    try:
        video_clip = VideoFileClip(video_path)
        video_clip.audio.write_audiofile(audio_path)
        video_clip.close()
    except:
        st.warning("No audio track found in video")
        audio_path = None
    
    return base64Frames, audio_path

def create_media_gallery():
    st.header("🎬 Media Gallery")
    
    tabs = st.tabs(["🖼️ Images", "🎵 Audio", "🎥 Video", "🎨 Scene Generator"])
    
    with tabs[0]:
        image_files = glob.glob("*.png") + glob.glob("*.jpg")
        if image_files:
            cols = st.columns(3)
            for idx, image_file in enumerate(image_files):
                with cols[idx % 3]:
                    st.image(image_file)
                    st.caption(os.path.basename(image_file))
    
    with tabs[1]:
        audio_files = glob.glob("*.mp3") + glob.glob("*.wav")
        for audio_file in audio_files:
            with st.expander(f"🎵 {os.path.basename(audio_file)}"):
                st.audio(audio_file)
    
    with tabs[2]:
        video_files = glob.glob("*.mp4")
        for video_file in video_files:
            with st.expander(f"🎥 {os.path.basename(video_file)}"):
                st.video(video_file)

    with tabs[3]:
        for collection_name, bikes in bike_collections.items():
            st.subheader(collection_name)
            cols = st.columns(len(bikes))
            
            for idx, (bike_name, details) in enumerate(bikes.items()):
                with cols[idx]:
                    st.markdown(f"""
                    <div class='bike-card'>
                        <h3>{details['emoji']} {bike_name}</h3>
                        <p>{details['prompt']}</p>
                    </div>
                    """, unsafe_allow_html=True)

def main():
    st.title("🚲 Bike Cinematic Universe")
    
    # Main navigation
    tab_main = st.radio("Choose Action:", 
                        ["📸 Upload Media", "🎬 View Gallery", "🎨 Generate Scene"],
                        horizontal=True)
    
    if tab_main == "📸 Upload Media":
        col1, col2 = st.columns(2)
        
        with col1:
            # Image upload
            uploaded_image = st.file_uploader("Upload Image", type=['png', 'jpg'])
            if uploaded_image:
                st.image(uploaded_image)
                prompt = st.text_input("Image Description:")
                if st.button("Process Image"):
                    filename = generate_filename(prompt, uploaded_image.type.split('/')[-1])
                    save_file(uploaded_image.getvalue(), filename, is_binary=True)
                    st.success(f"Saved as {filename}")
        
        with col2:
            # Audio/Video upload
            uploaded_media = st.file_uploader("Upload Audio/Video", type=['mp3', 'wav', 'mp4'])
            if uploaded_media:
                if uploaded_media.type.startswith('audio'):
                    st.audio(uploaded_media)
                else:
                    st.video(uploaded_media)
                if st.button("Save Media"):
                    filename = generate_filename("media", uploaded_media.type.split('/')[-1])
                    save_file(uploaded_media.getvalue(), filename, is_binary=True)
                    st.success(f"Saved as {filename}")
    
    elif tab_main == "🎬 View Gallery":
        create_media_gallery()
    
    else:  # Generate Scene
        st.header("🎨 Scene Generator")
        selected_collection = st.selectbox("Choose Collection", list(bike_collections.keys()))
        selected_bike = st.selectbox("Choose Bike", list(bike_collections[selected_collection].keys()))
        
        bike_details = bike_collections[selected_collection][selected_bike]
        st.markdown(f"""
        <div class='scene-card'>
            <h3>{bike_details['emoji']} {selected_bike}</h3>
            <p>{bike_details['prompt']}</p>
        </div>
        """, unsafe_allow_html=True)

if __name__ == "__main__":
    main()