import gradio as gr
from transformers.agents import CodeAgent, ReactCodeAgent, ReactJsonAgent, load_tool
import random

# Custom tools since some built-ins are flaky
class WittyQuoteTool:
    def __init__(self):
        self.name = "witty_quote"
        self.description = "Spouts a random witty ML quote! 😂"
        self.inputs = {}
        self.output_type = str

    def __call__(self):
        quotes = ["'AI is just spicy statistics' - Unknown 🤓", "'The only limit is your GPU budget' - Me 💸"]
        return random.choice(quotes)

class MLDevTool:
    def __init__(self):
        self.name = "ml_dev_helper"
        self.description = "Helps with ML dev tasks like coding or debugging! 🛠️"
        self.inputs = {"task": {"type": "str", "description": "What ML task to assist with"}}
        self.output_type = str

    def __call__(self, task):
        return f"🧠 Beep boop! Helping with: {task}. Tip: Always check your data first! 📊"

# Initialize 9 agents with stable tools
tools = [WittyQuoteTool(), MLDevTool()]  # Dropped problematic built-ins
try:
    tools.append(load_tool("speech_to_text"))  # Safer built-in option
except Exception as e:
    print(f"Skipping speech_to_text due to: {e}")

agents = {
    "CodeWizard": CodeAgent(tools=tools, system_prompt="You're a coding genius! 🧙‍♂️ Write concise code."),
    "StepByStep": ReactCodeAgent(tools=tools, system_prompt="Think step-by-step, oh wise one! 🤖"),
    "JsonJester": ReactJsonAgent(tools=tools, system_prompt="JSON all the things! 😂"),
    "DataDiva": ReactCodeAgent(tools=tools, system_prompt="Clean that data, superstar! 🌟"),
    "ModelMaestro": CodeAgent(tools=tools, system_prompt="Tune models like a pro! 🎻"),
    "DebugDemon": ReactJsonAgent(tools=tools, system_prompt="Squash bugs with flair! 👾"),
    "DocDynamo": ReactCodeAgent(tools=tools, system_prompt="Write docs that dazzle! 📜"),
    "VizVirtuoso": CodeAgent(tools=tools, system_prompt="Visualize like a boss! 📈"),
    "DeployDaredevil": ReactJsonAgent(tools=tools, system_prompt="Deploy with guts! 🚀"),
}

def run_ml_team(task):
    outputs = []
    for name, agent in agents.items():
        try:
            if isinstance(agent, CodeAgent):
                result = agent.run(task, return_generated_code=True)
                outputs.append(f"🎉 {name} says: Here's my code magic! ```python\n{result}\n```")
            else:
                result = agent.run(task)
                outputs.append(f"🔥 {name} reacts: {result}")
        except Exception as e:
            outputs.append(f"😱 {name} tripped over a wire: {str(e)}")
    return "\n\n".join(outputs)

# Gradio UI
with gr.Blocks(title="ML Agent Party 🎉") as demo:
    gr.Markdown("# ML Agent Party! 🤖🎈\n9 Agents to Boost Your ML Game with Sass & Smarts!")
    task_input = gr.Textbox(label="What's Your ML Quest? 🗡️", placeholder="E.g., 'Clean my data!'")
    output = gr.Markdown(label="Agent Squad Response 🎤")
    submit_btn = gr.Button("Unleash the Crew! 🚀")
    submit_btn.click(run_ml_team, inputs=task_input, outputs=output)
    gr.Markdown("---\nBuilt with ❤️ & 😂 by xAI's Grok & Transformers Agents")

demo.launch()