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

# Define the agents using open-source models

# 1. Comedy Script Generation Agent using GPT-Neo (open-source model)
comedy_writer = Agent(
    role="Comedy Show Scriptwriter",
    tools=[GPTNeoTool()],
    prompt=(
        "You are a comedy scriptwriter. Write a hilarious comedy show script. "
        "Each script should be around 5 minutes long, full of humor, and interactive."
        "\n- Theme: {theme}"
        "\n- Comedy Type: {comedy_type} (e.g., stand-up, skits, improv)"
        "\n- Characters: {characters}"
        "\n- Story: Write a funny and engaging plot that the audience will enjoy."
        "\n- Final Output: A 5-minute comedy script with multiple scenes."
    )
)

# 2. Children's Music Composer Agent using Coqui TTS and AudioCraft (open-source models)
music_composer = Agent(
    role="Children's Song Composer",
    tools=[CoquiTTSTool(), AudioCraftTool()],
    prompt=(
        "You are a children's music composer. Create a catchy and fun song for kids based on the given theme."
        "\n- Theme: {music_theme} (e.g., Baby Shark, Nursery Rhymes)"
        "\n- Lyrics: Write simple, repetitive, and easy-to-remember lyrics that kids will love."
        "\n- Tune: Create a playful and upbeat melody using instruments like piano, xylophone, and drums."
        "\n- Voices: Generate character voices using Coqui TTS in a friendly, fun tone."
        "\n- Final Output: A custom children's song with lyrics, vocals, and music."
    )
)

# 3. Video Generation Agent using Stable Diffusion for image generation and DeepFloyd IF for video interpolation
video_producer = Agent(
    role="Comedy and Music Video Producer",
    tools=[StableDiffusionTool(), DeepFloydIFTool()],
    prompt=(
        "You are a video producer. Create animated videos based on the provided comedy script or children's song."
        "\n- Script/Song: {script_or_song}"
        "\n- Generate key frames using Stable Diffusion for each scene in the script or song."
        "\n- Use DeepFloyd IF to interpolate between the frames to create smooth video transitions."
        "\n- Final Output: A short animated video with visuals synchronized to the script or song."
    )
)



# Comedy script generation task
generate_comedy_task = Task(
    description="Generate a 5-minute comedy script based on user input.",
    expected_output="A comedy script in text format, ready for video production.",
    agent=comedy_writer
)

# Children's music generation task
generate_music_task = Task(
    description="Create a children's song based on the provided theme.",
    expected_output="A song with lyrics, vocals, and background music.",
    agent=music_composer
)

# Video production task
generate_video_task = Task(
    description="Generate a video from the comedy script or children's song using open-source image and video generation tools.",
    expected_output="A short animated video synchronized to the script or song.",
    agent=video_producer
)