# app.py
import streamlit as st
from streamlit_chat import message
import asyncio
import json
import time
import uuid

# 匯入你的 BusinessAgent
from business_agent import BusinessAgent

# ==================== Streamlit 設定 ====================
st.set_page_config(page_title="BusinessAgent 測試台", layout="wide")
st.title("🧠 BusinessAgent 即時測試介面（Production 版）")

# ==================== 初始化 Agent（只做一次） ====================
if "agent" not in st.session_state:
    with st.spinner("正在初始化 BusinessAgent..."):
        st.session_state.agent = BusinessAgent(
            azure_openai_api_key=st.secrets["AZURE_OPENAI_KEY"],
            azure_openai_endpoint=st.secrets["AZURE_OPENAI_ENDPOINT"],
            azure_deployment=st.secrets["AZURE_DEPLOYMENT"],
            vector_search_client=None,  # 測試階段可用 Mock
            timeout_seconds=45.0,
        )

agent: BusinessAgent = st.session_state.agent

# ==================== Session State（極簡！） ====================
if "messages" not in st.session_state:
    st.session_state.messages = []        # 唯一來源：顯示 + 傳給 Agent
if "processing" not in st.session_state:
    st.session_state.processing = False
if "start_time" not in st.session_state:
    st.session_state.start_time = None

# ==================== 左側邊欄 ====================
with st.sidebar:
    st.header("📋 Metadata 設定")
    default_meta = {
        "request_id": f"req_{uuid.uuid4().hex[:10]}",
        "credit_card_name": "白金卡",
        "user_tier": "VIP",
        "channel": "web_test"
    }
    metadata_str = st.text_area(
        "Metadata JSON",
        value=json.dumps(default_meta, indent=2, ensure_ascii=False),
        height=250
    )

    st.header("🔍 Observation 即時監控")
    obs_box = st.expander("Observation", expanded=True)
    obs_placeholder = obs_box.empty()

# ==================== 聊天視窗 ====================
chat_container = st.container()

with chat_container:
    # 顯示歷史訊息
    for i, msg in enumerate(st.session_state.messages):
        if msg["role"] == "user":
            message(msg["content"], is_user=True, key=f"user_{i}")
        else:
            message(msg["content"], is_user=False, key=f"ai_{i}")

    # 處理中動態訊息
    if st.session_state.processing:
        elapsed = int(time.time() - st.session_state.start_time)
        st.info(f"🤔 您的問題正在處理中... ({elapsed}s)")

# ==================== 輸入框 ====================
user_input = st.chat_input("輸入您的問題...")

if user_input:
    # 加入使用者訊息
    st.session_state.messages.append({"role": "user", "content": user_input})
    with chat_container:
        message(user_input, is_user=True)
    
    st.session_state.processing = True
    st.session_state.start_time = time.time()
    st.rerun()

# ==================== 背景執行 Agent ====================
async def run_agent():
    try:
        metadata = json.loads(metadata_str)
    except:
        metadata = {}

    # 累積當前回覆（用於 token 串流）
    current_reply = ""

    async for event in agent.astream(
        history=st.session_state.messages.copy(),  # 唯一來源！
        request_context=metadata
    ):
        etype = event.get("type")

        # 1. 即時 token
        if etype == "token":
            current_reply += event["content"]
            # 更新最後一筆 AI 訊息（即時顯示）
            if st.session_state.messages and st.session_state.messages[-1]["role"] == "assistant":
                st.session_state.messages[-1]["content"] = current_reply
            else:
                st.session_state.messages.append({"role": "assistant", "content": current_reply})
            st.rerun()

        # 2. 流程結束（包含 quick_reply 直接結束）
        elif etype == "final":
            final_answer = event.get("reply") or event.get("answer", "")
            if final_answer:
                if st.session_state.messages and st.session_state.messages[-1]["role"] == "assistant":
                    st.session_state.messages[-1]["content"] = final_answer
                else:
                    st.session_state.messages.append({"role": "assistant", "content": final_answer})

            # 加入歷史（供下一輪使用）
            if final_answer:
                # 避免重複加入
                if not (st.session_state.messages and st.session_state.messages[-1]["role"] == "assistant"):
                    st.session_state.messages.append({"role": "assistant", "content": final_answer})

            # 顯示 metadata（後端存檔用）
            if "metadata" in event:
                obs_placeholder.success("流程完成！以下為可存檔 metadata：")
                obs_placeholder.json(event["metadata"], expanded=False)

            # 顯示完整 log（可選）
            if "logs" in event:
                with obs_box:
                    st.caption("完整執行 Log")
                    st.json(event["logs"], expanded=False)

        # 3. Observation（如果你有吐 obs）
        elif etype == "obs" and "data" in event:
            obs_placeholder.json(event["data"], expanded=False)

    # 結束處理狀態
    st.session_state.processing = False
    st.rerun()

# 啟動背景任務
if st.session_state.processing and "agent_task" not in st.session_state:
    st.session_state.agent_task = asyncio.create_task(run_agent())

# 自動刷新（token 即時顯示）
if st.session_state.processing:
    st.rerun()