import cv2
import streamlit as st
import tempfile
import base64
import os
from dotenv import load_dotenv
from openai import OpenAI
import assemblyai as aai
from moviepy.editor import *


# Load environment variables
load_dotenv()
aai.settings.api_key = os.getenv("ASSEMBLYAI_API_KEY")
OpenAI.api_key = os.getenv("OPENAI_API_KEY")
client = OpenAI()
 
def main():
    st.title('Insightly Video Content Moderation')
 
    # Video upload section
    uploaded_video = st.file_uploader('Upload a video', type=["mp4", "avi", "mov"])
 
    if uploaded_video is not None:
        # Save the video to a temp file
        tfile = tempfile.NamedTemporaryFile(delete=False)
        tfile.write(uploaded_video.read())
        video_file_path = tfile.name
        tfile.close()

        transcriber = aai.Transcriber()
        transcript = transcriber.transcribe(tfile.name)

        # Process the video and display frames in a grid layout
        base64_frames = video_to_base64_frames(video_file_path)
        display_frame_grid(base64_frames[::30])  # Display every 30th frame in a 3-column grid
 
    st.write("Actions:")  # Header for the actions/buttons section

   
    spacer_col1, col1, spacer_col2, col2, spacer_col3 = st.columns([1, 2, 1, 2, 1])

    with col1:
        if st.button("Description"):
            st.session_state['description'] = generate_description(base64_frames) if 'description' not in st.session_state else st.session_state['description']
        
    with col2:
        if st.button("Generate Transcript"):
            transcript = transcriber.transcribe(video_file_path)
            st.session_state['transcript'] = transcript.text if 'transcript' not in st.session_state else st.session_state['transcript']
    
    

    # If any value exists in session state then display it
    if 'description' in st.session_state and st.session_state['description']:
        st.subheader("Video Description")
        st.write(st.session_state['description'])

    
    if 'transcript' in st.session_state and st.session_state['transcript']:
        st.subheader("Video Transcript")
        st.write(st.session_state['transcript'])
    




            

def video_to_base64_frames(video_file_path):
    # Logic to extract all frames from the video and convert them to base64
    video = cv2.VideoCapture(video_file_path)
    base64_frames = []
 
    while video.isOpened():
        success, frame = video.read()
        if not success:
            break
 
        _, buffer = cv2.imencode('.jpg', frame)
        base64_frame = base64.b64encode(buffer).decode('utf-8')
        base64_frames.append(base64_frame)
 
    video.release()
    return base64_frames

#########################################
#Generate Video description
def generate_description(base64_frames):
    prompt_messages = [
        {
            "role": "user",
            "content": [
                "1. Generate a description for this sequence of video frames in about 100 words.\
                Return the following :  2. Frame by frame summary of what's happening in the video. 3. List of objects in the video. 4. Any restrictive content or sensitive content and if so which frame 5. What category can this video be tagged to?",
                *map(lambda x: {"image": x, "resize": 428}, base64_frames[0::30]),
            ],
        },
    ]
    response = client.chat.completions.create(
        model="gpt-4-vision-preview",
        messages=prompt_messages,
        max_tokens=3000,
    )
    return response.choices[0].message.content


######################## 
def display_frame_grid(base64_frames):
    cols_per_row = 3
    n_frames = len(base64_frames)
    for idx in range(0, n_frames, cols_per_row):
        cols = st.columns(cols_per_row)
        for col_index in range(cols_per_row):
            frame_idx = idx + col_index
            if frame_idx < n_frames:
                with cols[col_index]:
                    frame = base64_frames[frame_idx]
                    st.image(base64.b64decode(frame), caption=f'Frame {frame_idx * 30 + 1}', width=200)
 
if __name__ == '__main__':
    main()