import streamlit as st
import time
import random
import json
from datetime import datetime
import pytz
import platform
import uuid
import extra_streamlit_components as stx
from io import BytesIO
from PIL import Image
import base64
import cv2
import requests
from moviepy.editor import VideoFileClip
from gradio_client import Client
from openai import OpenAI
import openai
import os
from collections import deque

# Set page config
st.set_page_config(page_title="Personalized Real-Time Chat", page_icon="💬", layout="wide")

# Initialize cookie manager
cookie_manager = stx.CookieManager()

# File to store chat history and user data
CHAT_FILE = "chat_history.txt"

# Function to save chat history and user data to file
def save_data():
    with open(CHAT_FILE, 'w') as f:
        json.dump({
            'messages': st.session_state.messages,
            'users': st.session_state.users
        }, f)

# Function to load chat history and user data from file
def load_data():
    try:
        with open(CHAT_FILE, 'r') as f:
            data = json.load(f)
            st.session_state.messages = data['messages']
            st.session_state.users = data['users']
    except FileNotFoundError:
        st.session_state.messages = []
        st.session_state.users = []

# Load data at the start
load_data()

# Function to get or create user
def get_or_create_user():
    user_id = cookie_manager.get(cookie='user_id')
    if not user_id:
        user_id = str(uuid.uuid4())
        cookie_manager.set('user_id', user_id)
    
    user = next((u for u in st.session_state.users if u['id'] == user_id), None)
    if not user:
        user = {
            'id': user_id,
            'name': random.choice(['Alice', 'Bob', 'Charlie', 'David', 'Eve', 'Frank', 'Grace', 'Henry']),
            'browser': f"{platform.system()} - {st.session_state.get('browser_info', 'Unknown')}"
        }
        st.session_state.users.append(user)
        save_data()
    
    return user

# Initialize session state
if 'messages' not in st.session_state:
    st.session_state.messages = []
if 'users' not in st.session_state:
    st.session_state.users = []
if 'current_user' not in st.session_state:
    st.session_state.current_user = get_or_create_user()

# Sidebar for user information and settings
with st.sidebar:
    st.title("User Info")
    st.write(f"Current User: {st.session_state.current_user['name']}")
    st.write(f"Browser: {st.session_state.current_user['browser']}")
    
    new_name = st.text_input("Change your name:")
    if st.button("Update Name"):
        if new_name:
            for user in st.session_state.users:
                if user['id'] == st.session_state.current_user['id']:
                    user['name'] = new_name
                    st.session_state.current_user['name'] = new_name
                    save_data()
                    st.success(f"Name updated to {new_name}")
                    break
    st.title("Active Users")
    for user in st.session_state.users:
        st.write(f"{user['name']} ({user['browser']})")

# Main chat area
st.title("Personalized Real-Time Chat")

# Function to generate filenames
def generate_filename(prompt, file_type):
    central = pytz.timezone('US/Central')
    safe_date_time = datetime.now(central).strftime("%m%d_%H%M")
    replaced_prompt = prompt.replace(" ", "_").replace("\n", "_")
    safe_prompt = "".join(x for x in replaced_prompt if x.isalnum() or x == "_")[:90]
    return f"{safe_date_time}_{safe_prompt}.{file_type}"

# Function to create files
def create_file(filename, prompt, response, user_name, timestamp, is_image=False):
    with open(filename, "w", encoding="utf-8") as f:
        f.write(f"User: {user_name}\nTimestamp: {timestamp}\n\nPrompt:\n{prompt}\n\nResponse:\n{response}")

# Function to process text
def process_text(user_name, text_input):
    if text_input:
        timestamp = datetime.now(pytz.utc).strftime('%Y-%m-%d %H:%M:%S %Z')
        st.session_state.messages.append({"user": user_name, "message": text_input, "timestamp": timestamp})
        with st.chat_message(user_name):
            st.markdown(f"{user_name} ({timestamp}): {text_input}")
        with st.chat_message("Assistant"):
            completion = client.chat.completions.create(
                model=MODEL,
                messages=[
                    {"role": "user", "content": m["message"]}
                    for m in st.session_state.messages
                ],
                stream=False
            )
            return_text = completion.choices[0].message.content
            st.markdown(f"Assistant ({timestamp}): {return_text}")
            filename = generate_filename(text_input, "md")
            create_file(filename, text_input, return_text, user_name, timestamp)
            st.session_state.messages.append({"user": "Assistant", "message": return_text, "timestamp": timestamp})
            save_data()

# Function to process image
def process_image(user_name, image_input, user_prompt):
    image = Image.open(BytesIO(image_input))
    base64_image = base64.b64encode(image_input).decode("utf-8")
    
    response = client.chat.completions.create(
        model=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,
    )
    image_response = response.choices[0].message.content
    
    timestamp = datetime.now(pytz.utc).strftime('%Y-%m-%d %H:%M:%S %Z')
    st.session_state.messages.append({"user": user_name, "message": image_response, "timestamp": timestamp})
    
    with st.chat_message(user_name):
        st.image(image)
        st.markdown(f"{user_name} ({timestamp}): {user_prompt}")
    
    with st.chat_message("Assistant"):
        st.markdown(image_response)
    
    filename_md = generate_filename(user_prompt, "md")
    create_file(filename_md, user_prompt, image_response, user_name, timestamp)
    save_data()
    return image_response

