import streamlit as st
import os
from groq import Groq

# -----------------------
# PAGE CONFIG
# -----------------------
st.set_page_config(
    page_title="Krish AI Chatbot 🤖",
    page_icon="🤖",
    layout="centered"
)

# -----------------------
# MOBILE RESPONSIVE UI 🔥
# -----------------------
st.markdown("""
<style>

/* Main background */
.stApp {
    background-color: #0E1117;
}

/* Container padding (mobile fix) */
.block-container {
    padding-top: 1rem;
    padding-left: 1rem;
    padding-right: 1rem;
    max-width: 100%;
}

/* Title */
.title {
    text-align: center;
    font-size: 28px;
    font-weight: bold;
    color: white;
}

/* Subtitle */
.subtitle {
    text-align: center;
    color: #AAAAAA;
    font-size: 14px;
    margin-bottom: 15px;
}

/* Chat container */
.chat-container {
    display: flex;
    flex-direction: column;
    gap: 10px;
}

/* User bubble */
.chat-row-user {
    display: flex;
    justify-content: flex-end;
}

.chat-bubble-user {
    background-color: #4CAF50;
    color: white;
    padding: 10px 14px;
    border-radius: 15px;
    font-size: 14px;
    max-width: 85%;
    word-wrap: break-word;
}

/* Bot bubble */
.chat-row-bot {
    display: flex;
    justify-content: flex-start;
}

.chat-bubble-bot {
    background-color: #2A2F3A;
    color: white;
    padding: 10px 14px;
    border-radius: 15px;
    font-size: 14px;
    max-width: 85%;
    word-wrap: break-word;
}

/* Input box fix */
textarea {
    color: white !important;
}

/* -----------------------
   MOBILE OPTIMIZATION 📱
----------------------- */
@media (max-width: 600px) {

    .title {
        font-size: 22px;
    }

    .chat-bubble-user,
    .chat-bubble-bot {
        font-size: 13px;
        padding: 8px 12px;
        max-width: 90%;
    }

    .block-container {
        padding-left: 0.5rem;
        padding-right: 0.5rem;
    }
}

</style>
""", unsafe_allow_html=True)

# -----------------------
# HEADER
# -----------------------
st.markdown('<div class="title">💬 Krish AI Chatbot</div>', unsafe_allow_html=True)
st.markdown('<div class="subtitle">Smart, Fast & Beautiful AI Chat 🚀</div>', unsafe_allow_html=True)

# -----------------------
# API KEY (Secrets + fallback)
# -----------------------
api_key = os.getenv("GROQ_API_KEY")

if not api_key:
    api_key = st.sidebar.text_input("Enter Groq API Key 🔑", type="password")

if not api_key:
    st.warning("Please enter your Groq API key")
    st.stop()

client = Groq(api_key=api_key)

# -----------------------
# MODEL
# -----------------------
MODEL = "llama-3.1-8b-instant"

# -----------------------
# SESSION STATE
# -----------------------
if "messages" not in st.session_state:
    st.session_state.messages = []

# -----------------------
# DISPLAY CHAT
# -----------------------
for msg in st.session_state.messages:
    if msg["role"] == "user":
        st.markdown(f"""
        <div class="chat-row-user">
            <div class="chat-bubble-user">{msg["content"]}</div>
        </div>
        """, unsafe_allow_html=True)
    else:
        st.markdown(f"""
        <div class="chat-row-bot">
            <div class="chat-bubble-bot">{msg["content"]}</div>
        </div>
        """, unsafe_allow_html=True)

# -----------------------
# INPUT
# -----------------------
user_input = st.chat_input("Type your message...")

if user_input:
    # Save user message
    st.session_state.messages.append({"role": "user", "content": user_input})

    # Show instantly
    st.markdown(f"""
    <div class="chat-row-user">
        <div class="chat-bubble-user">{user_input}</div>
    </div>
    """, unsafe_allow_html=True)

    # Placeholder for streaming
    placeholder = st.empty()
    full_response = ""

    try:
        stream = client.chat.completions.create(
            model=MODEL,
            messages=st.session_state.messages,
            stream=True
        )

        for chunk in stream:
            content = chunk.choices[0].delta.content or ""
            full_response += content

            placeholder.markdown(f"""
            <div class="chat-row-bot">
                <div class="chat-bubble-bot">{full_response}▌</div>
            </div>
            """, unsafe_allow_html=True)

        # Final message
        placeholder.markdown(f"""
        <div class="chat-row-bot">
            <div class="chat-bubble-bot">{full_response}</div>
        </div>
        """, unsafe_allow_html=True)

        # Save
        st.session_state.messages.append(
            {"role": "assistant", "content": full_response}
        )

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