import tkinter as tk
from tkinter import simpledialog, messagebox
import ollama
import random

class Agent:
    def __init__(self, role):
        self.role = role

    def analyze_task(self, task):
        response = ollama.chat(
            model='llama3.2-vision',
            messages=[{
                'role': 'user',
                'content': f"As a {self.role}, analyze this task: {task}"
            }]
        )
        return response['message']['content']

class Supervisor:
    def __init__(self):
        self.agents = []

    def add_agent(self, agent):
        self.agents.append(agent)

    def assign_task(self, task):
        for agent in self.agents:
            analysis = agent.analyze_task(task)
            response = ollama.chat(
                model='llama3.2-vision',
                messages=[{
                    'role': 'user',
                    'content': f"As a supervisor, is this analysis satisfactory? Analysis: {analysis}"
                }]
            )
            if "yes" in response['message']['content'].lower():
                return f"Task solved by {agent.role}"
        return "Task not solved satisfactorily"

def create_multi_agent_system():
    roles = ['Detective', 'Scientist', 'Artist', 'Engineer', 'Philosopher']
    num_agents = simpledialog.askinteger("Input", "How many agents? (2-5)", minvalue=2, maxvalue=5)
    
    supervisor = Supervisor()
    
    for _ in range(num_agents):
        role = simpledialog.askstring("Input", f"Choose a role for agent {_+1}:\n" + "\n".join(roles))
        if role in roles:
            supervisor.add_agent(Agent(role))
            roles.remove(role)
        else:
            messagebox.showerror("Error", "Invalid role selected")
            return None
    
    return supervisor

def main():
    root = tk.Tk()
    root.withdraw()
    
    system = create_multi_agent_system()
    if system:
        task = simpledialog.askstring("Input", "Enter the task:")
        result = system.assign_task(task)
        messagebox.showinfo("Result", result)

if __name__ == "__main__":
    main()
