import os
import streamlit as st
from groq import Groq

# Page configuration
st.set_page_config(
    page_title="Groq Chatbot",
    page_icon="⚡",
    layout="centered",
    initial_sidebar_state="expanded"
)

# Custom CSS for modern dark theme
st.markdown("""
<style>
    /* Main container styling */
    .main {
        background: linear-gradient(135deg, #0f0c29 0%, #302b63 50%, #24243e 100%);
        color: #ffffff;
    }
    
    /* Chat container */
    .stChatMessage {
        background: rgba(255, 255, 255, 0.05);
        border-radius: 15px;
        backdrop-filter: blur(10px);
        border: 1px solid rgba(255, 255, 255, 0.1);
        margin: 10px 0;
        padding: 15px;
    }
    
    /* User message styling */
    .stChatMessage[data-testid="stChatMessageUser"] {
        background: linear-gradient(135deg, #667eea 0%, #764ba2 100%);
    }
    
    /* Assistant message styling */
    .stChatMessage[data-testid="stChatMessageAssistant"] {
        background: rgba(255, 255, 255, 0.08);
    }
    
    /* Input box styling */
    .stTextInput > div > div > input {
        background: rgba(255, 255, 255, 0.1);
        border: 2px solid rgba(102, 126, 234, 0.5);
        border-radius: 25px;
        color: white;
        padding: 15px 20px;
        font-size: 16px;
    }
    
    .stTextInput > div > div > input:focus {
        border-color: #667eea;
        box-shadow: 0 0 15px rgba(102, 126, 234, 0.5);
    }
    
    /* Button styling */
    .stButton > button {
        background: linear-gradient(135deg, #667eea 0%, #764ba2 100%);
        color: white;
        border: none;
        border-radius: 25px;
        padding: 10px 30px;
        font-weight: bold;
        transition: all 0.3s ease;
    }
    
    .stButton > button:hover {
        transform: translateY(-2px);
        box-shadow: 0 10px 20px rgba(102, 126, 234, 0.4);
    }
    
    /* Sidebar styling */
    .css-1d391kg {
        background: rgba(0, 0, 0, 0.3);
    }
    
    /* Headers */
    h1 {
        background: linear-gradient(135deg, #667eea 0%, #764ba2 100%);
        -webkit-background-clip: text;
        -webkit-text-fill-color: transparent;
        font-weight: 800;
    }
    
    /* Scrollbar styling */
    ::-webkit-scrollbar {
        width: 10px;
    }
    
    ::-webkit-scrollbar-track {
        background: rgba(255, 255, 255, 0.05);
    }
    
    ::-webkit-scrollbar-thumb {
        background: linear-gradient(135deg, #667eea 0%, #764ba2 100%);
        border-radius: 5px;
    }
    
    /* Typing indicator animation */
    @keyframes pulse {
        0%, 100% { opacity: 1; }
        50% { opacity: 0.5; }
    }
    
    .typing-indicator {
        animation: pulse 1.5s infinite;
        color: #667eea;
    }
</style>
""", unsafe_allow_html=True)

# Initialize Groq client
@st.cache_resource
def get_groq_client():
    api_key = os.environ.get("GROQ_API_KEY")
    if not api_key:
        st.error("⚠️ GROQ_API_KEY not found in environment variables!")
        st.stop()
    return Groq(api_key=api_key)

# Initialize session state for chat history
if "messages" not in st.session_state:
    st.session_state.messages = [
        {"role": "assistant", "content": "👋 Hello! I'm your AI assistant powered by Groq. How can I help you today?"}
    ]

if "model" not in st.session_state:
    st.session_state.model = "llama-3.3-70b-versatile"

# Sidebar configuration
with st.sidebar:
    st.markdown("### ⚙️ Settings")
    
    # Model selection
    model_options = {
        "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"
    }
    
    selected_model = st.selectbox(
        "Choose Model",
        options=list(model_options.keys()),
        index=0
    )
    st.session_state.model = model_options[selected_model]
    
    st.markdown("---")
    
    # Temperature slider
    temperature = st.slider("Temperature", 0.0, 1.0, 0.7, 0.1)
    
    # Max tokens slider
    max_tokens = st.slider("Max Tokens", 256, 4096, 1024, 128)
    
    st.markdown("---")
    
    # Clear chat button
    if st.button("🗑️ Clear Chat", use_container_width=True):
        st.session_state.messages = [
            {"role": "assistant", "content": "👋 Chat cleared! How can I help you?"}
        ]
        st.rerun()
    
    st.markdown("---")
    st.markdown("### 📊 Stats")
    st.write(f"Messages: {len(st.session_state.messages)}")
    
    # API Key status
    if os.environ.get("GROQ_API_KEY"):
        st.success("🔑 API Key: Connected")
    else:
        st.error("🔑 API Key: Missing")

# Main chat interface
st.title("⚡ Groq Chatbot")
st.markdown("*Lightning-fast AI conversations*")

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

# Chat input
if prompt := st.chat_input("Type your message here..."):
    # Add user message to history
    st.session_state.messages.append({"role": "user", "content": prompt})
    
    # Display user message
    with st.chat_message("user"):
        st.markdown(prompt)
    
    # Get AI response
    with st.chat_message("assistant"):
        message_placeholder = st.empty()
        
        try:
            client = get_groq_client()
            
            # Show typing indicator
            message_placeholder.markdown('<span class="typing-indicator">🤔 Thinking...</span>', unsafe_allow_html=True)
            
            # Call Groq API
            chat_completion = client.chat.completions.create(
                messages=[
                    {"role": m["role"], "content": m["content"]}
                    for m in st.session_state.messages
                ],
                model=st.session_state.model,
                temperature=temperature,
                max_tokens=max_tokens,
                stream=False
            )
            
            response = chat_completion.choices[0].message.content
            
            # Display response with typewriter effect simulation
            message_placeholder.markdown(response)
            
            # Add to history
            st.session_state.messages.append({"role": "assistant", "content": response})
            
        except Exception as e:
            message_placeholder.error(f"❌ Error: {str(e)}")

# Footer
st.markdown("---")
st.markdown(
    "<div style='text-align: center; color: rgba(255,255,255,0.5);'>"
    "Powered by Groq ⚡ | Built with Streamlit ❤️"
    "</div>",
    unsafe_allow_html=True
)