"""中医知识图谱问答 —— 流式页面（调后端 /process_stream）

和首页（app.py，调 /process）问的是同一张图，区别只在**什么时候给反馈**：
首页要等 5-30 秒才一次性出答案，中间只有一个转圈；这一页把图里每个节点的产出
边跑边显示出来，用户能看到"正在抽取实体""正在查图谱"，等待就不再是黑盒。

Streamlit 1.57 的多页机制：`pages/` 目录下每个 .py 自动成为一页，侧边栏出现导航，
入口仍然是 app.py。所以这一页只管自己的界面，不用改 app.py。

运行（和首页同一条命令，起来后从侧边栏切页）:

    python -X utf8 -m _005_fastapi.main       # 1. 后端
    streamlit run _006_streamlit/app.py       # 2. 前端

网络那一层在 api_client.ask_backend_stream，这里只管界面。
"""

import sys
import uuid
from pathlib import Path

# 和 app.py 同一个理由：streamlit run 只把**脚本所在目录**放进 sys.path。
# 这一页在 pages/ 下，比 app.py 还深一层，所以要往上走两级才是仓库根。
# 这几行必须排在项目内的 import 之前，下面的 E402 是故意的。
_REPO_ROOT = str(Path(__file__).resolve().parent.parent.parent)
if _REPO_ROOT not in sys.path:
    sys.path.insert(0, _REPO_ROOT)

import streamlit as st

from _006_streamlit.api_client import (
    DEFAULT_BASE_URL,
    DEFAULT_TIMEOUT,
    BackendError,
    ask_backend_stream,
    check_health,
)

# 示例问题挑的是**要查图谱**的那种：闲聊只走两个节点，看不出流式的效果
SAMPLE_QUESTIONS = [
    "人参有什么功效？",
    "四君子汤由哪些药材组成？",
    "我最近总是拉肚子，有什么方剂可以调理？",
]

# AgentState 里的键名对用户没有意义，翻译成人看得懂的中文。
# 键取自 _004_langgraph_more_nodes/agent_state.py —— 那边加字段，这里补一行，
# 没补也不会坏，下面的 fallback 会直接显示原始键名。
FIELD_LABELS = {
    "input_semantic_trans": "改写后的问题",
    "is_zhongyi_intent": "是否中医相关",
    "user_input_herbs": "问题里的中药",
    "user_input_formulas": "问题里的方剂",
    "user_input_symptoms": "问题里的症状",
    "user_input_diseases": "问题里的疾病",
    "user_input_effects": "问题里的功效",
    "user_input_sources": "问题里的出处",
    "matched_herbs": "图谱中的中药",
    "matched_formulas": "图谱中的方剂",
    "matched_symptoms": "图谱中的症状",
    "matched_diseases": "图谱中的疾病",
    "matched_effects": "图谱中的功效",
    "matched_sources": "图谱中的出处",
    "cypher_query": "生成的查询语句",
    "is_all_validate_cypher": "查询语句校验",
    "cypher_results": "图谱返回的数据",
}

st.set_page_config(page_title="流式问答 · 中医知识图谱", page_icon="🌊", layout="centered")

def init_state() -> None:
    """初始化 session_state。rerun 会重跑整个脚本，所以每个键都要判存在再写。"""
    if "stream_messages" not in st.session_state:
        # [{"role", "content", "thinking": [(节点名, 增量), ...]}]。
        # thinking 也要存：rerun 会重跑整个脚本，不存的话一提下一个问题，
        # 上一轮的思考过程就消失了。用和首页不同的键，两页历史互不打扰
        st.session_state.stream_messages = []
    if "user_id" not in st.session_state:
        # 和首页共用这个键：同一个浏览器会话就是同一个用户，
        # 后端拿它当 thread_id，切页不该换身份
        st.session_state.user_id = f"streamlit-{uuid.uuid4().hex[:12]}"
    if "stream_pending" not in st.session_state:
        # 点示例按钮时把问题存这里，交给下一次 rerun 处理。
        # 按钮回调发生在页面渲染中途，直接在那儿发请求会让聊天记录的渲染顺序错乱
        st.session_state.stream_pending = None


def render_sidebar() -> tuple[str, int]:
    """侧边栏：后端地址、超时、状态灯、清空按钮。返回 (base_url, timeout)"""
    with st.sidebar:
        st.header("⚙️ 设置")
        base_url = st.text_input(
            "后端地址",
            value=DEFAULT_BASE_URL,
            help="_005_fastapi 服务的地址。端口被占用时改这里，不用改代码",
            key="stream_base_url",
        )
        timeout = st.number_input(
            "超时（秒）",
            min_value=10,
            max_value=600,
            value=DEFAULT_TIMEOUT,
            step=10,
            help="流式接口的超时是**读一行**的间隔，不是总时长，给到 120 秒很宽裕",
            key="stream_timeout",
        )

        # 手动触发，不自动探活：每次 rerun 都发一个请求属于白费
        if st.button("检查后端连接", use_container_width=True, key="stream_health"):
            if check_health(base_url):
                st.success("后端正常")
            else:
                st.error("连不上后端，请先启动 _005_fastapi")

        st.divider()
        if st.button("🗑️ 清空对话", use_container_width=True, key="stream_clear"):
            st.session_state.stream_messages = []
            st.rerun()

        st.caption(f"会话 user_id：`{st.session_state.user_id}`")

    return base_url, int(timeout)


