import os
import gradio as gr
from groq import Groq

# ===============================
# 1. Groq client
# ===============================
client = Groq(api_key=os.environ.get("GROQ_API_KEY"))

def groq_chat(query, history):
    messages = [{"role": "system", "content": "You are a warm, empathetic psychologist bot. Speak naturally, human-like, and supportive."}]
    for user, bot in history:
        if user: messages.append({"role": "user", "content": user})
        if bot: messages.append({"role": "assistant", "content": bot})
    messages.append({"role": "user", "content": query})

    response = client.chat.completions.create(
        model="llama-3.3-70b-versatile",
        messages=messages
    )
    return response.choices[0].message.content.strip()

# ===============================
# 2. User sessions + flow
# ===============================
user_sessions = {}

def get_name_meaning(name):
    return "a unique and beautiful soul 🌸"

def chatbot(user_input, history, user_id="default"):
    session = user_sessions.get(user_id, {"stage": 0, "name": None})
    response = ""

    if session["stage"] == 0:
        response = (
            "🌟 Hey friend! I’m really happy you’re here 🤝 "
            "I want to know you better—what’s your beautiful name?"
        )
        session["stage"] = 1

    elif session["stage"] == 1 and not session["name"]:
        session["name"] = user_input.strip().capitalize()
        meaning = get_name_meaning(session["name"])
        response = (
            f"✨ That’s wonderful! **{session['name']}** means *{meaning}*.\n\n"
            "🤗 You’re safe here with me. This is your space to open up.\n\n"
            "Tell me—how are you truly feeling today?"
        )
        session["stage"] = 2

    elif session["stage"] == 2:
        response = (
            f"💭 Thanks for sharing that, {session['name']}.\n"
            "What’s been weighing most on your heart or mind recently?"
        )
        session["stage"] = 3

    elif session["stage"] == 3:
        response = (
            f"🤝 I hear you, {session['name']}. Your feelings really matter 🌈\n"
            "If you could change one thing in your life right now, what would it be?"
        )
        session["stage"] = 4

    else:
        groq_response = groq_chat(user_input, history)
        response = groq_response if groq_response else (
            f"🌟 Remember {session['name']}, you’re not alone here. "
            "I’m here to listen and walk beside you 💕"
        )

    user_sessions[user_id] = session
    return response

# ===============================
# 3. Gradio Interface (Warm + Visible)
# ===============================
with gr.Blocks(css="""
    body {background-color: #EAF4F8;}
    .gradio-container {color: #2D2D2D; font-family: 'Segoe UI', sans-serif;}
    h2 {color: #4E9A8E; text-align: center; font-weight: bold;}
    .message.user {
        background-color: #FFFFFF !important;
        color: #000000 !important;
        border-radius: 16px; 
        box-shadow: 0px 4px 10px rgba(0,0,0,0.12);
        font-size: 15px;
        line-height: 1.6;
        font-weight: 600;
    }
    .message.bot {
        background-color: #B8A2D1 !important;
        color: #2E3A3A !important;
        border-radius: 16px; 
        box-shadow: 0px 4px 10px rgba(0,0,0,0.12);
        font-size: 15px;
        line-height: 1.6;
        font-weight: 600;
    }
    button {
        background-color: #4E9A8E !important;
        color: white !important;
        border-radius: 12px !important;
        font-weight: bold !important;
        padding: 8px 16px !important;
    }
    textarea, input {
        background-color: #FFFFFF !important;
        color: #000000 !important;
        border: 1px solid #CCC !important;
        border-radius: 10px !important;
        padding: 10px !important;
        font-size: 14px;
    }
""") as demo:

    gr.Markdown("<h2>🌿 MindEase: Your Safe Companion 🌸</h2>")

    chatbot_ui = gr.Chatbot(height=500, bubble_full_width=False, show_label=False)

    with gr.Row():
        msg = gr.Textbox(placeholder="Type here and press Enter...", scale=10)
        clear = gr.Button("Clear")

    def respond(message, history):
        bot_message = chatbot(message, history, user_id="user1")
        history.append((message, bot_message))
        return "", history

    msg.submit(respond, [msg, chatbot_ui], [msg, chatbot_ui])
    clear.click(lambda: None, None, chatbot_ui, queue=False)

# 🚀 Launch
demo.launch()