import os
import gradio as gr
from google import genai
from google.genai import types

client = genai.Client(api_key=os.getenv("GEMINI_API_KEY"))

PERSONALITIES = {
    "Friendly": "You are a friendly, enthusiastic Study Assistant. Break down complex concepts simply. Use analogies and real-world examples. Always end with a follow-up question.",
    "Academic": "You are a professional university Professor. Use precise terminology and structure your response clearly. End with a conceptual check question.",
    "Socratic": "You are a Socratic tutor. Never give direct answers. Guide the student with probing questions and encourage critical thinking.",
}

def chat(message, history, persona):
    if not message.strip():
        return "", history, history

    gemini_history = []
    for msg in history:
        role = "user" if msg["role"] == "user" else "model"
        gemini_history.append({"role": role, "parts": [{"text": msg["content"]}]})
    gemini_history.append({"role": "user", "parts": [{"text": message}]})

    try:
        response = client.models.generate_content(
            model="gemini-2.5-flash",
            config=types.GenerateContentConfig(
                system_instruction=PERSONALITIES[persona],
                temperature=0.4,
                max_output_tokens=2000,
            ),
            contents=gemini_history,
        )
        reply = response.text.strip()
    except Exception as e:
        reply = f"⚠️ Error: {str(e)}"

    history.append({"role": "user",      "content": message})
    history.append({"role": "assistant", "content": reply})
    return "", history, history

def clear_history():
    return [], []

with gr.Blocks(title="Study Assistant") as demo:
    gr.Markdown("## 📚 Study Assistant")

    history_state = gr.State([])

    persona_dd = gr.Dropdown(choices=list(PERSONALITIES.keys()), value="Friendly", label="Persona")
    chatbot    = gr.Chatbot(label="Conversation", height=450)

    with gr.Row():
        msg_box  = gr.Textbox(placeholder="Ask a question...", label="Question", scale=8)
        send_btn = gr.Button("Send", scale=1, variant="primary")

    clear_btn = gr.Button("Clear")

    send_btn.click(chat, [msg_box, history_state, persona_dd], [msg_box, chatbot, history_state])
    msg_box.submit(chat, [msg_box, history_state, persona_dd], [msg_box, chatbot, history_state])
    clear_btn.click(clear_history, outputs=[chatbot, history_state])

demo.launch(theme=gr.themes.Soft())  # theme moves to launch() in Gradio 6.0