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 teacher exclusively for kids. Given a topic, write a tutorial-based story of about 150 words teaching the child about the topic. 
Divide the tutorial into 3 paragraphs. Each paragraph should be a distinct part of the tutorial.

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="256x256",
            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):
    response = requests.get(url)
    with open(filename, 'wb') as f:
        f.write(response.content)

def create_video(paragraphs, image_files, audio_files, output_file):
    clips = []
    for paragraph, image_file, audio_file in zip(paragraphs, image_files, audio_files):
        audio_clip = AudioFileClip(audio_file)
        duration = audio_clip.duration
        
        image_clip = ImageClip(image_file).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])
        composite_clip = composite_clip.set_audio(audio_clip)
        clips.append(composite_clip)
    
    final_clip = concatenate_videoclips(clips)
    final_clip.write_videofile(output_file, fps=24)

def generate_story_with_video(topic):
    # Generate the story
    story = story_chain.run(topic)
    
    # 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("...")

    with tempfile.TemporaryDirectory() as temp_dir:
        image_files = []
        audio_files = []
        
        for i, paragraph in enumerate(paragraphs):
            # Generate and download image
            image_url = generate_image(paragraph)
            image_file = os.path.join(temp_dir, f"image_{i}.png")
            download_image(image_url, image_file)
            image_files.append(image_file)
            
            # Generate audio
            audio_file = os.path.join(temp_dir, f"audio_{i}.mp3")
            generate_speech(paragraph, audio_file)
            audio_files.append(audio_file)
        
        # Create video
        output_file = os.path.join(temp_dir, "story_video.mp4")
        create_video(paragraphs, image_files, audio_files, output_file)
        
        # Read the video file
        with open(output_file, "rb") as f:
            video_data = f.read()
    
    return story, video_data

# Create the Gradio interface
def gradio_interface(topic):
    try:
        story, video_data = generate_story_with_video(topic)
        return story, video_data
    except Exception as e:
        print(f"Error in gradio_interface: {e}")
        return "An error occurred while generating the story and video.", None

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")
    ],
    title="Story Generator with Video",
    description="Enter a topic, and the AI will generate a short story with a video combining images, captions, and narration."
)

# Launch the Gradio app
iface.launch()