import streamlit as st
import openai
from openai import OpenAI
import os
import base64
import cv2
from moviepy.editor import VideoFileClip

# Set API key and organization ID from environment variables
openai.api_key = os.getenv('OPENAI_API_KEY')
openai.organization = os.getenv('OPENAI_ORG_ID')
client = OpenAI(api_key= os.getenv('OPENAI_API_KEY'), organization=os.getenv('OPENAI_ORG_ID'))

# Define the model to be used
#MODEL = "gpt-4o"
MODEL = "gpt-4o-2024-05-13"

def process_text():
    text_input = st.text_input("Enter your text:")
    if text_input:
        completion = client.chat.completions.create(
            model=MODEL,
            messages=[
                {"role": "system", "content": "You are a helpful assistant. Help me with my math homework!"},
                {"role": "user", "content": f"Hello! Could you solve {text_input}?"}
            ]
        )
        st.write("Assistant: " + completion.choices[0].message.content)

def process_text_old():
    text_input = st.text_input("Enter your text:")
    if text_input:
        response = openai.Completion.create(
            model=MODEL,
            prompt=f"You are a helpful assistant. Help me with my math homework! {text_input}",
            max_tokens=100,
            temperature=0.5,
        )
        st.write("Assistant: " + response.choices[0].text.strip())

def process_image(image_input):
    if image_input:
        base64_image = base64.b64encode(image_input.read()).decode("utf-8")
        response = openai.Completion.create(
            model=MODEL,
            prompt=f"You are a helpful assistant that responds in Markdown. Help me with my math homework! What's the area of the triangle? [image: data:image/png;base64,{base64_image}]",
            max_tokens=100,
            temperature=0.5,
        )
        st.markdown(response.choices[0].text.strip())

def process_audio(audio_input):
    if audio_input:
        transcription = openai.Audio.transcriptions.create(
            model="whisper-1",
            file=audio_input,
        )
        response = openai.Completion.create(
            model=MODEL,
            prompt=f"You are generating a transcript summary. Create a summary of the provided transcription. Respond in Markdown. The audio transcription is: {transcription['text']}",
            max_tokens=100,
            temperature=0.5,
        )
        st.markdown(response.choices[0].text.strip())

def process_video(video_input):
    if video_input:
        base64Frames, audio_path = process_video_frames(video_input)
        transcription = openai.Audio.transcriptions.create(
            model="whisper-1",
            file=open(audio_path, "rb"),
        )
        frames_text = " ".join([f"[image: data:image/jpg;base64,{frame}]" for frame in base64Frames])
        response = openai.Completion.create(
            model=MODEL,
            prompt=f"You are generating a video summary. Create a summary of the provided video and its transcript. Respond in Markdown. These are the frames from the video. {frames_text} The audio transcription is: {transcription['text']}",
            max_tokens=500,
            temperature=0.5,
        )
        st.markdown(response.choices[0].text.strip())

def process_video_frames(video_path, seconds_per_frame=2):
    base64Frames = []
    base_video_path, _ = os.path.splitext(video_path.name)
    video = cv2.VideoCapture(video_path.name)
    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()
    audio_path = f"{base_video_path}.mp3"
    clip = VideoFileClip(video_path.name)
    clip.audio.write_audiofile(audio_path, bitrate="32k")
    clip.audio.close()
    clip.close()
    return base64Frames, audio_path

def main():
    st.title("Omni Demo")
    option = st.selectbox("Select an option", ("Text", "Image", "Audio", "Video"))
    if option == "Text":
        process_text()
    elif option == "Image":
        image_input = st.file_uploader("Upload an image", type=["jpg", "jpeg", "png"])
        process_image(image_input)
    elif option == "Audio":
        audio_input = st.file_uploader("Upload an audio file", type=["mp3", "wav"])
        process_audio(audio_input)
    elif option == "Video":
        video_input = st.file_uploader("Upload a video file", type=["mp4"])
        process_video(video_input)

if __name__ == "__main__":
    main()