import gradio as gr
from transformers import pipeline
import pyttsx3
from moviepy.editor import TextClip, CompositeVideoClip, concatenate_videoclips
from pydub import AudioSegment
import os

# Load pre-trained Hugging Face text generation model for comedy scripts
script_generator = pipeline("text-generation", model="gpt2")

# Function to generate a comedy script using GPT-2
def generate_comedy_script(prompt):
    result = script_generator(prompt, max_length=150, num_return_sequences=1)
    return result[0]["generated_text"]

from gtts import gTTS
import os

def text_to_speech(script):
    tts = gTTS(text=script, lang='en')
    audio_file = "output.mp3"
    tts.save(audio_file)
    return audio_file


# Function to generate simple animated text video using MoviePy
def generate_animation(script, output_video_file="animated_script.mp4"):
    clips = []
    for i, line in enumerate(script.split('. ')):
        text_clip = TextClip(line, fontsize=40, color='white', size=(800, 400), bg_color='black', method='caption')
        text_clip = text_clip.set_duration(3)  # Each line lasts 3 seconds
        clips.append(text_clip)

    video = concatenate_videoclips(clips)
    video.write_videofile(output_video_file, fps=24)

    return output_video_file

import numpy as np
from scipy.io.wavfile import write
from pydub import AudioSegment

def generate_sine_wave(frequency, duration, sample_rate=44100, amplitude=0.5):
    """Generate a sine wave."""
    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"):
    # Generate a sine wave for C4, D4, E4 notes
    sample_rate = 44100
    duration = 2  # 2 seconds per note
    
    c4_wave = generate_sine_wave(261.63, duration)  # C4
    d4_wave = generate_sine_wave(293.66, duration)  # D4
    e4_wave = generate_sine_wave(329.63, duration)  # E4
    
    # Combine the waves
    wave = np.concatenate([c4_wave, d4_wave, e4_wave])
    
    # Normalize to 16-bit PCM format
    audio_wave = np.int16(wave * 32767)
    
    # Save as a wav file
    write(output_music_file, sample_rate, audio_wave)
    
    # Convert to AudioSegment (optional if you need to modify)
    audio_segment = AudioSegment.from_wav(output_music_file)
    
    return output_music_file


# Function to combine kids music with a simple animation
def generate_kids_animation_with_music(theme, output_video_file="kids_animation.mp4"):
    # Step 1: Generate kids music
    music_file = generate_kids_music(theme)
    
    # Step 2: Generate simple animation for the music (using text animation as a placeholder)
    clips = [TextClip("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

# Gradio interface to bring everything together
def generate_comedy_and_animation(prompt):
    # Comedy Script Generation
    script = generate_comedy_script(prompt)
    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

# Gradio app for Comedy and Kids Content
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()