import os
import requests
import numpy as np
import gradio as gr
from transformers import AutoTokenizer, AutoModelForCausalLM

# Load the TinyLlama model and tokenizer
tokenizer = AutoTokenizer.from_pretrained("TinyLlama/TinyLlama-1.1B-Chat-v1.0")
model = AutoModelForCausalLM.from_pretrained("TinyLlama/TinyLlama-1.1B-Chat-v1.0")

class Agent:
    def __init__(self, id, api_key=None):
        self.id = id
        self.task = None
        self.results = None
        self.api_key = api_key

    def execute_task(self):
        if self.task:
            print(f"Agent {self.id} is making an API call to '{self.task}'")
            headers = {"Authorization": f"Bearer {self.api_key}"} if self.api_key else {}
            try:
                response = requests.get(self.task, headers=headers, timeout=10)  # Add timeout
                if response.status_code == 200:
                    self.results = response.json()
                else:
                    self.results = f"Error: Unable to fetch data, status code {response.status_code}"
                print(f"Agent {self.id} received: {self.results}")
            except Exception as e:
                self.results = f"Error: {str(e)}"
                print(f"Agent {self.id} encountered an error: {str(e)}")

    def communicate(self, other_agents):
        pass

class Swarm:
    def __init__(self, num_agents, fractal_pattern, api_key=None):
        self.agents = [Agent(i, api_key) for i in range(num_agents)]
        self.fractal_pattern = fractal_pattern
        print(f"Swarm created with {num_agents} agents using the {fractal_pattern} pattern.")

    def assign_tasks(self, tasks):
        for i, task in enumerate(tasks):
            self.agents[i % len(self.agents)].task = task
            print(f"Task assigned to Agent {self.agents[i % len(self.agents)].id}: {task}")

    def execute(self):
        for agent in self.agents:
            agent.execute_task()
        for agent in self.agents:
            agent.communicate(self.agents)

    def gather_results(self):
        return [agent.results for agent in self.agents if agent.results]

def generate_tasks(api_url, num_tasks):
    return [api_url] * num_tasks

def run_swarm(api_url, api_key, num_agents, num_tasks):
    try:
        tasks = generate_tasks(api_url, num_tasks)
        print(f"Generated tasks: {tasks}")
    except Exception as e:
        return f"Error generating tasks: {str(e)}"

    try:
        swarm = Swarm(num_agents=num_agents, fractal_pattern="Pentagonal", api_key=api_key)
        swarm.assign_tasks(tasks)
        swarm.execute()
    except Exception as e:
        return f"Error executing swarm tasks: {str(e)}"

    try:
        results = swarm.gather_results()
    except Exception as e:
        return f"Error gathering results: {str(e)}"

    print("\nAll results retrieved by the swarm:")
    for i, result in enumerate(results):
        print(f"Result {i + 1}: {result}")

    return results

def gradio_interface(prompt, api_url, api_key, num_agents, num_tasks):
    if prompt == "Autobots, Assemble!":
        results = run_swarm(api_url, api_key, num_agents, num_tasks)
        return "\n".join(str(result) for result in results)
    else:
        return "Prompt not recognized. Please use 'Autobots, Assemble!' to execute the swarm."

# Default values for the inputs
default_prompt = "Autobots, Assemble!"
default_api_url = "https://meowfacts.herokuapp.com/"
default_api_key = ""
default_num_agents = 5
default_num_tasks = 2

iface = gr.Interface(
    fn=gradio_interface,
    inputs=[
        gr.Textbox(label="Prompt", placeholder="Enter the prompt", value=default_prompt),
        gr.Textbox(label="API URL", placeholder="Enter the API URL", value=default_api_url),
        gr.Textbox(label="API Key (Optional)", placeholder="Enter the API Key", value=default_api_key),
        gr.Number(label="Number of Agents", value=default_num_agents, precision=0),
        gr.Number(label="Number of Tasks", value=default_num_tasks, precision=0)
    ],
    outputs=gr.Textbox(label="Results"),
    title="Swarm Model Processing and Result Gatherer",
    description="""
        This Gradio app demonstrates a swarm of agents making API calls and gathering results.
        - When the prompt 'Autobots, Assemble!' is entered, a swarm is created and the API calls are executed.
        - Each agent makes an API call to the specified URL and retrieves data.
        - The results from all agents are gathered and displayed.
        - Enter the prompt 'Autobots, Assemble!', API URL, API Key (optional), number of agents, and number of tasks to see the process in action.
        - By default, the app uses the prompt 'Autobots, Assemble!' and the API URL 'https://meowfacts.herokuapp.com/' with 5 agents and 2 tasks.
    """
)

iface.launch()