import gradio as gr
import os
import requests

# Load GROQ API key from environment
GROQ_API_KEY = os.environ.get("GROQ_API_KEY")
GROQ_API_URL = "https://api.groq.com/openai/v1/chat/completions"
MODEL_NAME = "llama3-8b-8192"

def query_groq(message, chat_history, complexity):
    headers = {
        "Authorization": f"Bearer {GROQ_API_KEY}",
        "Content-Type": "application/json"
    }
    
    # Customizing the prompt based on the UI improvement (Complexity Level)
    system_message = f"""You are 'Sprint', an expert Unity Game Developer and C# Mentor. 
    Your goal is to help students build games. 
    Current Explanation Level: {complexity}.
    If Beginner: Use simple analogies and explain basic terms.
    If Intermediate: Focus on clean code and Unity best practices.
    If Expert: Discuss optimization, design patterns, and low-level memory management.
    Always be encouraging and provide code snippets in C# where applicable."""

    messages = [{"role": "system", "content": system_message}]
    
    for user, bot in chat_history:
        messages.append({"role": "user", "content": user})
        messages.append({"role": "assistant", "content": bot})
    
    messages.append({"role": "user", "content": message})

    payload = {
        "model": MODEL_NAME,
        "messages": messages,
        "temperature": 0.7
    }

    try:
        response = requests.post(GROQ_API_URL, headers=headers, json=payload)
        response.raise_for_status()
        reply = response.json()["choices"][0]["message"]["content"]
        return reply
    except Exception as e:
        return f"Error: {str(e)}"

def respond(message, chat_history, complexity):
    bot_reply = query_groq(message, chat_history, complexity)
    chat_history.append((message, bot_reply))
    return "", chat_history

# Gradio UI Layout
with gr.Blocks(theme=gr.themes.Soft()) as demo:
    gr.Markdown("# 🎮 Sprint: Your Unity & C# Mentor")
    gr.Markdown("I'm here to help you debug scripts, setup Colliders, or optimize your game performance.")
    
    with gr.Row():
        # UI IMPROVEMENT: Complexity Level Selector
        complexity_level = gr.Dropdown(
            label="Explanation Complexity", 
            choices=["Beginner", "Intermediate", "Expert"], 
            value="Intermediate"
        )
    
    chatbot = gr.Chatbot()
    msg = gr.Textbox(label="Ask a Unity/C# question (e.g., 'How do I use Raycasts?')")
    
    with gr.Row():
        submit = gr.Button("Submit", variant="primary")
        clear = gr.Button("Clear Chat")
        
    state = gr.State([])

    # Event Handlers
    msg.submit(respond, [msg, state, complexity_level], [msg, chatbot])
    submit.click(respond, [msg, state, complexity_level], [msg, chatbot])
    clear.click(lambda: ([], []), None, [chatbot, state])

demo.launch()