import streamlit as st
from google import genai
import os
# Set GEMINI_API_KEY from GOOGLE_API_KEY if not already set
if not os.getenv('GEMINI_API_KEY'):
    os.environ["GEMINI_API_KEY"] = os.getenv("GOOGLE_API_KEY")

# initializes chat client
if 'client' not in st.session_state:
    st.session_state['client'] = genai.Client()

st.set_page_config(page_title="AI Chat Demo", page_icon="💬", layout="centered")

# Initialize chat history in session state
if "messages" not in st.session_state:
    st.session_state["messages"] = [
        {"role": "ai", "content": "👋 Welcome! How can I help you today?"}
    ]
    st.session_state["dummy_idx"] = 0


st.title("AI Chat Demo")

# Chat history display

chat_container = st.container()
with chat_container:
    for msg in st.session_state["messages"]:
        if msg["role"] == "user":
            st.markdown(f"<div style='text-align: right; color: #1a73e8;'><b>You:</b> {msg['content']}</div>", unsafe_allow_html=True)
        else:
            st.markdown(f"<div style='text-align: left; color: #444;'><b>AI:</b> {msg['content']}</div>", unsafe_allow_html=True)

# User input
with st.form(key="chat_form", clear_on_submit=True):
    user_input = st.text_input("Type your message:", "", key="input")
    submitted = st.form_submit_button("Send")

_ = st.button("Clear Chat", on_click=lambda: st.session_state.clear(), key="clear_chat")


if submitted and user_input.strip():
    # Add user message
    st.session_state["messages"].append({"role": "user", "content": user_input.strip()})

    # Prepare the content for the AI model by joining all messages with role tags
    content = " /n ".join(f'<{msg["role"]}> {msg["content"]}' for msg in st.session_state["messages"])

    # Generate AI response using the prepared content
    ai_message = st.session_state['client'].models.generate_content(
        model="gemini-2.5-flash",
        contents=content
    ).text

    st.session_state["messages"].append({"role": "ai", "content": ai_message})

    st.rerun()