# Function to process audio
def process_audio(user_name, audio_input, text_input):
    if audio_input:
        transcription = client.audio.transcriptions.create(
            model="whisper-1",
            file=audio_input,
        )
        timestamp = datetime.now(pytz.utc).strftime('%Y-%m-%d %H:%M:%S %Z')
        st.session_state.messages.append({"user": user_name, "message": transcription.text, "timestamp": timestamp})
        with st.chat_message(user_name):
            st.markdown(f"{user_name} ({timestamp}): {transcription.text}")
        with st.chat_message("Assistant"):
            st.markdown(transcription.text)
            filename = generate_filename(transcription.text, "wav")
            create_file(filename, text_input, transcription.text, user_name, timestamp)
            st.session_state.messages.append({"user": "Assistant", "message": transcription.text, "timestamp": timestamp})
            save_data()

# Function to process video
def process_video(user_name, video_input, user_prompt):
    if isinstance(video_input, str):
        with open(video_input, "rb") as video_file:
            video_input = video_file.read()
    base64Frames, audio_path = extract_video_frames(video_input)
    transcript = process_audio_for_video(video_input)
    response = client.chat.completions.create(
        model=MODEL,
        messages=[
            {"role": "system", "content": "You are generating a video summary. Create a summary of the provided video and its transcript. Respond in Markdown"},
            {"role": "user", "content": [
                "These are the frames from the video.",
                *map(lambda x: {"type": "image_url", "image_url": {"url": f'data:image/jpg;base64,{x}', "detail": "low"}}, base64Frames),
                {"type": "text", "text": f"The audio transcription is: {transcript}"},
                {"type": "text", "text": user_prompt}
            ]}
        ],
        temperature=0,
    )
    video_response = response.choices[0].message.content
    st.markdown(video_response)
    timestamp = datetime.now(pytz.utc).strftime('%Y-%m-%d %H:%M:%S %Z')
    filename_md = generate_filename(user_prompt, "md")
    create_file(filename_md, user_prompt, video_response, user_name, timestamp)
    st.session_state.messages.append({"user": user_name, "message": video_response, "timestamp": timestamp})
    save_data()
    return video_response

# Function to extract video frames
def extract_video_frames(video_path, seconds_per_frame=2):
    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)
    curr_frame = 0
    while curr_frame < total_frames - 1:
        video.set(cv2.CAP_PROP_POS_FRAMES, curr_frame)
        success, frame = video.read()
        if not success:
            break
        _, buffer = cv2.imencode(".jpg", frame)
        base64Frames.append(base64.b64encode(buffer).decode("utf-8"))
        curr_frame += frames_to_skip
    video.release()
    return base64Frames, None

# Function to process audio for video
def process_audio_for_video(video_input):
    try:
        transcription = client.audio.transcriptions.create(
            model="whisper-1",
            file=video_input,
        )
        return transcription.text
    except:
        return ''

# Initialize OpenAI client
openai.api_key = os.getenv('OPENAI_API_KEY')
openai.organization = os.getenv('OPENAI_ORG_ID')
client = OpenAI(api_key=openai.api_key, organization=openai.organization)
MODEL = "gpt-4o-2024-05-13"
should_save = st.sidebar.checkbox("💾 Save", value=True, help="Save your session data.")

# Function to display chat messages
def display_messages():
    for msg in st.session_state.messages:
        with st.chat_message(msg['user']):
            st.markdown(f"**{msg['user']}** ({msg['timestamp']}): {msg['message']}")

# Display messages
display_messages()

# Main function
def main():
    st.markdown("##### GPT-4o Omni Model: Text, Audio, Image, & Video")
    option = st.selectbox("Select an option", ("Text", "Image", "Audio", "Video"))
    if option == "Text":
        text_input = st.text_input("Enter your text:")
        if text_input:
            process_text(st.session_state.current_user['name'], text_input)
    elif option == "Image":
        text_input = st.text_input("Enter text prompt to use with Image context:")
        uploaded_files = st.file_uploader("Upload images", type=["png", "jpg", "jpeg"], accept_multiple_files=True)
        for image_input in uploaded_files:
            image_bytes = image_input.read()
            process_image(st.session_state.current_user['name'], image_bytes, text_input)
    elif option == "Audio":
        text_input = st.text_input("Enter text prompt to use with Audio context:")
        uploaded_files = st.file_uploader("Upload an audio file", type=["mp3", "wav"], accept_multiple_files=True)
        for audio_input in uploaded_files:
            process_audio(st.session_state.current_user['name'], audio_input, text_input)
    elif option == "Video":
        video_input = st.file_uploader("Upload a video file", type=["mp4"])
        text_input = st.text_input("Enter text prompt to use with Video context:")
        if video_input and text_input:
            process_video(st.session_state.current_user['name'], video_input, text_input)

# Add buttons for quick access to different UI options
with st.sidebar:
    if st.button("📝 Add Text"):
        st.session_state.ui_option = "Text"
    if st.button("🖼️ Add Image"):
        st.session_state.ui_option = "Image"
    if st.button("🎵 Add Audio"):
        st.session_state.ui_option = "Audio"
    if st.button("🎥 Add Video"):
        st.session_state.ui_option = "Video"

if 'ui_option' in st.session_state:
    main()
else:
    st.session_state.ui_option = "Text"
    main()