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
from huggingface_hub import InferenceClient
from bs4 import BeautifulSoup
import textract
from xml.etree import ElementTree as ET
from urllib.parse import quote
import time
from collections import deque

# 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()

# Initialize OpenAI client
client = OpenAI(
    api_key=os.getenv('OPENAI_API_KEY'),
    organization=os.getenv('OPENAI_ORG_ID')
)

# Initialize session state
if "openai_model" not in st.session_state:
    st.session_state["openai_model"] = "gpt-4o-2024-05-13"
if "messages" not in st.session_state:
    st.session_state.messages = []

# Hugging Face settings
API_URL = os.getenv('API_URL')
HF_KEY = os.getenv('HF_KEY')
headers = {
    "Authorization": f"Bearer {HF_KEY}",
    "Content-Type": "application/json"
}

# 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": "✨"
        }
    },
    "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": "🦗"
        }
    }
}

# File handling functions
def generate_filename(prompt, file_type):
    """Generate a safe filename using the prompt and file type."""
    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 create_and_save_file(content, file_type="md", prompt=None, is_image=False, should_save=True):
    """Create and save file with proper handling of different types."""
    if not should_save:
        return None
    
    filename = generate_filename(prompt if prompt else content, file_type)
    
    if file_type == "md":
        title_from_content = extract_markdown_title(content)
        if title_from_content:
            filename = generate_filename(title_from_content, file_type)
    
    with open(filename, "w", encoding="utf-8") as f:
        if is_image:
            f.write(content)
        else:
            f.write(prompt + "\n\n" + content)
    
    return filename

def extract_markdown_title(content):
    """Extract the first markdown title from content."""
    title_match = re.search(r'^\s*#\s*(.+)', content, re.MULTILINE)
    if title_match:
        return title_match.group(1).strip()
    return None

# HTML5 Speech Synthesis
@st.cache_resource
def SpeechSynthesis(result):
    documentHTML5 = f'''
    <!DOCTYPE html>
    <html>
    <head>
        <title>Read It Aloud</title>
        <script type="text/javascript">
            function readAloud() {{
                const text = document.getElementById("textArea").value;
                const speech = new SpeechSynthesisUtterance(text);
                window.speechSynthesis.speak(speech);
            }}
        </script>
    </head>
    <body>
        <h1>🔊 Read It Aloud</h1>
        <textarea id="textArea" rows="10" cols="80">{result}</textarea>
        <br>
        <button onclick="readAloud()">🔊 Read Aloud</button>
    </body>
    </html>
    '''
    st.components.v1.html(documentHTML5, width=1280, height=300)

# Process functions for different media types
def process_text(text_input):
    """Process text input with GPT-4o."""
    if text_input:
        st.session_state.messages.append({"role": "user", "content": text_input})
        
        with st.chat_message("user"):
            st.markdown(text_input)
        
        with st.chat_message("assistant"):
            completion = client.chat.completions.create(
                model=st.session_state["openai_model"],
                messages=[
                    {"role": m["role"], "content": m["content"]}
                    for m in st.session_state.messages
                ],
                stream=False
            )
            return_text = completion.choices[0].message.content
            st.write("Assistant: " + return_text)
            
            create_and_save_file(return_text, file_type="md", prompt=text_input)
            st.session_state.messages.append({"role": "assistant", "content": return_text})

def process_image(image_input, user_prompt):
    """Process image with GPT-4o vision."""
    if isinstance(image_input, str):
        with open(image_input, "rb") as image_file:
            image_input = image_file.read()
            
    base64_image = base64.b64encode(image_input).decode("utf-8")
    
    response = client.chat.completions.create(
        model=st.session_state["openai_model"],
        messages=[
            {"role": "system", "content": "You are a helpful assistant that responds in Markdown."},
            {"role": "user", "content": [
                {"type": "text", "text": user_prompt},
                {"type": "image_url", "image_url": {
                    "url": f"data:image/png;base64,{base64_image}"
                }}
            ]}
        ],
        temperature=0.0,
    )
    
    return response.choices[0].message.content

def process_audio(audio_input, text_input=''):
    """Process audio with GPT-4o and Whisper."""
    if isinstance(audio_input, str):
        with open(audio_input, "rb") as file:
            audio_input = file.read()

    transcription = client.audio.transcriptions.create(
        model="whisper-1",
        file=audio_input,
    )
    
    st.session_state.messages.append({"role": "user", "content": transcription.text})
    
    with st.chat_message("assistant"):
        st.markdown(transcription.text)
        SpeechSynthesis(transcription.text)
        
        filename = generate_filename(transcription.text, "wav")
        create_and_save_file(audio_input.getvalue(), "wav", transcription.text, True)

def process_video(video_path, seconds_per_frame=1):
    """Process video files for frame extraction and audio."""
    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():
    """Create the media gallery interface."""
    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))
                    
                    # Add prompt input for GPT-4o analysis
                    prompt = st.text_input(f"Analyze image {idx}", 
                                         "Describe this image in detail and list key elements.")
                    if st.button(f"Analyze {idx}"):
                        analysis = process_image(image_file, prompt)
                        st.markdown(analysis)
    
    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)
                if st.button(f"Transcribe {audio_file}"):
                    process_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)
                if st.button(f"Analyze {video_file}"):
                    frames, audio = process_video(video_file)
                    if audio:
                        st.audio(audio)

    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", "🤖 Chat"],
                        horizontal=True)
    
    if tab_main == "📸 Upload Media":
        col1, col2 = st.columns(2)
        
        with col1:
            uploaded_image = st.file_uploader("Upload Image", type=['png', 'jpg'])
            if uploaded_image:
                st.image(uploaded_image)
                prompt = st