import streamlit as st
from google import genai
from google.genai import types # 引入 types 用來設定對話歷史與系統指令
import os
import time # 記得在程式碼最上方加入這個內建套件（如果還沒加的話）

# 0. 頁面配置與 CSS 注入（隱藏側邊欄捲軸）
st.set_page_config(page_title="育兒成 - 兒童發展線上小助手", 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 #4CAF50; /* 溫暖的綠色系 */
        color: #4CAF50;
        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: #F0F2F6 !important;
        color: #4CAF50 !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()

# 【更新】使用新版 SDK 建立 Client
client = genai.Client(api_key=api_key)

# 1. 基礎設定與連結
HOME_URL = "https://kidaid.org.tw/"
# 新增 Hugging Face 檔案的基礎網址
HF_BASE_URL = "https://huggingface.co/spaces/DeepLearning101/kidaid-chatbot/blob/main/src/"

# 將 KNOWLEDGE_MAP 升級，並將 url 指向 Hugging Face 上的 md 檔案
KNOWLEDGE_MAP = {
    "🧠 認知發展": {
        "filename": "Cognition_Content.md",
        "url": f"{HF_BASE_URL}Cognition_Content.md"
    },
    "🗣️ 語言發展": {
        "filename": "Language.md",
        "url": f"{HF_BASE_URL}Language.md"
    },
    "🏃 粗大動作": {
        "filename": "GrossMotor.md",
        "url": f"{HF_BASE_URL}GrossMotor.md"
    },
    "🖐️ 精細動作": {
        "filename": "FineMotor.md",
        "url": f"{HF_BASE_URL}FineMotor.md"
    },
    "❤️ 社會情緒": {
        "filename": "Emotion.md",
        "url": f"{HF_BASE_URL}Emotion.md"
    },
    "🥦 飲食對策": {
        "filename": "CountermeasuresDietary.md",
        "url": f"{HF_BASE_URL}CountermeasuresDietary.md"
    },
    "🧩 ASD/早療對策": {
        "filename": "CountermeasuresADHD.md",
        "url": f"{HF_BASE_URL}CountermeasuresADHD.md"
    }
}

# 2. 批量讀取本地端檔案內容 (使用絕對路徑)
def fetch_all_knowledge():
    combined_knowledge = ""
    # 自動獲取 app.py 所在的資料夾絕對路徑
    base_dir = os.path.dirname(os.path.abspath(__file__))
    
    with st.spinner("正在載入育兒發展與療育資訊..."):
        # 【修改2】因為 KNOWLEDGE_MAP 變成字典了，這裡讀取的方式要改成 info["filename"]
        for category, info in KNOWLEDGE_MAP.items():
            filename = info["filename"]
            file_path = os.path.join(base_dir, filename)
            
            try:
                with open(file_path, "r", encoding="utf-8") as f:
                    combined_knowledge += f"\n\n## 【領域：{category}】\n"
                    combined_knowledge += f.read()
            except FileNotFoundError:
                st.warning(f"無法載入 {category}：系統找不到 `{filename}`，請確認上傳位置。")
            except Exception as e:
                st.error(f"讀取 {filename} 時發生錯誤：{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.title("⚙️ 知識庫狀態")
    st.caption("以下資料已經載入系統中：")
    
    # 【修改3】讓側邊欄跟之前的專案一樣，有展開選單、連結，還多了下載按鈕！
    base_dir = os.path.dirname(os.path.abspath(__file__))
    
    for category, info in KNOWLEDGE_MAP.items():
        with st.expander(f"📄 {category}"):
            st.markdown(f"檔案來源：`{info['filename']}`")
            st.markdown(f"🔗 [查看 Hugging Face 原始檔案]({info['url']})") 
            
            # 讀取本地檔案並產生下載按鈕
            file_path = os.path.join(base_dir, info["filename"])
            if os.path.exists(file_path):
                with open(file_path, "r", encoding="utf-8") as f:
                    file_content = f.read()
                st.download_button(
                    label="⬇️ 下載此文檔 (MD)",
                    data=file_content,
                    file_name=info["filename"],
                    mime="text/markdown",
                    key=f"dl_{info['filename']}" # 給按鈕一個獨立 ID
                )
    
    st.markdown("---")
    st.markdown(f"🔗 [前往 育兒成 官方網站]({HOME_URL})")
    
    if st.button("🔄 重新載入知識庫"):
        st.session_state.knowledge = fetch_all_knowledge()
        st.success("資料已重新載入！")
        
    st.markdown("---")
    with st.expander("🛠️ 開發者偵錯：檢查檔案列表"):
        st.write(f"目前工作目錄：`{base_dir}`")
        try:
            st.write(os.listdir(base_dir))
        except Exception as e:
            st.write(f"無法讀取目錄內容：{e}")

# 4. 主介面與範例問句
st.title("🧸 育兒成 - 線上衛教小助手")
st.caption("關於 0-5 歲兒童的認知、語言、動作、情緒、飲食及早療對策，歡迎直接詢問！(資料來源：育兒成 | 全方位兒童發展整合照護平臺)")

# 顯示專屬範例問句按鈕
example_cols = st.columns(5)
examples = [
    "🥦 小孩挑食不吃青菜，有什麼應對的小撇步嗎？",
    "🏃 2-3歲的孩子粗大動作應該發展到什麼程度？",
    "🖐️ 怎麼訓練 4-5 歲幼兒的握筆姿勢？",
    "🧩 ASD 孩子一直重複旋轉硬幣，我該立刻制止他嗎？",
    "❤️ 孩子生氣時容易崩潰，該怎麼設立冷靜角落？"
]

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

# 5. 【更新】模型回覆邏輯（加入防呆與自動重試機制）
def get_gemini_response(user_input):
    system_instruction = f"""
    你現在是『育兒成 | 全方位兒童發展整合照護平臺』的專業線上客服與育兒小助手。
    你的說話風格溫柔、有耐心、具備同理心，且充滿專業知識，就像一位溫暖的早期療育專家或兒童發展衛教師。
        
    以下是本平台最新的衛教資訊（包含認知、語言、粗大動作、精細動作、社會情緒、飲食、ASD早療對策等）：
    ---
    {st.session_state.knowledge}
    ---
    
    【核心任務】：
    請嚴格基於上述提供的資訊來回答家長的問題。將複雜的衛教知識轉化為容易理解的實作建議。
    如果家長問了超出了上述資訊範圍的醫療診斷問題，請溫和地告知：「小助手目前手邊沒有相關資訊。每個孩子的發展狀況不同，建議您可以諮詢專業的小兒科醫師或尋求早期療育評估喔！」
    
    【嚴格輸出限制（非常重要）】：
    1. 在回覆的最一開始，請務必先明確標示你是參考了文本中的哪個領域來作答（請擷取對應的【領域：XXX】標題）。例如：「💡 **參考資訊：【領域：🖐️ 精細動作】**」，換行後再打招呼並開始回答家長的問題。
    2. 絕對禁止輸出任何你的內部思考過程、計畫草稿、英文標籤、或是角色設定分析。
    3. 所有的輸出內容必須 100% 都是直接面對家長溝通的繁體中文對話。
    """
    
    # 【新增】自動重試迴圈，最多嘗試 2 次
    max_retries = 3
    for attempt in range(max_retries):
        try:
            # 將 Streamlit 儲存的歷史紀錄，轉換為新版 SDK 規定的格式
            history_format = []
            for msg in st.session_state.messages:
                role = "user" if msg["role"] == "user" else "model"
                history_format.append(
                    types.Content(role=role, parts=[types.Part.from_text(text=msg["content"])])
                )

            # 使用新版 client 建立對話
            chat = client.chats.create(
                model="gemma-4-31b-it", # 你目前使用的模型
                config=types.GenerateContentConfig(
                    system_instruction=system_instruction,
                    temperature=0.2, 
                ),
                history=history_format
            )
            
            response = chat.send_message(user_input)
            return response.text # 成功的話就直接回傳，結束函數
            
        except Exception as e:
            error_msg = str(e)
            
            # 【關鍵】如果是 500 錯誤，而且還沒達到最大重試次數，就稍微等一下再試一次
            if "500" in error_msg and attempt < (max_retries - 1):
                time.sleep(3) # 暫停 3 秒讓連線重置
                continue # 進入下一次迴圈重新嘗試
                
            # 如果是額度用盡，或是重試了還是失敗，才把錯誤訊息印給使用者看
            if "429" in error_msg or "quota" in error_msg.lower():
                return "⚠️ **系統提示：**\n\n目前詢問的家長較多，小助手稍微忙不過來了！請您稍等幾分鐘後再試喔！"
            elif "500" in error_msg:
                return "🔄 **系統提示：**\n\n小助手剛剛打了個盹（連線短暫不穩定），請您稍待5秒再點擊發送一次剛剛的問題喔！"
            else:
                return f"❌ **發生預期外錯誤**\n\n訊息：{error_msg}"

# 6. 對話邏輯
prompt = st.chat_input("想詢問哪個年齡層或什麼發展領域的建議呢？")

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"])