import streamlit as st
import requests
import json
from typing import Optional, List, Dict, Any
import threading
from queue import Queue
import time

# 後端 API 端點
BACKEND_URL = "http://localhost:8000/generate"

# 設定頁面為寬版，適合聊天室
st.set_page_config(page_title="RAG 聊天室", layout="wide")

st.title("🤖 RAG 聊天室")

# 初始化 session state
if "messages" not in st.session_state:
    st.session_state.messages = []

# 側邊欄：設定
with st.sidebar:
    st.header("設定")
    backend_url = st.text_input("後端 URL", value=BACKEND_URL)
    
    # 新增 streaming 選項
    streaming_enabled = st.checkbox("啟用 Streaming (步驟 + 逐字回應)", value=True)
    
    st.header("元數據 (metadata)")
    metadata_json = st.text_area(
        "JSON 格式 (可選): {'key': 'value'}",
        value="{}",
        height=100
    )
    try:
        metadata: Optional[Dict[str, Any]] = json.loads(metadata_json) if metadata_json.strip() else None
    except json.JSONDecodeError:
        st.error("metadata JSON 格式錯誤，請檢查！")
        metadata = None

# 聊天歷史：自動從 messages 轉換為 chat_history 格式
chat_history = [{"role": msg["role"], "content": msg["content"]} for msg in st.session_state.messages]

# 顯示聊天歷史
for message in st.session_state.messages:
    with st.chat_message(message["role"]):
        st.markdown(message["content"])

# 輸入框：像聊天室一樣在底部
if prompt := st.chat_input("輸入您的問題..."):
    # 追加使用者訊息到歷史
    st.session_state.messages.append({"role": "user", "content": prompt})
    
    # 顯示使用者訊息 (右邊)
    with st.chat_message("user"):
        st.markdown(prompt)
    
    # 準備請求資料
    request_data = {
        "query": prompt,
        "streaming": streaming_enabled,
        "chat_history": chat_history,  # 包含最新使用者訊息
        "metadata": metadata
    }
    
    # 顯示機器人訊息 placeholder (左邊)
    with st.chat_message("assistant"):
        # 狀態 placeholder (用於步驟顯示)
        status_placeholder = st.empty()
        
        # 回應內容 placeholder
        response_placeholder = st.empty()
        
        # 非同步處理 (使用 threading)
        queue = Queue()
        
        def fetch_response():
            try:
                if streaming_enabled:
                    # Streaming=True: SSE 解析
                    response = requests.post(
                        backend_url,
                        json=request_data,
                        stream=True,
                        headers={"Content-Type": "application/json"},
                        timeout=60
                    )
                    response.raise_for_status()
                    
                    full_response = ""
                    for line in response.iter_lines(decode_unicode=True):
                        if line:
                            line_str = line.decode('utf-8') if isinstance(line, bytes) else line
                            if line_str.startswith("data: "):
                                data = line_str[6:].strip()
                                if data == "[DONE]":
                                    break
                                # 檢查是否為步驟消息
                                if data in ["開始處理查詢...", "文件檢索完成", "上下文建構完成", "生成完成"]:
                                    queue.put(("status", data))
                                else:
                                    # 假設為 LLM chunk，追加到內容
                                    full_response += data
                                    queue.put(("content", data))
                    queue.put(("end", full_response))
                
                else:
                    # Streaming=False: 等待完整 JSON
                    response = requests.post(
                        backend_url,
                        json=request_data,
                        headers={"Content-Type": "application/json"},
                        timeout=60
                    )
                    response.raise_for_status()
                    result = response.json()
                    full_response = result.get("response", "無回應")
                    queue.put(("content", full_response))
                    queue.put(("end", full_response))
                
            except requests.exceptions.RequestException as e:
                error_msg = f"錯誤: {str(e)}"
                queue.put(("error", error_msg))
                queue.put(("end", ""))
        
        # 啟動背景執行緒
        thread = threading.Thread(target=fetch_response, daemon=True)
        thread.start()
        
        # 處理 queue 中的訊息
        full_display = ""
        while True:
            try:
                msg = queue.get(timeout=0.1)
                if msg[0] == "end":
                    break
                
                msg_type, content = msg
                if msg_type == "status":
                    # 更新狀態
                    status_placeholder.markdown(f"**狀態：** {content}")
                elif msg_type == "content":
                    # 追加內容 (逐字或完整)
                    full_display += content
                    response_placeholder.markdown(full_display)
                elif msg_type == "error":
                    response_placeholder.markdown(f"**錯誤：** {content}")
                
                # 強制重新渲染
                time.sleep(0.01)
                st.rerun()
                
            except:
                pass  # 逾時繼續
        
        # 最終追加到歷史 (僅成功時)
        if full_display:
            st.session_state.messages.append({"role": "assistant", "content": full_display})
    
    # 清除輸入
    st.rerun()

# 說明
with st.expander("使用說明"):
    st.write("""
    - **聊天室介面**：左邊為機器人 (assistant) 回答，右邊為使用者 (user) 問題。
    - **Streaming=True**：顯示執行步驟狀態 (status_placeholder)，並逐字即時顯示回答。
    - **Streaming=False**：無步驟顯示，直接等待並顯示完整回答。
    - **聊天歷史**：自動維護多輪對話，傳送到後端。
    - **Metadata**：側邊欄設定，任意 JSON 傳遞額外資訊。
    - 確保後端在 localhost:8000 運行。
    - 安裝依賴: `pip install streamlit requests`
    - 執行: `streamlit run chat_app.py`
    """)