import gradio as gr
import requests
import os
from dotenv import load_dotenv
# Load .env (for local testing)
load_dotenv()

# Get API Key
GROQ_API_KEY = os.getenv("GROQ_API_KEY")
if not GROQ_API_KEY:
    raise ValueError("Missing GROQ_API_KEY")

# Groq API Config
GROQ_API_URL = "https://api.groq.com/openai/v1/chat/completions"
MODEL = "llama3-8b-8192"

headers = {
    "Authorization": f"Bearer {GROQ_API_KEY.strip()}",
    "Content-Type": "application/json"
}

# Core Chat Function with updated system prompt
def echo(message, history):
    messages = [
        {
            "role": "system",
            "content": (
                "You are an AI tutor specializing in large language models (LLMs). "
                "Always give short, clear, beginner-friendly answers. "
                "When a user refers to 'RAG', always interpret it as 'Retrieval-Augmented Generation' in the context of LLMs, "
                "not as the traffic-light system (Red, Amber, Green). "
                "Avoid jargon, and break complex ideas into bite-sized explanations. "
                "If the user asks a follow-up question, respond with full awareness of the prior conversation."
            )
        }
    ]

    # Add chat history
    for user_msg, bot_msg in history:
        messages.append({"role": "user", "content": user_msg})
        messages.append({"role": "assistant", "content": bot_msg})
    
    # Append current message
    messages.append({"role": "user", "content": message})

    try:
        response = requests.post(
            GROQ_API_URL,
            headers=headers,
            json={
                "model": MODEL,
                "messages": messages
            }
        )

        if response.status_code == 200:
            reply = response.json()["choices"][0]["message"]["content"]
            return reply
        else:
            return f"Error {response.status_code}: {response.text}"
    except Exception as e:
        return f"Error: {str(e)}"

# Gradio Chat UI
demo = gr.ChatInterface(
    fn=echo,
    title="LLM Mentor",
    examples=["What is an LLM?", "What is RAG in LLMs?", "Explain NLP simply"],
    theme="default"
)

# For Hugging Face compatibility
demo.launch(ssr_mode=False)