"""
可视化交互测试台 - Streamlit

左侧边栏：切换 Blackboard.is_in_air（模拟在地面/空中）
主界面：输入自然语言指令，分步展示四层流水线输出。
"""

from __future__ import annotations

import json
import sys
import time
from pathlib import Path

# 确保 src 在路径中
sys.path.insert(0, str(Path(__file__).resolve().parent / "src"))

import streamlit as st

from drone_planning.core.blackboard import DroneStateBlackboard
from drone_planning.execution.tree_wrapper import parse_json_to_tree, tree_to_ascii, wrap_and_build_tree
from drone_planning.pipeline.composer import _load_node_schema, build_system_prompt
from drone_planning.pipeline.planner import plan
from drone_planning.pipeline.router import route
from drone_planning.rag.retriever import RAGRetriever

st.set_page_config(page_title="无人机行为规划测试台", page_icon="🚁", layout="wide")

# 缩小环节耗时、intents、entities 等数据的字体
st.markdown("""
<style>
/* 环节耗时、intents、entities 等 metric 字体 */
div[data-testid="stMetric"] { font-size: 0.8rem !important; }
div[data-testid="stMetric"] label { font-size: 0.75rem !important; }
div[data-testid="stMetric"] [data-testid="stMetricValue"] { font-size: 0.85rem !important; }
/* JSON、code 块字体 */
div[data-testid="stJson"], .stJson { font-size: 0.8rem !important; }
div[data-testid="stCode"] code, .stCode code { font-size: 0.8rem !important; }
/* expander 内标题与内容 */
.streamlit-expanderHeader { font-size: 0.9rem !important; }
.streamlit-expanderContent { font-size: 0.85rem !important; }
/* caption 字体 */
.stCaptionContainer { font-size: 0.8rem !important; }
/* subheader 略小 */
h3 { font-size: 1rem !important; }
</style>
""", unsafe_allow_html=True)

st.title("🚁 无人机行为规划 - 四层流水线测试台")
st.caption("输入自然语言指令，观察 Layer 1~4 的逐步输出。切换「是否在空中」可验证 Layer 4 的起飞包装逻辑。")

# ---------------------------------------------------------------------------
# 左侧边栏：Blackboard 状态
# ---------------------------------------------------------------------------
with st.sidebar:
    st.header("⚙️ 黑板状态")
    is_in_air = st.toggle("是否在空中 (is_in_air)", value=False, help="关=在地面，开=已起飞。关掉时 Layer 4 会在业务树前插入 SystemCheck + Takeoff")
    bb = DroneStateBlackboard()
    bb.is_in_air = is_in_air
    st.info(f"当前: {'已起飞' if is_in_air else '在地面'}")

