import gradio as gr
from transformers import pipeline
from diffusers import StableDiffusionPipeline
import torch
from PIL import Image, ImageDraw, ImageFont
import scipy.io.wavfile
from TTS.api import TTS
from moviepy.editor import CompositeVideoClip, ImageClip, AudioFileClip, concatenate_videoclips
import os


# Use DistilGPT-2 for text generation
script_generator = pipeline("text-generation", model="distilgpt2", truncation=True, max_length=100)

# Use Stable Diffusion for image generation
image_generator = StableDiffusionPipeline.from_pretrained("stabilityai/stable-diffusion-2-1-base", torch_dtype=torch.float32).to("cpu")

# Use Coqui TTS for text-to-speech
tts = TTS(model_name="tts_models/en/ljspeech/tacotron2-DDC", progress_bar=False, gpu=False)

# Placeholder for music generation
def generate_fun_music(prompt, output_music_file="fun_music.wav"):
    # Generate silence as placeholder
    duration = 3  # seconds
    sample_rate = 44100
    silence = np.zeros(int(sample_rate * duration), dtype=np.int16)
    scipy.io.wavfile.write(output_music_file, sample_rate, silence)
    return output_music_file


def generate_comedy_script(prompt):
    script = script_generator(prompt)[0]['generated_text']
    return script


def text_to_speech(script):
    output_audio = 'output.wav'
    tts.tts_to_file(text=script, file_path=output_audio)
    return output_audio

def create_images_from_script(script):
    lines = script.split('. ')
    image_paths = []
    for i, line in enumerate(lines):
        img = image_generator(line).images[0]
        img_path = f'/tmp/image_{i}.png'
        img.save(img_path)
        image_paths.append(img_path)
    return image_paths

def generate_text_video(script):
    image_paths = create_images_from_script(script)
    clips = []
    for img_path in image_paths:
        image_clip = ImageClip(img_path).set_duration(3).set_position(('center', 'center'))
        clips.append(image_clip)
    final_video = concatenate_videoclips(clips, method="compose")
    final_video.write_videofile("/tmp/final_video.mp4", fps=24)
    return "/tmp/final_video.mp4"

def combine_audio_video(video_path, audio_path):
    video = VideoFileClip(video_path)
    audio = AudioFileClip(audio_path)
    final_video = video.set_audio(audio)
    final_video.write_videofile("/tmp/final_comedy_video.mp4", fps=24)
    return "/tmp/final_comedy_video.mp4"

def generate_comedy_and_animation(prompt):
    script = generate_comedy_script(prompt)
    audio_file = text_to_speech(script)
    video_file = generate_text_video(script)
    fun_music = generate_fun_music(prompt)
    final_video = combine_audio_video(video_file, fun_music)
    return script, audio_file, final_video