import os
import time
import gradio as gr
from groq import Groq
import re

# Initialize Groq
client = Groq(api_key=os.getenv("GROQ_API_KEY"))
MODEL = "llama-3.3-70b-versatile"

# ------------------ Direct Answer ------------------
def answer_anything(question):
    response = client.chat.completions.create(
        model=MODEL,
        messages=[
            {"role": "system", "content": "You are a helpful assistant."},
            {"role": "user", "content": question}
        ],
        temperature=0.5
    )
    return response.choices[0].message.content

# ------------------ Autonomous Agent ------------------
class AutonomousAgent:
    def __init__(self, goal):
        self.goal = goal
        self.memory = []
        self.completed_tasks = []

    def think(self, prompt):
        response = client.chat.completions.create(
            model=MODEL,
            messages=[
                {"role": "system", "content": "You are an autonomous AI agent that works without human intervention."},
                {"role": "user", "content": prompt}
            ],
            temperature=0.5,
        )
        return response.choices[0].message.content

    def plan_tasks(self):
        prompt = f"""
        Goal: {self.goal}
        Create a step-by-step task plan to achieve this goal.
        Use short, clear tasks.
        """
        return self.think(prompt)

    def execute_task(self, task):
        memory_context = "\n".join([f"Task: {m['task']} Result: {m['result']}" for m in self.memory])
        prompt = f"""
        Based on memory:
        {memory_context}

        Execute this task fully:
        {task}

        Output the result clearly.
        """
        return self.think(prompt)

    def run(self, max_steps=6, show_steps=False):
        tasks_text = self.plan_tasks()
        tasks = [re.sub(r'^\d+\.\s*', '', t).strip() for t in tasks_text.split("\n") if t.strip()]

        results = []
        for i, task in enumerate(tasks[:max_steps], 1):
            result = self.execute_task(task)
            self.memory.append({"task": task, "result": result})
            self.completed_tasks.append(task)
            if show_steps:
                results.append(f"STEP {i}: {task}\nRESULT: {result}\n{'-'*40}")

        # Return either detailed steps or just the final result
        if show_steps:
            return "\n".join(results)
        else:
            return result  # last task result only

# ------------------ Gradio Interface ------------------
def run(user_input, mode, max_steps, show_steps):
    if mode == "Direct Answer":
        return answer_anything(user_input)
    elif mode == "Autonomous Agent":
        agent = AutonomousAgent(user_input)
        return agent.run(max_steps=max_steps, show_steps=show_steps)

interface = gr.Interface(
    fn=run,
    inputs=[
        gr.Textbox(label="Enter your question or goal"),
        gr.Radio(choices=["Direct Answer", "Autonomous Agent"], label="Mode"),
        gr.Slider(1, 10, value=6, step=1, label="Max Steps (for Autonomous Agent)"),
        gr.Checkbox(label="Show steps (for Autonomous Agent)", value=True)
    ],
    outputs="text",
    title="🤖 Combined AI Agent",
    description="Choose 'Direct Answer' for quick responses, or 'Autonomous Agent' for multi-step goal execution."
)

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