def format_value(value) -> str:
    """把状态值渲染成一行 markdown。

    值的形状不统一：布尔、字符串、字符串列表（实体、cypher）、嵌套的查询结果。
    列表用顿号连起来比原始 JSON 好读；cypher 和查询结果太长，截断，
    这里是"思考过程"不是调试面板，看个大概就够。
    """
    if isinstance(value, bool):
        return "是" if value else "否"
    if isinstance(value, list):
        # 全是字符串的列表（实体名、cypher 语句）直接拼
        if all(isinstance(v, str) for v in value):
            text = "、".join(value)
        else:
            # cypher_results 这种嵌套结构，只报个条数，展开了没人看
            return f"共 {len(value)} 组结果"
        return text if len(text) <= 300 else text[:300] + "…"
    text = str(value)
    return text if len(text) <= 300 else text[:300] + "…"


def render_thinking(container, node: str, data: dict) -> None:
    """把一个节点的状态增量画进容器。

    每个节点一个小标题 + 若干"字段：值"。用 container.markdown 而不是 st.markdown，
    是为了把内容写进 status 里那块固定区域，而不是追加到页面末尾。
    """
    lines = [f"**{node}**"]
    for key, value in data.items():
        # 没登记的键显示原始键名，加字段时不至于漏掉不显示
        label = FIELD_LABELS.get(key, key)
        lines.append(f"- {label}：{format_value(value)}")
    container.markdown("\n".join(lines))


def handle_question(question: str, base_url: str, timeout: int) -> None:
    """跑一轮流式问答，边收边画。

    先把用户消息写进历史再发请求，这样请求失败时提问也留在页面上，
    用户不用重新打一遍。
    """
    st.session_state.stream_messages.append({"role": "user", "content": question})
    with st.chat_message("user"):
        st.markdown(question)

    with st.chat_message("assistant"):
        # st.status 给的是一块带标题和状态图标的固定区域，思考过程往里写。
        # expanded=True 且跑完也不收，这一页存在的意义就是让人看见中间过程
        status = st.status("正在思考…", expanded=True)
        answer = ""
        is_error = False
        # 边收边记，跑完存进历史，rerun 时重放
        thinking_log: list = []

        try:
            for message in ask_backend_stream(
                question,
                user_id=st.session_state.user_id,
                base_url=base_url,
                timeout=timeout,
            ):
                kind = message.get("type")
                if kind == "thinking":
                    node = message.get("node") or "未知节点"
                    data = message.get("data") or {}
                    # 节点名直接当进度文案，比"处理中"有信息量
                    status.update(label=f"正在执行：{node}")
                    render_thinking(status, node, data)
                    thinking_log.append((node, data))
                elif kind == "final":
                    answer = (message.get("output") or "").strip()
                elif kind == "error":
                    # 后端图跑挂了，HTTP 仍是 200，错误文案在 output 里
                    answer = (message.get("output") or "").strip()
                    is_error = True
        except BackendError as exc:
            # 连不上/超时/参数不合法：流根本没开始
            answer = f"⚠️ {exc}"
            is_error = True

        if not answer:
            # 流正常结束但没收到 final，属于后端不该有的行为，兜一下
            answer = "⚠️ 后端没有返回最终答案，请重试。"
            is_error = True

        # 跑的时候展开（用户要看进度），跑完收起（答案才是结论，中间过程让位）。
        # 点一下标题还能再展开，信息没丢
        status.update(
            label="思考失败" if is_error else "思考完成",
            state="error" if is_error else "complete",
            expanded=False,
        )
        # 回答里带 markdown（后端的 LLM 会输出加粗、列表），用 markdown 渲染
        st.markdown(answer)

    st.session_state.stream_messages.append(
        {
            "role": "assistant",
            "content": answer,
            "thinking": thinking_log,
            "is_error": is_error,
        }
    )


def main() -> None:
    init_state()
    base_url, timeout = render_sidebar()

    st.title("🌊 流式问答")
    st.caption("和首页问的是同一张图，但会把每个节点的思考过程实时显示出来")

    # 空历史时给几个示例，用户不必自己想第一个问题该怎么问
    if not st.session_state.stream_messages:
        st.markdown("**试试这些问题：**")
        for i, sample in enumerate(SAMPLE_QUESTIONS):
            # key 必须唯一，且不能和首页的 sample-N 撞
            if st.button(sample, key=f"stream-sample-{i}", use_container_width=True):
                # 存起来交给下一次 rerun 处理，别在渲染中途发请求
                st.session_state.stream_pending = sample
                st.rerun()

    # 重放历史。rerun 会重跑整个脚本，页面上的消息全靠这一段重新画出来，
    # 往轮的思考过程装在收起的 expander 里 —— 想回看点开就有，不占版面
    for item in st.session_state.stream_messages:
        with st.chat_message(item["role"]):
            thinking = item.get("thinking")
            if thinking:
                with st.expander("💭 思考过程", expanded=False):
                    for node, data in thinking:
                        render_thinking(st, node, data)
            st.markdown(item["content"])

    # 输入框固定在页面底部，回车提交
    question = st.chat_input("请输入你的问题，例如：人参有什么功效？")

    # 示例按钮排在输入框之后处理，两条来源汇成一个入口
    if st.session_state.stream_pending:
        question = st.session_state.stream_pending
        st.session_state.stream_pending = None

    if question and question.strip():
        handle_question(question.strip(), base_url, timeout)


main()