from crewai import Agent, Task, Crew, Process
import gradio as gr
from langchain_groq import ChatGroq  # Assuming this is your custom Groq integration
import os
import dotenv

# Load API keys from .env
dotenv.load_dotenv()

# Retrieve API key from .env file
api_key = os.getenv("Groq_API_KEY")

# Initialize Groq-powered LLaMA model using the API
llama_model = ChatGroq(
    api_key=api_key,
    model='llama-3.2-3b-preview'  # Assuming this is the correct model name
)

# Define Plot Creator Agent
plot_creator_agent = Agent(
    role="Plot Creator",
    goal="""Create a compelling fictional plot based on the user's input about the genre, characters, and setting.""",
    backstory="""You are a master storyteller, capable of weaving intricate plots for any genre. You take user inputs and craft the foundation of the story.""",
    verbose=True,
    allow_delegation=False,
    llm=llama_model  # Use the Groq LLaMA model for generating the plot
)

# Define Story Narrator Agent
story_narrator_agent = Agent(
    role="Story Narrator",
    goal="""Write a full fictional story based on the plot provided by the Plot Creator Agent. 
            The story should include vivid descriptions, engaging dialogue, and follow the narrative structure.""",
    backstory="""You are a gifted writer who turns plot outlines into rich and immersive narratives, breathing life into the characters and setting.""",
    verbose=True,
    allow_delegation=False,
    llm=llama_model  # Use the Groq LLaMA model for generating the story
)

# Function to generate a fictional story based on user input
def generate_story(genre, characters, setting):
    # Create tasks for the plot and story writing
    plot_creation_task = Task(
        description=f"Create a plot for a fictional story in the '{genre}' genre with the following characters: {characters}, "
                    f"and the story is set in: {setting}.",
        agent=plot_creator_agent,
        expected_output="A detailed plot with the main conflict, setting, and character arcs."
    )

    story_narration_task = Task(
        description=f"Write a fictional story based on the generated plot. "
                    f"Include rich narrative elements such as character development, conflict resolution, and engaging dialogue.",
        agent=story_narrator_agent,
        expected_output="A 3-5 paragraph fictional story that fully immerses the reader in the characters and plot."
    )

    # Create a crew that orchestrates the agents and tasks
    crew = Crew(
        agents=[plot_creator_agent, story_narrator_agent],
        tasks=[plot_creation_task, story_narration_task],
        verbose=True,
        process=Process.sequential  # First generate plot, then write story
    )

    # Execute the tasks and return the final story output
    output = crew.kickoff(inputs={'genre': genre, 'characters': characters, 'setting': setting})
    return output.raw  # Return the generated story

# Gradio Interface for Fictional Story Writer
iface = gr.Interface(
    fn=generate_story,  # Function to call when user submits input
    inputs=[
        gr.Textbox(label="Genre", placeholder="e.g., fantasy, sci-fi"),
        gr.Textbox(label="Characters (comma-separated)", placeholder="e.g., knight, dragon, wizard"),
        gr.Textbox(label="Setting", placeholder="e.g., a magical forest, outer space")
    ],
    outputs="text",  # Output is a text response (generated story)
    title="Fictional Story Writer",
    description="Generate a fictional story based on your input of genre, characters, and setting."
)

# Launch the Gradio app
iface.launch()