# ---------------------------------------------------------------------------
# 主界面：输入与执行
# ---------------------------------------------------------------------------
instruction = st.text_area("输入自然语言指令", height=80, placeholder="例如：飞到大门然后拍照、起飞、飞到广场然后搜索汽车")
run_text = instruction.strip() if instruction else ""
if run_text:
    st.divider()
    if st.button("执行规划", type="primary", key="run_plan"):
        # 再次同步 blackboard（用户可能在点击前改了侧边栏）
        bb = DroneStateBlackboard()
        bb.is_in_air = is_in_air

        layer1_result = None
        rag_context = None
        layer2_prompt = None
        layer2_meta = None
        tool_call_log: list[dict] = []
        layer3_json = None
        layer4_ascii = None
        fast_path = False
        error_msg = None
        timing: dict[str, float] = {}

        try:
            # Layer 1: Router
            t0 = time.perf_counter()
            layer1_result = route(run_text)
            timing["Layer1_Router"] = (time.perf_counter() - t0) * 1000
            fast_path = layer1_result.is_fast_path

            if fast_path:
                st.success("✅ Layer 1 触发 Fast-Path 短路，后续层不执行")
            else:
                # RAG 检索（阶段二）
                t0 = time.perf_counter()
                retriever = RAGRetriever()
                rag_context = retriever.retrieve_context(
                    intents=layer1_result.intents,
                    entities=layer1_result.entities,
                    user_text=run_text,
                )
                timing["RAG_检索"] = (time.perf_counter() - t0) * 1000

                # Layer 2: Composer
                t0 = time.perf_counter()
                schema_json = _load_node_schema()
                layer2_prompt, layer2_meta = build_system_prompt(
                    intents=layer1_result.intents,
                    entities=layer1_result.entities,
                    schema_json=schema_json,
                    rag_context=rag_context,
                )
                timing["Layer2_Composer"] = (time.perf_counter() - t0) * 1000

                # Layer 3: Planner（支持 Tool 循环）
                t0 = time.perf_counter()
                layer3_json = plan(
                    system_prompt=layer2_prompt,
                    user_text=run_text,
                    tool_call_log=tool_call_log,
                )
                timing["Layer3_Planner"] = (time.perf_counter() - t0) * 1000

                # Layer 4: 解析 + 包装（会读取 blackboard.is_in_air）
                t0 = time.perf_counter()
                root_node = layer3_json.get("root") or layer3_json
                business_tree = parse_json_to_tree(root_node)
                final_tree = wrap_and_build_tree(business_tree)
                layer4_ascii = tree_to_ascii(final_tree)
                timing["Layer4_Execution"] = (time.perf_counter() - t0) * 1000

        except Exception as e:
            error_msg = str(e)
            import traceback
            st.error(f"❌ 执行出错: {error_msg}")
            st.code(traceback.format_exc(), language="text")

        # -----------------------------------------------------------------------
        # 各环节耗时
        # -----------------------------------------------------------------------
        if timing:
            st.subheader("⏱️ 各环节耗时 (ms)")
            cols = st.columns(len(timing))
            for i, (name, ms) in enumerate(timing.items()):
                with cols[i]:
                    st.metric(name, f"{ms:.0f} ms")
            st.caption(f"总计: {sum(timing.values()):.0f} ms")

        # -----------------------------------------------------------------------
        # 分步展示
        # -----------------------------------------------------------------------
        st.subheader("📋 各层输出")

        # Layer 1
        with st.expander("**Layer 1 (Router)** - 意图与实体", expanded=True):
            if layer1_result:
                col1, col2, col3 = st.columns(3)
                with col1:
                    st.metric("intents", ", ".join(layer1_result.intents) or "-")
                with col2:
                    ent_str = json.dumps(layer1_result.entities, ensure_ascii=False)
                    st.metric("entities", ent_str[:60] + "..." if len(ent_str) > 60 else ent_str or "-")
                with col3:
                    st.metric("Fast-Path", "是" if fast_path else "否")
                st.json(layer1_result.entities)
            else:
                st.warning("Layer 1 未执行或出错")

        if not fast_path:
            # RAG 检索结果（Layer 1 与 Layer 2 之间）
            with st.expander("**RAG 检索结果** - 地图坐标、规则、Few-shot 示例", expanded=True):
                if rag_context:
                    col1, col2 = st.columns(2)
                    with col1:
                        st.markdown("**地图坐标 (map_context)**")
                        if rag_context.get("map_context"):
                            st.code(rag_context["map_context"], language=None)
                        else:
                            st.caption("无匹配")
                    with col2:
                        st.markdown("**规则约束 (rule_context)**")
                        if rag_context.get("rule_context"):
                            st.code(rag_context["rule_context"], language=None)
                        else:
                            st.caption("无匹配")
                    st.markdown("**Few-shot 示例**")
                    if rag_context.get("few_shot_examples"):
                        for i, ex in enumerate(rag_context["few_shot_examples"]):
                            st.json({"instruction": ex.get("instruction"), "tree_json": ex.get("tree_json")})
                    else:
                        st.caption("无相似示例")
                else:
                    st.caption("RAG 未执行")

            # Layer 2
            with st.expander("**Layer 2 (Composer)** - 精简 Prompt 与选中节点", expanded=True):
                if layer2_meta:
                    st.metric("Prompt 长度", f"{layer2_meta.get('prompt_length', 0)} 字符")
                    st.metric("选中节点", ", ".join(layer2_meta.get("selected_actions", [])))
                    st.metric("基准点坐标", json.dumps(layer2_meta.get("base_location_coords", {}), ensure_ascii=False))
                if layer2_prompt:
                    preview = layer2_prompt[:2000] + ("..." if len(layer2_prompt) > 2000 else "")
                    st.code(preview, language="text")

            # 坐标计算与 Tool Call 日志（仅 LLM 通过 mcp 工具计算）
            rel_descs = (rag_context or {}).get("relative_descriptions") or []
            has_any = bool(rel_descs) or bool(tool_call_log)
            with st.expander("**坐标计算与 Tool Call 日志**", expanded=has_any):
                if rel_descs:
                    st.markdown("**RAG 提供的相对描述**（需 LLM 调用 calculate_relative_coordinate）")
                    for rd in rel_descs:
                        st.json(rd)
                if tool_call_log:
                    st.markdown("**LLM Tool Call**（大模型调用 mcp 计算）")
                    for entry in tool_call_log:
                        st.markdown(f"第 {entry.get('round', '?')} 次 `{entry.get('tool', '')}`")
                        st.json({"arguments": entry.get("arguments"), "result": entry.get("result")})
                if not rel_descs and not tool_call_log:
                    st.caption("无相对描述（指令仅含绝对地点时，直接使用 base_location_coords）")

            # Layer 3
            with st.expander("**Layer 3 (Planner)** - 原生业务树 JSON", expanded=True):
                if layer3_json:
                    st.json(layer3_json)
                else:
                    st.warning("Layer 3 未执行")

            # Layer 4
            with st.expander("**Layer 4 (Execution)** - 安全包装后的 py_trees ASCII 结构", expanded=True):
                if layer4_ascii:
                    st.text("当 is_in_air=False 时，此处应看到 [Safe_Execution] -> [SystemCheck] -> [Takeoff] -> 业务树")
                    st.code(layer4_ascii, language=None)
                else:
                    st.warning("Layer 4 未执行")