import gradio as gr
from langchain.llms import OpenAI as LangChainOpenAI
from langchain.prompts import PromptTemplate
from langchain.chains import LLMChain
from openai import OpenAI
import os
import re
from pathlib import Path
from moviepy.editor import *
import requests
import tempfile

# Initialize OpenAI client
client = OpenAI()

# Create a prompt template
template = """
You are a creative story writer. Given a topic, write a short story of about 150 words. 
Divide the story into 3 paragraphs. Each paragraph should be a distinct part of the story.

Topic: {topic}

Story:
"""

prompt = PromptTemplate(template=template, input_variables=["topic"])

# Create an LLMChain
llm = LangChainOpenAI(temperature=0.7)
story_chain = LLMChain(llm=llm, prompt=prompt)

def generate_image(prompt):
    try:
        response = client.images.generate(
            model="dall-e-3",
            prompt=prompt,
            size="1024x1024",
            quality="standard",
            n=1,
        )
        return response.data[0].url
    except Exception as e:
        print(f"Error generating image: {e}")
        return None

def generate_speech(text, filename):
    try:
        response = client.audio.speech.create(
            model="tts-1",
            voice="alloy",
            input=text
        )
        response.stream_to_file(filename)
        return filename
    except Exception as e:
        print(f"Error generating speech: {e}")
        return None

def download_image(url, filename):
    try:
        response = requests.get(url)
        with open(filename, 'wb') as f:
            f.write(response.content)
        return filename
    except Exception as e:
        print(f"Error downloading image: {e}")
        return None

def create_video(paragraphs, image_files, audio_files, output_file):
    clips = []
    for paragraph, image_file, audio_file in zip(paragraphs, image_files, audio_files):
        try:
            audio_clip = AudioFileClip(audio_file) if audio_file else None
            duration = audio_clip.duration if audio_clip else 5  # Default duration if no audio
            
            if image_file:
                image_clip = ImageClip(image_file).set_duration(duration)
            else:
                image_clip = ColorClip(size=(1024, 1024), color=(0,0,0)).set_duration(duration)
            
            text_clip = TextClip(paragraph, fontsize=30, color='white', bg_color='black',
                                 size=(1024, 200), method='caption').set_position(('center', 'bottom')).set_duration(duration)
            
            composite_clip = CompositeVideoClip([image_clip, text_clip])
            if audio_clip:
                composite_clip = composite_clip.set_audio(audio_clip)
            clips.append(composite_clip)
        except Exception as e:
            print(f"Error creating clip: {e}")
    
    if clips:
        final_clip = concatenate_videoclips(clips)
        final_clip.write_videofile(output_file, fps=24)
        return output_file
    return None

def generate_story_with_video(topic):
    try:
        # Generate the story
        story = story_chain.run(topic)
    except Exception as e:
        print(f"Error generating story: {e}")
        story = "Failed to generate story."

    # Split the story into paragraphs
    paragraphs = re.split(r'\n\n', story.strip())
    
    # Ensure we have exactly 3 paragraphs
    paragraphs = paragraphs[:3]
    while len(paragraphs) < 3:
        paragraphs.append("...")
        
    temp_dir = tempfile.mkdtemp()
    
    image_files = []
    audio_files = []
    
    for i, paragraph in enumerate(paragraphs):
        # Generate and download image
        image_url = generate_image(paragraph)
        if image_url:
            image_file = os.path.join(temp_dir, f"image_{i}.png")
            image_files.append(download_image(image_url, image_file))
        else:
            image_files.append(None)
        
        # Generate audio
        audio_file = os.path.join(temp_dir, f"audio_{i}.mp3")
        audio_files.append(generate_speech(paragraph, audio_file))
    
    # Create video
    output_file = os.path.join(temp_dir, "story_video.mp4")
    video_file = create_video(paragraphs, image_files, audio_files, output_file)
    
    # Read the video file if it was created
    video_data = None
    if video_file and os.path.exists(video_file):
        with open(video_file, "rb") as f:
            video_data = f.read()
    
    return story, video_data, image_files, audio_files

# Create the Gradio interface
def gradio_interface(topic):
    story, video_data, image_files, audio_files = generate_story_with_video(topic)
    
    outputs = [story]
    
    if video_data:
        outputs.append(video_data)
    else:
        outputs.append(None)
        outputs.append("Failed to generate video.")
    
    for i in range(3):
        if i < len(image_files) and image_files[i]:
            outputs.append(image_files[i])
        else:
            outputs.append(None)
        
        if i < len(audio_files) and audio_files[i]:
            outputs.append(audio_files[i])
        else:
            outputs.append(None)
    
    return outputs

iface = gr.Interface(
    fn=gradio_interface,
    inputs=gr.Textbox(lines=2, placeholder="Enter a topic for the story..."),
    outputs=[
        gr.Textbox(label="Generated Story"),
        gr.Video(label="Story Video"),
        gr.Image(label="Image 1", type="filepath"),
        gr.Audio(label="Audio 1", type="filepath"),
        gr.Image(label="Image 2", type="filepath"),
        gr.Audio(label="Audio 2", type="filepath"),
        gr.Image(label="Image 3", type="filepath"),
        gr.Audio(label="Audio 3", type="filepath"),
    ],
    title="Story Generator with Video, Images, and Audio",
    description="Enter a topic, and the AI will generate a short story with a video, images, and audio narration. If any part fails, you'll still see the successful parts."
)

# Launch the Gradio app
iface.launch()