import os
from crewai import Crew

from agents import CustomAgents
from tasks import CustomTasks

class CustomCrew:
    def __init__(self, var1, var2):
        self.var1 = var1
        self.var2 = var2

    def run(self):
        agents = CustomAgents()
        tasks = CustomTasks()

        # Define agents
        custom_agent_1 = agents.agent_1_name()
        custom_agent_2 = agents.agent_2_name()

        # Define tasks
        custom_task_1 = tasks.task_1_name(custom_agent_1, self.var1, self.var2)
        custom_task_2 = tasks.task_2_name(custom_agent_2)

        # Create and run the crew
        crew = Crew(
            agents=[custom_agent_1, custom_agent_2],
            tasks=[custom_task_1, custom_task_2],
            verbose=True,
        )
        result = crew.kickoff()
        return result

########################################
# GRADIO INTERFACE
########################################

import gradio as gr

def run_crew(var1, var2, history):
    # Append a user message
    history.append({"role": "user", "content": f"Variable 1: {var1}\nVariable 2: {var2}"})
    
    # Run the crew and ensure result is a string
    custom_crew = CustomCrew(var1, var2)
    result = custom_crew.run()
    if result is None:
        result = "No result was returned."
    else:
        result = str(result)
    
    # Append the result as an assistant message
    history.append({"role": "assistant", "content": result})
    return history, history

with gr.Blocks(css="footer {visibility: hidden}") as demo:
    gr.Markdown("# Crew AI Gradio Interface")
    gr.Markdown("Enter your two variables below and click 'Run Crew' to execute.")

    with gr.Row():
        var1_input = gr.Textbox(label="Variable 1", placeholder="Enter the first variable...")
        var2_input = gr.Textbox(label="Variable 2", placeholder="Enter the second variable...")

    run_button = gr.Button("Run Crew")

    chatbot = gr.Chatbot(label="Conversation", type="messages", height=300)
    clear_button = gr.Button("Clear")

    state = gr.State([])

    run_button.click(
        run_crew,
        inputs=[var1_input, var2_input, state],
        outputs=[chatbot, state],
    )

    clear_button.click(lambda: [], None, chatbot, queue=False)
    clear_button.click(lambda: [], None, state, queue=False)

if __name__ == "__main__":
    demo.launch()