import streamlit as st
import google.generativeai as genai
from datetime import datetime
import os

class GeminiChatApp:
    def __init__(self):
        # 設置頁面配置
        st.set_page_config(
            page_title="Gemini AI 聊天助手",
            page_icon="💬",
            layout="wide"
        )
        
        # 設置API
        self.API_KEY = "AIzaSyAgej7dbfnw880SZvz-qsNMdFfD4W_Kv5U"  # 建議使用st.secrets或環境變數存放API密鑰
        
        # 初始化session state來存儲對話歷史
        if "messages" not in st.session_state:
            st.session_state.messages = []
            
        if "chat" not in st.session_state:
            self.initialize_gemini()
            
        # 創建UI元素
        self.create_ui()

    def initialize_gemini(self):
        """初始化Gemini API和模型"""
        try:
            # 配置API
            genai.configure(api_key=self.API_KEY)

            # 配置生成模型參數
            generation_config = {
                'temperature': 0.7,
                'top_p': 1.0,
                'max_output_tokens': 2048
            }

            # 安全設定
            safety_settings = [
                {
                    "category": "HARM_CATEGORY_HARASSMENT",
                    "threshold": "BLOCK_MEDIUM_AND_ABOVE"
                },
                {
                    "category": "HARM_CATEGORY_HATE_SPEECH",
                    "threshold": "BLOCK_MEDIUM_AND_ABOVE"
                },
                {
                    "category": "HARM_CATEGORY_SEXUALLY_EXPLICIT",
                    "threshold": "BLOCK_MEDIUM_AND_ABOVE"
                },
                {
                    "category": "HARM_CATEGORY_DANGEROUS_CONTENT",
                    "threshold": "BLOCK_MEDIUM_AND_ABOVE"
                }
            ]

            # 嘗試初始化不同的模型版本
            try:
                model = genai.GenerativeModel(
                    model_name="gemini-1.5-pro",
                    generation_config=generation_config,
                    safety_settings=safety_settings
                )
            except Exception:
                try:
                    model = genai.GenerativeModel(
                        model_name="gemini-pro",
                        generation_config=generation_config,
                        safety_settings=safety_settings,
                        api_version="v1"
                    )
                except Exception:
                    model = genai.GenerativeModel(
                        model_name="gemini-pro",
                        generation_config=generation_config,
                        safety_settings=safety_settings
                    )

            # 創建聊天會話
            st.session_state.chat = model.start_chat(history=[])
            st.session_state.model = model
            
            # 顯示歡迎訊息
            if len(st.session_state.messages) == 0:
                st.session_state.messages.append({"role": "assistant", "content": "歡迎使用Gemini AI聊天助手!\n請在下方輸入框輸入訊息並按下發送按鈕或Enter鍵。"})

        except Exception as e:
            st.error(f"無法初始化Gemini API: {str(e)}\n請檢查API金鑰和網絡連接。")

    def create_ui(self):
        """創建Streamlit UI元素"""
        # 頁面標題
        st.title("Gemini AI 聊天助手")
        
        # 側邊欄
        with st.sidebar:
            st.header("選項")
            
            # 清除聊天按鈕
            if st.button("清除對話", key="clear"):
                self.clear_chat()
                
            # 下載聊天按鈕
            if st.button("下載對話", key="download"):
                self.download_chat()
                
            # 顯示版本資訊
            try:
                version_info = f"Using google-generativeai {genai.__version__}"
            except:
                version_info = "Gemini API"
            st.caption(version_info)
        
        # 顯示聊天訊息
        self.display_messages()
        
        # 使用者輸入區域
        user_input = st.text_area("輸入您的訊息:", key="user_input", height=100)
        
        # 發送按鈕
        if st.button("發送", key="send") or user_input and st.session_state.get("_last_user_input") != user_input:
            if user_input:
                st.session_state["_last_user_input"] = user_input
                self.send_message(user_input)
    
    def display_messages(self):
        """顯示對話訊息"""
        messages_container = st.container()
        
        with messages_container:
            for message in st.session_state.messages:
                if message["role"] == "user":
                    st.markdown(f"**您**: {message['content']}")
                else:
                    st.markdown(f"**Gemini**: {message['content']}")
                st.divider()
    
    def send_message(self, user_message):
        """發送用戶輸入的訊息至Gemini API"""
        if not user_message.strip():
            return
        
        # 添加用戶訊息到對話歷史
        st.session_state.messages.append({"role": "user", "content": user_message})
        
        # 顯示思考中訊息
        thinking_placeholder = st.empty()
        thinking_placeholder.markdown("**Gemini**: 思考中...")
        
        try:
            # 獲取Gemini回應
            response = st.session_state.chat.send_message(user_message)
            response_text = response.text
            
            # 添加回應到對話歷史
            st.session_state.messages.append({"role": "assistant", "content": response_text})
            
            # 清除思考中訊息
            thinking_placeholder.empty()
            
            # 重新加載頁面以顯示新訊息
            st.rerun()
            
        except Exception as e:
            error_message = f"發生錯誤: {str(e)}"
            thinking_placeholder.empty()
            st.error(error_message)
    
    def clear_chat(self):
        """清除聊天歷史"""
        st.session_state.messages = []
        st.session_state.messages.append({"role": "assistant", "content": "聊天已清除，開始新的對話!"})
        
        # 重新初始化聊天會話
        try:
            st.session_state.chat = st.session_state.model.start_chat(history=[])
        except Exception as e:
            st.error(f"無法重新初始化聊天會話: {str(e)}")
        
        st.rerun()
    
    def download_chat(self):
        """下載聊天內容"""
        if not st.session_state.messages:
            st.warning("聊天內容為空，無需下載。")
            return
        
        # 創建時間戳記檔名
        timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
        filename = f"gemini_chat_{timestamp}.txt"
        
        # 準備聊天內容
        chat_content = ""
        for message in st.session_state.messages:
            if message["role"] == "user":
                chat_content += f"您: {message['content']}\n\n"
            else:
                chat_content += f"Gemini: {message['content']}\n\n"
        
        # 提供下載連結
        st.download_button(
            label="點擊下載對話",
            data=chat_content,
            file_name=filename,
            mime="text/plain",
        )

if __name__ == "__main__":
    app = GeminiChatApp()