# app.py
import os
import streamlit as st
from groq import Groq

# Initialize Groq client
client = Groq(api_key=os.getenv("GROQ_API_KEY"))

# ---------------- UI Setup ----------------
st.set_page_config(page_title="Advanced Groq Chatbot", layout="centered")
st.title("🤖 Groq Chatbot")
st.markdown("Chat using lightning-fast open-source LLMs on Groq!")

# Model selector (friendly names)
model_options = {
    "💬 Fast Chat (LLaMA 3 8B)": "llama3-8b-8192",
    "🧠 Deep Reasoning (LLaMA 3 70B)": "llama3-70b-8192",
    "📊 Math Solver (Mixtral)": "mixtral-8x7b-32768"
}
model_choice = st.selectbox("Choose assistant type:", list(model_options.keys()))
model_id = model_options[model_choice]

# Personality selector
style_options = {
    "Friendly": "You are a friendly assistant who replies casually.",
    "Professional": "You are a formal and helpful assistant.",
    "Technical": "You are a technical assistant providing accurate information."
}
personality = st.selectbox("Choose assistant personality:", list(style_options.keys()))
system_prompt = style_options[personality]

# Initialize session history
if "messages" not in st.session_state:
    st.session_state.messages = [{"role": "system", "content": system_prompt}]

# Clear chat button
if st.button("🔄 Clear Chat"):
    st.session_state.messages = [{"role": "system", "content": system_prompt}]
    st.rerun()

# Show chat history
for msg in st.session_state.messages[1:]:
    with st.chat_message(msg["role"]):
        st.markdown(msg["content"])

# Chat input
if prompt := st.chat_input("Type your message..."):
    st.chat_message("user").markdown(prompt)
    st.session_state.messages.append({"role": "user", "content": prompt})

    # Call Groq with streaming
    try:
        response = client.chat.completions.create(
            model=model_id,
            messages=st.session_state.messages,
            temperature=0.7,
            stream=True
        )

        reply_text = ""
        with st.chat_message("assistant"):
            reply_box = st.empty()
            for chunk in response:
                content = chunk.choices[0].delta.content or ""
                reply_text += content
                reply_box.markdown(reply_text)

        st.session_state.messages.append({"role": "assistant", "content": reply_text})

    except Exception as e:
        st.error(f"❌ Error: {e}")