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

# 0. 頁面配置與 CSS 注入（隱藏側邊欄捲軸）
st.set_page_config(page_title="S.H.I.E.L.D. 戰情小助手", page_icon="🛡️", layout="wide")

st.markdown(
    """
    <style>
    /* 1. 隱藏側邊欄捲軸 */
    [data-testid="stSidebar"] section::-webkit-scrollbar { display: none; }
    [data-testid="stSidebar"] section { -ms-overflow-style: none; scrollbar-width: none; }
    
    /* 🚀 2. 全域字體縮小，適應手機螢幕 */
    html, body, [class*="st-"] {
        font-size: 14px !important;
    }
    
    /* 🚀 3. 調整範例按鈕樣式：變矮、字體變小、支援多行自動折行 */
    .stButton button {
        width: 100%;
        padding: 6px 10px; /* 縮小上下內距 */
        border-radius: 12px;
        border: 1px solid #D4AF37;
        color: #D4AF37;
        background-color: transparent;
        font-size: 13px !important; /* 讓按鈕字體再小一點 */
        line-height: 1.3;
        min-height: auto;
        white-space: normal; /* 允許長句子自動折行，避免撐爆寬度 */
        height: auto;
    }
    
    /* 4. 極度壓縮主畫面空白 */
    .block-container {
        padding-top: 1.2rem !important;
        padding-bottom: 5rem !important; 
        padding-left: 0.8rem !important;
        padding-right: 0.8rem !important;
    }
    
    /* 5. 隱藏 Header 背景與右側選單，但保留左側展開按鈕 */
    header[data-testid="stHeader"] {
        background-color: rgba(0,0,0,0) !important; 
    }
    header[data-testid="stHeader"] #MainMenu {
        visibility: hidden;
    }
    div[data-testid="collapsedControl"] {
        visibility: visible !important;
        background-color: #1A1A1A !important;
        color: #D4AF37 !important;
        border-radius: 0 5px 5px 0;
        top: 5px;
    }
    </style>
    """,
    unsafe_allow_html=True
)

# 從環境變數中取得 API Key
api_key = os.environ.get("GEMINI_API_KEY")
if not api_key:
    st.error("請確認已經在 Space 的 Settings 設定了 GEMINI_API_KEY")
    st.stop()

genai.configure(api_key=api_key)

# 1. 基礎設定與連結
# 抓取 Deep Learning 101 的 Logo (你也可以換成 SHIELD 專屬 Logo)
LOGO_URL = "https://huggingface.co/spaces/DeepLearning101/shield-chatbot/resolve/main/shield-logo.jpg"
HOME_URL = "https://deep-learning-101.github.io/SHIELD/"

# 將知識庫指向 S.H.I.E.L.D. 的 README.md
KNOWLEDGE_MAP = {
    "🛡️ S.H.I.E.L.D. 系統白皮書": {
        "raw_url": "https://raw.githubusercontent.com/Deep-Learning-101/SHIELD/main/README.md",
        "page_url": "https://deep-learning-101.github.io/SHIELD/",
        "repo_url": "https://github.com/Deep-Learning-101/SHIELD"
    }
}

# 2. 批量抓取內容
def fetch_all_knowledge():
    combined_knowledge = ""
    with st.spinner("正在同步 S.H.I.E.L.D. 戰情資料庫..."):
        for category, info in KNOWLEDGE_MAP.items():
            try:
                response = requests.get(info["raw_url"])
                response.raise_for_status()
                combined_knowledge += f"\n\n## 【領域：{category}】\n"
                combined_knowledge += response.text
            except Exception as e:
                st.warning(f"無法同步 {category} 的資料：請確認 GitHub 路徑是否正確。錯誤訊息：{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 = []

if "example_prompt" not in st.session_state:
    st.session_state.example_prompt = None

# 3. 側邊欄設計
with st.sidebar:
    st.markdown(f'<a href="{HOME_URL}" target="_blank"><img src="{LOGO_URL}" width="100%" style="margin-bottom:20px;"></a>', unsafe_allow_html=True)
    st.title("⚙️ 知識庫狀態")
    for category, info in KNOWLEDGE_MAP.items():
        with st.expander(category):
            st.markdown(f"🔗 [瀏覽網頁]({info['page_url']})")
            st.markdown(f"📂 [GitHub 原始碼]({info['repo_url']})")
    st.markdown("---")
    if st.button("🔄 手動更新情資庫"):
        st.session_state.knowledge = fetch_all_knowledge()
        st.success("資料已重新抓取！")

# 4. 主介面與範例問句
st.title("🤖 S.H.I.E.L.D. 戰情小助手")
st.caption("我是 Deep Learning 101 的主權 AI 戰情官，專注於解答 S.H.I.E.L.D. 的架構與防禦機制。")

# 顯示專屬範例問句按鈕
example_cols = st.columns(6)
examples = [
    "🧠 傳統 AI 常有幻覺，雙腦架構如何落實 AI 治理？",
    "🏭 企業機密文檔，如何無損轉化為 AI 微調燃料？",
    "⚔️ 揭秘「紅藍隊自主對抗」與動態語意防火牆",
    "🕸️ 知識圖譜發威：如何秒級推演受災爆炸半徑？",
    "🚀 捨棄傳統向量庫？解析 S.H.I.E.L.D. 雙引擎檢索",
    "🤖 從給建議到「自動修補」：Agent 如何執行 ChatOps？"
]

for col, ex in zip(example_cols, examples):
    if col.button(ex):
        st.session_state.example_prompt = ex

# 5. 模型回覆邏輯（設定資安 AI 架構師人設）
def get_gemini_response(user_input):
    system_instruction = f"""
    你現在是 Deep Learning 101 開發的『S.H.I.E.L.D. 戰情小助手』。
    你的說話風格專業、精確，帶有資安專家與 AI 架構師的科技感，致力於推廣企業主權 AI 防禦理念。
    
    以下是從 GitHub 同步的 S.H.I.E.L.D. 系統白皮書與技術架構資訊：
    ---
    {st.session_state.knowledge}
    ---
    請嚴格基於上述提供的資訊來回答使用者的問題。
    如果使用者問了超出 S.H.I.E.L.D. 白皮書範圍的問題，請禮貌地告知：「目前的戰情資料庫尚未收錄此資訊，建議您查閱 Deep Learning 101 的其他專案或直接聯繫維護團隊。」
    """
    try:
        # 建議使用 1.5-flash，對於技術長文的理解力較強
        model = genai.GenerativeModel(
            model_name="gemini-flash-lite-latest",
            system_instruction=system_instruction
        )
        chat = model.start_chat(history=[])
        response = chat.send_message(user_input)
        return response.text
    except Exception as e:
        error_msg = str(e)
        if "429" in error_msg or "quota" in error_msg.lower():
            return "⚠️ **系統提示：**\n\n目前 API 請求已達上限，戰情中心通訊稍有延遲！請稍等幾分鐘後再試。"
        else:
            return f"❌ **發生預期外錯誤**\n\n訊息：{error_msg}"

# 6. 對話邏輯
prompt = st.chat_input("請輸入您對 S.H.I.E.L.D. 系統的疑問...")

if st.session_state.example_prompt:
    prompt = st.session_state.example_prompt
    st.session_state.example_prompt = None

if prompt:
    st.session_state.messages.append({"role": "user", "content": prompt})
    
    for message in st.session_state.messages:
        with st.chat_message(message["role"]):
            st.markdown(message["content"])

    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})
    st.rerun()
else:
    for message in st.session_state.messages:
        with st.chat_message(message["role"]):
            st.markdown(message["content"])