import streamlit as st
from crewai import Agent, Task, Crew, Process
from crewai_tools import GPTNeoTool, CoquiTTSTool, AudioCraftTool, StableDiffusionTool, MoviePyTool, WebSocketTool, StreamlitTool

# Import your CrewAI setup
# Assuming the code you provided is in a file named crew_setup.py
from crew_setup import comedy_writer, music_composer, video_producer, live_manager

def generate_comedy_script(theme, comedy_type, characters):
    task = Task(
        description=f"Generate a 5-minute comedy script based on theme: {theme}, comedy type: {comedy_type}, and characters: {characters}",
        agent=comedy_writer
    )
    return Crew(agents=[comedy_writer], tasks=[task], process=Process.sequential).kickoff()

def generate_children_song(music_theme):
    task = Task(
        description=f"Create a children's song based on the theme: {music_theme}",
        agent=music_composer
    )
    return Crew(agents=[music_composer], tasks=[task], process=Process.sequential).kickoff()

def generate_video(script_or_song):
    task = Task(
        description=f"Generate a video from the provided content: {script_or_song}",
        agent=video_producer
    )
    return Crew(agents=[video_producer], tasks=[task], process=Process.sequential).kickoff()

def main():
    st.title("Comedy Show and Children's Music Generator")
    
    tab1, tab2, tab3 = st.tabs(["Comedy Script", "Children's Song", "Video Generation"])
    
    with tab1:
        st.header("Generate Comedy Script")
        theme = st.text_input("Enter the theme for the comedy script:")
        comedy_type = st.selectbox("Select comedy type:", ["Stand-up", "Skits", "Improv"])
        characters = st.text_input("Enter character names (comma-separated):")
        
        if st.button("Generate Comedy Script"):
            with st.spinner("Generating comedy script..."):
                script = generate_comedy_script(theme, comedy_type, characters)
                st.text_area("Generated Comedy Script:", value=script, height=300)
    
    with tab2:
        st.header("Generate Children's Song")
        music_theme = st.text_input("Enter the theme for the children's song:")
        
        if st.button("Generate Children's Song"):
            with st.spinner("Generating children's song..."):
                song = generate_children_song(music_theme)
                st.text_area("Generated Children's Song:", value=song, height=300)
                # Here you would typically also play the generated audio
                st.audio(song['audio_file'], format='audio/wav')
    
    with tab3:
        st.header("Generate Video")
        content_type = st.radio("Select content type:", ["Comedy Script", "Children's Song"])
        content = st.text_area("Enter the script or song lyrics:")
        
        if st.button("Generate Video"):
            with st.spinner("Generating video..."):
                video = generate_video(content)
                st.video(video)

if __name__ == "__main__":
    main()