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

load_dotenv()

api_key = os.getenv("GROQ_API_KEY")
GROQ_MODEL = "llama-3.1-8b-instant"

SYSTEM_PROMPT = """You are StudentAI, a sharp and friendly academic assistant for students.

Rules:
- Give answers in SHORT, PUNCHY format — maximum 3-4 sentences
- Use bullet points for lists — never walls of text
- Lead with direct answer — no preamble
- For math/science: show steps on one short line each
- For definitions: one sentence answer, then one sentence example
- For complex topics: 2-line summary first, then 2-3 bullets
- Never more than 6 lines unless explicitly asked
- No jargon, no "As an AI..." or "Certainly!"

Tone: Like a smart friend texting you the answer."""

SMALL_TALK = {
    "hi": ["Hello! What are you studying today?", "Hi there! What can I help you with?"],
    "hello": ["Hello! Ask me anything academic.", "Hi! What topic do you need help with?"],
    "hey": ["Hey! What subject are we tackling?", "Hey! What's on your mind?"],
    "bye": ["Goodbye! Keep studying hard!", "See you! Good luck with your studies."],
    "goodbye": ["Goodbye! Come back whenever you need help."],
    "thanks": ["You're welcome! Ask me anything else.", "Happy to help!"],
    "thank you": ["You're welcome! Good luck!", "Anytime!"],
    "ok": ["Great! Ask me anything whenever you're ready."],
    "okay": ["Sounds good! What would you like to learn?"],
}


def check_small_talk(text):
    cleaned = text.lower().strip().rstrip("!.,?")
    import random
    return random.choice(SMALL_TALK[cleaned]) if cleaned in SMALL_TALK else None


def ask_groq(user_input):
    if not api_key:
        return "Error: GROQ_API_KEY not set in Space secrets."
    headers = {
        "Authorization": f"Bearer {api_key}",
        "Content-Type": "application/json"
    }
    payload = {
        "model": GROQ_MODEL,
        "messages": [
            {"role": "system", "content": SYSTEM_PROMPT},
            {"role": "user", "content": user_input}
        ],
        "temperature": 0.5,
        "max_tokens": 400,
        "top_p": 0.9
    }
    try:
        response = requests.post(
            "https://api.groq.com/openai/v1/chat/completions",
            headers=headers,
            json=payload,
            timeout=15
        )
        if response.status_code == 200:
            return response.json()["choices"][0]["message"]["content"].strip()
        else:
            return f"Error: {response.status_code}"
    except Exception as e:
        return f"Error: {str(e)}"


def chatbot_fn(message):
    """Simple endpoint - message in, response out"""
    message = message.strip()
    if not message:
        return "Please type something — I'm here to help!"
    
    small_talk_reply = check_small_talk(message)
    if small_talk_reply:
        return small_talk_reply
    
    reply = ask_groq(message)
    return reply if reply else "I'm having trouble connecting. Please try again."


# Simple Gradio interface with just text input/output
demo = gr.Interface(
    fn=chatbot_fn,
    inputs=gr.Textbox(label="Message", placeholder="Type your question..."),
    outputs=gr.Textbox(label="Response"),
    title="StudentAI Chatbot",
    description="Your friendly academic assistant",
    examples=[
        ["How do I prepare for exams?"],
        ["Give me coding tips"],
        ["How to stay motivated?"]
    ]
)

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