import os
import google.generativeai as genai
import gradio as gr
import requests
from moviepy.editor import ImageClip, AudioFileClip

# Configure Google Gemini API
genai.configure(api_key=os.getenv("GEMINI_API_KEY"))

# Play.ht API keys
API_KEY = os.getenv('PLAY_API_KEY')
USER_ID = os.getenv('PLAY_USER_ID')

# theme selection
theme = gr.themes.Base(
    primary_hue="emerald",
)

# Function to upload image to Gemini and get roasted text
def upload_to_gemini(path, mime_type="image/jpeg"):
    file = genai.upload_file(path, mime_type=mime_type)
    return file

def generate_roast(image_path):
    # Upload the image to Gemini and get the text
    uploaded_file = upload_to_gemini(image_path)
    generation_config = {
        "temperature": 1,
        "top_p": 0.95,
        "top_k": 40,
        "max_output_tokens": 8192,
        "response_mime_type": "text/plain",
    }
    model = genai.GenerativeModel(
        model_name="gemini-1.5-flash-002",
        generation_config=generation_config,
        system_instruction="You are a professional satirist and fashion expert. You will be given a profile picture. Your duty is to roast in less than 50 words whatever is given to you in the funniest way possible!",
    )
    
    chat_session = model.start_chat(
        history=[{"role": "user", "parts": [uploaded_file]}]
    )
    response = chat_session.send_message("Roast this image!")
    return response.text

# Function to convert text to speech with Play.ht
def text_to_speech(text):
    url = "https://api.play.ht/api/v2/tts/stream"
    payload = {
        "voice": "s3://voice-cloning-zero-shot/d9ff78ba-d016-47f6-b0ef-dd630f59414e/female-cs/manifest.json",
        "output_format": "mp3",
        "text": text,
    }
    headers = {
        "accept": "audio/mpeg",
        "content-type": "application/json",
        "Authorization": API_KEY,
        "X-User-ID": USER_ID
    }
    
    response = requests.post(url, json=payload, headers=headers)
    if response.status_code == 200:
        audio_path = "output_audio.mp3"
        with open(audio_path, "wb") as audio_file:
            audio_file.write(response.content)
        return audio_path
    else:
        print(f"Error: {response.status_code} - {response.text}")
        return None  # Return None if there's an error

# Function to create video from image and audio
def generate_video(image_path, audio_path):
    if audio_path is None:
        return None  # Skip video generation if there's no valid audio file

    image_clip = ImageClip(image_path).set_duration(AudioFileClip(audio_path).duration)
    audio_clip = AudioFileClip(audio_path)
    video_clip = image_clip.set_audio(audio_clip)
    video_output_path = "output_video.mp4"
    video_clip.write_videofile(video_output_path, codec="libx264", audio_codec="aac")
    return video_output_path

# Gradio Interface
with gr.Blocks(theme=theme) as demo:
    gr.Markdown("# Image to Text-to-Speech Roasting App")
    gr.Markdown("Upload an image, and the AI will roast it, convert the roast to audio, and create a video output.")
    
    with gr.Row():
        with gr.Column():
            image_input = gr.Image(type="filepath", label="Upload Image")
        with gr.Column():
            output_text = gr.Textbox(label="Roast Text")
            audio_output = gr.Audio(label="Roast Audio")
            video_output = gr.Video(label="Roast Video")
    
    def process_image(image):
        roast_text = generate_roast(image)
        audio_path = text_to_speech(roast_text)
        
        if audio_path is None:
            return "Error generating audio. Please try again.", None, None
        
        video_path = generate_video(image, audio_path)
        return roast_text, audio_path, video_path
    
    submit_button = gr.Button("Generate Roast")
    submit_button.click(process_image, inputs=image_input, outputs=[output_text, audio_output, video_output])

# Launch the app
demo.launch(debug=True)