import os
import groq
from dotenv import load_dotenv
import gradio as gr


# 1. Environment and Groq Client Setup

load_dotenv()  # Load environment variables from .env

# Retrieve the API key and model identifier from environment variables
GROQ_API_KEY = os.getenv("GROQ_API_KEY")
GROQ_MODEL_ID = os.getenv("GROQ_MODEL_ID", "llama-3.3-70b-versatile")

# Initialize the Groq client (adjust this based on your actual Groq API)
groq_client = groq.Client(api_key=GROQ_API_KEY)


# 2. Define a Simple Agent Class Using Groq

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

    def build_prompt(self, query):
        # Construct a prompt that provides context for the agent
        prompt = (
            f"Agent Name: {self.name}\n"
            f"Role: {self.role}\n"
            f"Instructions: {self.instructions}\n"
            f"User Query: {query}\n"
            "Response:"
        )
        return prompt

    def respond(self, query):
        prompt = self.build_prompt(query)
        # Call Groq Cloud API to generate a response (adjust parameters as needed)
        response = groq_client.generate(
            model_id=GROQ_MODEL_ID,
            prompt=prompt,
            max_tokens=200
        )
        # Assume the API returns a dict with a 'text' field for the generated response
        return response.get("text", "")


# 3. Define Specific Agents

websearch_agent = Agent(
    name="WebSearch Agent",
    role="Searches the web for financial information and the latest news.",
    instructions="Return results with sources in APA format."
)

youtube_agent = Agent(
    name="YouTube Agent",
    role="Analyzes YouTube videos that are less than 22 minutes in duration.",
    instructions="Provide concise and engaging answers based on the video content."
)

finance_agent = Agent(
    name="Finance Agent",
    role="Analyzes stock prices, analyst recommendations, and fundamentals.",
    instructions="Format your response using markdown tables where applicable."
)


# 4. Build a Multi-Agent Coordinator

class MultiAgent:
    def __init__(self, agents):
        self.agents = agents

    def handle_query(self, query):
        # For the given query, ask each agent and collect their responses
        responses = {}
        for agent in self.agents:
            responses[agent.name] = agent.respond(query)
        return responses

# Create a multi-agent system with the defined agents
multi_agent = MultiAgent([websearch_agent, youtube_agent, finance_agent])


# 5. Gradio Chatbot Front End Function

def chat_response(user_query):
    responses = multi_agent.handle_query(user_query)
    # Format the responses into a single markdown string for display
    result = ""
    for agent_name, response in responses.items():
         result += f"**{agent_name}**:\n{response}\n\n"
    return result


# 6. Create the Gradio Interface

demo = gr.Interface(
    fn=chat_response,
    inputs=gr.Textbox(placeholder="Enter your message here...", lines=2),
    outputs="markdown",
    title="Groq Agent Chatbot",
    description="A chatbot powered by Groq Cloud for GPT inference with multiple agents."
)

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