import gradio as gr
from transformers import pipeline
from gtts import gTTS
from moviepy.editor import TextClip, CompositeVideoClip, concatenate_videoclips
import numpy as np
from scipy.io.wavfile import write
from pydub import AudioSegment

script_generator = pipeline("text-generation", model="gpt2", truncation=True, max_length=100)


def text_to_speech(script):
    if isinstance(script, list):
        texts = [item['generated_text'] for item in script if isinstance(item, dict) and 'generated_text' in item]
        script = " ".join(texts)
        print("Final script for TTS:", script)  # Debugging line

    if not script.strip():
        raise ValueError("No text to convert to speech.")
        
    tts = gTTS(text=script, lang='en')
    audio_file = 'output.mp3'
    tts.save(audio_file)
    return audio_file


def generate_animation(script):
    lines = script.split('\n')
    clips = []
    for line in lines:
        text_clip = TextClip(line, fontsize=40, color='white', size=(800, 400), bg_color='black', method='caption')
        text_clip = text_clip.set_duration(3)
        clips.append(text_clip)
    final_video = concatenate_videoclips(clips)
    return final_video

def generate_sine_wave(frequency, duration, sample_rate=44100, amplitude=0.5):
    t = np.linspace(0, duration, int(sample_rate * duration), False)
    wave = amplitude * np.sin(2 * np.pi * frequency * t)
    return wave

def generate_kids_music(theme, output_music_file="kids_music.wav"):
    sample_rate = 44100
    duration = 2
    c4_wave = generate_sine_wave(261.63, duration)
    d4_wave = generate_sine_wave(293.66, duration)
    e4_wave = generate_sine_wave(329.63, duration)
    wave = np.concatenate([c4_wave, d4_wave, e4_wave])
    audio_wave = np.int16(wave * 32767)
    write(output_music_file, sample_rate, audio_wave)
    return output_music_file

def generate_kids_animation_with_music(theme, output_video_file="kids_animation.mp4"):
    music_file = generate_kids_music(theme)
    clips = [TextClip(f"Kids Music: {theme}", fontsize=70, color='yellow', size=(800, 400), bg_color='blue', method='caption').set_duration(5)]
    video = CompositeVideoClip(clips)
    video.write_videofile(output_video_file, fps=24)
    return output_video_file, music_file

def generate_comedy_and_animation(prompt):
    script = script_generator(prompt)
    print("Generated script:", script)  # Debugging line
    audio_file = text_to_speech(script)
    video_file = generate_animation(script)
    return script, audio_file, video_file

def generate_kids_content(theme):
    video_file, music_file = generate_kids_animation_with_music(theme)
    return music_file, video_file

with gr.Blocks() as app:
    gr.Markdown("## AI Comedy and Kids Content Generator")

    with gr.Tab("Generate Comedy Animation"):
        prompt_input = gr.Textbox(label="Comedy Prompt")
        generate_btn = gr.Button("Generate Comedy Script and Animation")
        comedy_script = gr.Textbox(label="Generated Script")
        comedy_audio = gr.Audio(label="Generated Audio")
        comedy_video = gr.Video(label="Generated Animation")

        generate_btn.click(
            generate_comedy_and_animation,
            inputs=prompt_input,
            outputs=[comedy_script, comedy_audio, comedy_video]
        )

    with gr.Tab("Generate Kids Music Animation"):
        theme_input = gr.Textbox(label="Kids Music Theme")
        generate_music_btn = gr.Button("Generate Kids Music and Animation")
        kids_music_audio = gr.Audio(label="Generated Music")
        kids_music_video = gr.Video(label="Generated Kids Animation")

        generate_music_btn.click(
            generate_kids_content,
            inputs=theme_input,
            outputs=[kids_music_audio, kids_music_video]
        )

app.launch()