import os
import streamlit as st
from groq import Groq

# =====================================================
# Page Config
# =====================================================

st.set_page_config(
    page_title="Zainab's AI Assistant",
    page_icon="🚀",
    layout="wide",
)

# =====================================================
# Setup Groq Client
# =====================================================

API_KEY = os.environ.get("GROQ_API_KEY")
client = Groq(api_key=API_KEY) if API_KEY else None

AVAILABLE_MODELS = {
    "Llama 3.3 70B": "llama-3.3-70b-versatile",
    "Llama 3.1 70B": "llama-3.1-70b-versatile",
    "Llama 3.1 8B": "llama-3.1-8b-instant",
    "Mixtral 8x7B": "mixtral-8x7b-32768",
    "Gemma 2 9B": "gemma2-9b-it",
}

# =====================================================
# Sidebar
# =====================================================

with st.sidebar:
    st.title("⚙️ Settings")

    model = st.selectbox(
        "Model",
        list(AVAILABLE_MODELS.keys()),
        index=0,
    )

    temperature = st.slider("Temperature", 0.0, 1.0, 0.7)

    max_tokens = st.slider("Max Tokens", 256, 4096, 1024)

    system_prompt = st.text_area(
        "System Prompt",
        "You are a helpful, professional AI assistant.",
        height=90,
    )

    st.divider()

    if client:
        st.success("🟢 API Connected")
    else:
        st.error("🔴 GROQ_API_KEY not found")

    if st.button("🗑️ Clear Chat"):
        st.session_state.messages = []

# =====================================================
# Session State
# =====================================================

if "messages" not in st.session_state:
    st.session_state.messages = []

# =====================================================
# Header
# =====================================================

st.markdown(
    """
    <h1 style="text-align:center;"> Zainab's AI Assistant </h1>
    <p style="text-align:center;">Ask Anything</p>
    """,
    unsafe_allow_html=True,
)

# =====================================================
# Chat History
# =====================================================

for msg in st.session_state.messages:
    with st.chat_message(msg["role"]):
        st.markdown(msg["content"])

# =====================================================
# Input
# =====================================================

prompt = st.chat_input("Type your message...")

# =====================================================
# Chat Logic
# =====================================================

if prompt:

    if not client:
        st.error("Add GROQ_API_KEY in Hugging Face Secrets")
    else:

        clean_prompt = prompt.strip()

        st.session_state.messages.append({
            "role": "user",
            "content": clean_prompt,
        })

        with st.chat_message("user"):
            st.markdown(clean_prompt)

        # Build Groq payload safely
        payload = [{
            "role": "system",
            "content": system_prompt.strip()
        }]

        for m in st.session_state.messages:
            if (
                isinstance(m, dict)
                and isinstance(m.get("role"), str)
                and isinstance(m.get("content"), str)
            ):
                payload.append({
                    "role": m["role"],
                    "content": m["content"],
                })

        model_id = AVAILABLE_MODELS[model]

        # ---- Streaming response ----
        with st.chat_message("assistant"):
            placeholder = st.empty()
            full = ""

            stream = client.chat.completions.create(
                model=model_id,
                messages=payload,
                temperature=temperature,
                max_tokens=max_tokens,
                stream=True,
            )

            for chunk in stream:
                delta = chunk.choices[0].delta
                if delta and delta.content:
                    full += delta.content
                    placeholder.markdown(full)

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