import gradio as gr
from openai import OpenAI
from groq import Groq
import os
import time

# Load API keys
OPENAI_API_KEY = "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"
GROQ_API_KEY = os.getenv("GROQ_API_KEY")  # Retrieved from Spaces Secrets

# Setup clients
openai_client = OpenAI(api_key=OPENAI_API_KEY)
groq_client = Groq(api_key=GROQ_API_KEY)

# Define API query functions with error handling
def query_openai(query):
    try:
        if not query or not isinstance(query, str):
            return "Error: Please provide a valid query"
        for attempt in range(3):  # Retry up to 3 times for rate limits
            try:
                response = openai_client.chat.completions.create(
                    model="gpt-3.5-turbo",  # Using gpt-3.5-turbo as tested
                    messages=[
                        {"role": "system", "content": "You are an expert consultant in cybersecurity named OpenAI Thor. Introduce yourself as Thor when first responding and not in the later responses so you do it only one. In order to help your users, it is best to get some insights on the problem, organisation etc or get a context."},
                        {"role": "user", "content": query}
                    ],
                    max_tokens=100
                )
                return response.choices[0].message.content
            except Exception as e:
                if "429" in str(e) and attempt < 2:  # Rate limit error, wait and retry
                    time.sleep(2 ** attempt)  # Exponential backoff: 1s, 2s
                    continue
                return f"OpenAI Error: {str(e)}"  # Return error if retries fail or other issue
    except Exception as e:
        return f"OpenAI Error: {str(e)}"

def query_groq(query):
    try:
        if not query or not isinstance(query, str):
            return "Error: Please provide a valid query"
        chat_completion = groq_client.chat.completions.create(
            messages=[
                {"role": "system", "content": "You are an expert consultant in cybersecurity named Groq Thor. In order to help your users, it is best to get some insights on the problem, organisation etc or get a context."},
                {"role": "user", "content": query}
            ],
            model="llama3-8b-8192"
        )
        return chat_completion.choices[0].message.content
    except Exception as e:
        return f"Groq Error: {str(e)}"

# Unified chatbot function
def chatbot(query, provider):
    if provider == "OpenAI":
        return query_openai(query)
    elif provider == "Groq":
        return query_groq(query)
    else:
        return "Invalid provider selected."

# Gradio UI setup
interface = gr.Interface(
    fn=chatbot,
    inputs=[
        gr.Textbox(lines=2, placeholder="Enter your question here...", label="Query"),
        gr.Dropdown(choices=["OpenAI", "Groq"], label="Provider")
    ],
    outputs=gr.Textbox(label="Chatbot Response"),
    title="Cybersecurity Chatbot",
    description="Ask cybersecurity questions to Thor, powered by OpenAI or Groq. Select a provider and type your query."
)

# Launch the app
interface.launch()