import streamlit as st
from groq import Groq
import os

# Page config
st.set_page_config(page_title="Saeed's Chat Bot", page_icon=":robot:")

# Initialize Groq client
try:
    client = Groq(api_key=os.environ.get("GROQ_API_KEY"))
except Exception as e:
    st.error(f"Error initializing Groq client: {e}")
    st.stop()

# Header with image
st.markdown(
    """
    <h2>
        <img src="https://cdn-uploads.huggingface.co/production/uploads/69cb5c18f952e82f9367b6ba/WuYfTyW9yjam_A31AfuHV.png" width="50" style="vertical-align: middle; margin-right:10px;">
        Saeed's AI Chat Bot
    </h2>
    """,
    unsafe_allow_html=True
)

# Initialize chat history
if "messages" not in st.session_state:
    st.session_state.messages = []

# Display previous messages
for message in st.session_state.messages:
    with st.chat_message(message["role"]):
        st.markdown(message["content"])

# --- Custom Input Section (Image + Textbox + Button) ---
col1, col2, col3 = st.columns([1, 10, 2])

with col1:
    st.image(
        "https://cdn-uploads.huggingface.co/production/uploads/69cb5c18f952e82f9367b6ba/WuYfTyW9yjam_A31AfuHV.png",
        width=40
    )

with col2:
    prompt = st.text_input(
        "How Saeed's BOT can help you?",
        label_visibility="collapsed",
        placeholder="How Saeed's BOT can help you?"
    )

with col3:
    send = st.button("Send")

# --- Chat Logic ---
if send and prompt:
    # Save user message
    st.session_state.messages.append({"role": "user", "content": prompt})

    # Show user message
    with st.chat_message("user"):
        st.markdown(prompt)

    # Assistant response
    with st.chat_message("assistant"):
        message_placeholder = st.empty()
        full_response = ""

        try:
            stream = client.chat.completions.create(
                model="llama-3.1-8b-instant",
                messages=[
                    {"role": m["role"], "content": m["content"]}
                    for m in st.session_state.messages
                ],
                max_tokens=1024,
                temperature=0.7,
                stream=True,
            )

            for chunk in stream:
                full_response += (chunk.choices[0].delta.content or "")
                message_placeholder.markdown(full_response + "▌")

            message_placeholder.markdown(full_response)

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

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