import streamlit as st
import google.generativeai as genai
import requests
import os

st.set_page_config(page_title="AI 新知小助手", page_icon="📚", layout="wide")

# 改從環境變數中取得 API Key
api_key = os.environ.get("GEMINI_API_KEY")

if not api_key:
    st.error("請確認已經在 Space 的 Settings -> Variables and secrets 中設定了 GEMINI_API_KEY")
    st.stop()

genai.configure(api_key=api_key)

# 1. 定義你要鎖定的 GitHub 檔案清單
# 請把 YOUR_ACCOUNT/YOUR_REPO 換成你真實的 GitHub 資訊
BASE_URL = "https://raw.githubusercontent.com/Deep-Learning-101/deep-learning-101.github.io/main/"
TARGET_FILES = {
    "大型語言模型 (LLM)": f"{BASE_URL}Large-Language-Model.md",
    "自然語言處理 (NLP)": f"{BASE_URL}Natural-Language-Processing.md",
    "語音處理 (Speech)": f"{BASE_URL}Speech-Processing.md",
    "電腦視覺 (CV)": f"{BASE_URL}Computer-Vision.md"
}

# 2. 批量抓取並格式化內容
def fetch_all_knowledge():
    combined_knowledge = ""
    with st.spinner("正在從 GitHub 同步 4 大領域最新資訊..."):
        for category, url in TARGET_FILES.items():
            try:
                response = requests.get(url)
                response.raise_for_status()
                # 在每個檔案內容前加上明確的標題，幫助 AI 分類記憶
                combined_knowledge += f"\n\n## 【領域：{category}】\n"
                combined_knowledge += response.text
            except Exception as e:
                st.warning(f"無法同步 {category} 的資料：{e}")
    return combined_knowledge

# 初始化 Session State
if "knowledge" not in st.session_state:
    st.session_state.knowledge = fetch_all_knowledge()

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

# 側邊欄設計
with st.sidebar:
    st.title("⚙️ 知識庫狀態")
    st.write("目前收錄以下領域：")
    for category in TARGET_FILES.keys():
        st.markdown(f"- {category}")
        
    st.markdown("---")
    if st.button("🔄 手動更新知識庫"):
        st.session_state.knowledge = fetch_all_knowledge()
        st.success("所有領域資料已重新抓取！")

# 主介面
st.title("📚 AI 演算法與論文社群助手")
st.caption("知識庫涵蓋 LLM、NLP、Speech、CV。歡迎直接提問！")

def get_gemini_response(user_input):
    system_instruction = f"""
    你是一位專業的 AI 技術分析專家。
    以下是我從 GitHub 同步的四個領域 (LLM, NLP, Speech, CV) 的最新技術簡介與重點：
    ---
    {st.session_state.knowledge}
    ---
    請嚴格基於上述提供的資訊來回答使用者的問題。
    如果使用者詢問了上述資訊中沒有涵蓋的細節，請明確告知：「目前這份懶人包中尚未收錄關於此問題的詳細資訊」。
    """
    model = genai.GenerativeModel(
        model_name="gemini-1.5-pro",
        system_instruction=system_instruction
    )
    
    chat = model.start_chat(history=[])
    response = chat.send_message(user_input)
    return response.text

# 顯示對話紀錄
for message in st.session_state.messages:
    with st.chat_message(message["role"]):
        st.markdown(message["content"])

# 接收輸入
if prompt := st.chat_input("例如：請幫我比較一下 CV 跟 NLP 目前常用的技術？"):
    st.session_state.messages.append({"role": "user", "content": prompt})
    with st.chat_message("user"):
        st.markdown(prompt)

    with st.chat_message("assistant"):
        response_text = get_gemini_response(prompt)
        st.markdown(response_text)
    
    st.session_state.messages.append({"role": "assistant", "content": response_text})