import streamlit as st
import google.generativeai as genai

# 設置頁面配置
st.set_page_config(page_title="Gemini AI 對話", page_icon="🤖")

# 添加自定義 CSS 樣式
st.markdown("""
<style>
.stApp {
    background-color: #f0f2f6;
}
.stTextInput>div>div>input {
    border-radius: 20px;
}
.stButton>button {
    border-radius: 20px;
    background-color: #4CAF50;
    color: white;
}
</style>
""", unsafe_allow_html=True)

# 設定 Gemini API
genai.configure(api_key='AIzaSyAs-lbHwPZVWMGmjKJ654KZ_9M-reyna08')

# 側邊欄設置
with st.sidebar:
    st.title("設置")
    model_name = st.selectbox("選擇模型", ["gemini-pro", "other-model"])
    temperature = st.slider("Temperature", 0.0, 1.0, 0.7)

# 主頁面標題
st.title("🤖 AI 對話助手")

# 初始化對話歷史
if 'messages' not in st.session_state:
    st.session_state.messages = []

# 清除對話歷史的按鈕
if st.button("清除對話歷史"):
    st.session_state.messages = []
    st.experimental_rerun()

# 定義消息顯示函數
def message(content, is_user=False):
    if is_user:
        st.markdown(f'<div style="background-color: #DCF8C6; padding: 10px; border-radius: 10px; margin-bottom: 10px; text-align: right;">{content}</div>', unsafe_allow_html=True)
    else:
        st.markdown(f'<div style="background-color: #FFFFFF; padding: 10px; border-radius: 10px; margin-bottom: 10px;">{content}</div>', unsafe_allow_html=True)

# 顯示對話歷史
for msg in st.session_state.messages:
    message(msg["content"], msg["role"] == "user")

# 用戶輸入
if prompt := st.chat_input("說點什麼..."):
    # 添加用戶訊息到對話歷史
    st.session_state.messages.append({"role": "user", "content": prompt})
    message(prompt, is_user=True)

    # 使用 Gemini API 生成回應
    try:
        model = genai.GenerativeModel(model_name)
        with st.spinner("AI正在思考中..."):
            full_response = ""
            for response in model.generate_content(prompt, stream=True):
                full_response += response.text
            message(full_response)
        
        # 添加 AI 回應到對話歷史
        st.session_state.messages.append({"role": "assistant", "content": full_response})
    except Exception as e:
        st.error(f"發生錯誤: {str(e)}")