from __future__ import annotations

import asyncio
from datetime import datetime

import streamlit as st

from nexus.sandbox.visualization import (
    build_graph_progress_fig,
    build_graph_fig,
    build_heatmap_fig,
    build_ledger_area_fig,
    build_radar_fig,
    build_sunburst_fig,
    build_timeline_fig,
    build_timeseries_fig,
)
from nexus.sandbox.workflow import NexusSandboxWorkflow


def _init_state() -> None:
    if "messages" not in st.session_state:
        st.session_state.messages = []
    if "events" not in st.session_state:
        st.session_state.events = []
    if "running" not in st.session_state:
        st.session_state.running = False
    if "agent_traces" not in st.session_state:
        st.session_state.agent_traces = []


def _theme() -> None:
    st.markdown(
        """
        <style>
        .main > div {padding-top: 1rem;}
        .block-container {padding-top: 1.1rem; padding-bottom: 1rem;}
        .stChatMessage {border-radius: 12px;}
        </style>
        """,
        unsafe_allow_html=True,
    )


def run_app() -> None:
    st.set_page_config(page_title="NEXUS Sandbox", page_icon="N", layout="wide")
    _theme()
    _init_state()

    st.title("NEXUS Sandbox")
    st.caption("Two-pane live sandbox using NEXUS workflow with local llama-server (Qwen) only.")

    with st.sidebar:
        st.subheader("Local LLM Settings")
        base_url = st.text_input("llama-server URL", value="http://localhost:8080")
        model = st.text_input("Model", value="qwen2.5-3b-instruct-q4_k_m.gguf")
        st.markdown("This sandbox uses only local provider mode.")

    chat_col, viz_col = st.columns([1.1, 1.3], gap="large")

    with chat_col:
        st.subheader("Agent Chat")
        for msg in st.session_state.messages:
            with st.chat_message(msg["role"]):
                st.markdown(msg["content"])
        if st.session_state.agent_traces:
            with st.expander("Latest Agent Outputs", expanded=True):
                st.markdown("\n\n".join(st.session_state.agent_traces[-6:]))

    with viz_col:
        st.subheader("Live Visualization")
        main_plot_placeholder = st.empty()
        graph_placeholder = st.empty()
        metric_cols = st.columns(3)
        radar_placeholder = st.empty()
        heatmap_placeholder = st.empty()
        timeline_placeholder = st.empty()
        sunburst_placeholder = st.empty()
        ledger_placeholder = st.empty()
        graph_progress_placeholder = st.empty()
        pillar_meta_placeholder = st.empty()

        if st.session_state.events:
            main_plot_placeholder.plotly_chart(
                build_timeseries_fig(st.session_state.events),
                use_container_width=True,
                key="timeseries_initial",
            )
            graph_placeholder.plotly_chart(
                build_graph_fig(st.session_state.events),
                use_container_width=True,
                key="graph_initial",
            )
            latest = st.session_state.events[-1]
            metric_cols[0].metric("Cumulative Tokens", latest.get("cumulative_tokens", 0))
            metric_cols[1].metric("Tool Calls", latest.get("tool_calls", 0))
            metric_cols[2].metric("Hot Tokens", latest.get("hot_tokens", 0))
            radar_placeholder.plotly_chart(build_radar_fig(latest), use_container_width=True, key="radar_initial")
            heatmap_placeholder.plotly_chart(
                build_heatmap_fig(st.session_state.events),
                use_container_width=True,
                key="heatmap_initial",
            )
            timeline_placeholder.plotly_chart(
                build_timeline_fig(st.session_state.events),
                use_container_width=True,
                key="timeline_initial",
            )
            sunburst_placeholder.plotly_chart(
                build_sunburst_fig(latest),
                use_container_width=True,
                key="sunburst_initial",
            )
            ledger_placeholder.plotly_chart(
                build_ledger_area_fig(st.session_state.events),
                use_container_width=True,
                key="ledger_initial",
            )
            graph_progress_placeholder.plotly_chart(
                build_graph_progress_fig(st.session_state.events),
                use_container_width=True,
                key="graph_progress_initial",
            )
            with pillar_meta_placeholder.container():
                st.markdown("### NEXUS Pillar Telemetry")
                st.json(latest.get("pillars", {}))

    prompt = st.chat_input("Ask NEXUS agents...")
    if prompt and not st.session_state.running:
        st.session_state.running = True
        st.session_state.messages.append({"role": "user", "content": prompt})
        with chat_col:
            with st.chat_message("user"):
                st.markdown(prompt)

        workflow = NexusSandboxWorkflow(base_url=base_url, model=model)

        with chat_col:
            agent_msg = st.chat_message("assistant")
            answer_placeholder = agent_msg.empty()
            status_placeholder = st.empty()
            trace_placeholder = st.empty()

        async def consume() -> None:
            async for event in workflow.run(prompt):
                row = {
                    "step": len(st.session_state.events) + 1,
                    "stage": event.stage,
                    "node_id": event.node_id,
                    "agent_id": event.agent_id,
                    "status": event.status,
                    "timestamp": event.timestamp.isoformat(),
                    "message": event.message,
                    "cumulative_tokens": event.cumulative_tokens,
                    "hot_tokens": event.hot_tokens,
                    "warm_entries": event.warm_entries,
                    "cold_count": event.cold_count,
                    "tool_calls": event.tool_calls,
                    "checkpoint_count": event.checkpoint_count,
                    "graph_nodes": event.graph_nodes,
                    "completed_nodes": event.completed_nodes,
                    "ready_nodes": event.ready_nodes,
                    "mutation_count": event.mutation_count,
                    "circuit_state": event.circuit_state,
                    "planner_used": event.planner_used,
                    "researcher_used": event.researcher_used,
                    "synthesizer_used": event.synthesizer_used,
                    "duration_ms": event.duration_ms,
                    "agent_output": event.agent_output,
                    "pillars": event.meta.get("pillars", {}),
                }
                st.session_state.events.append(row)
                status_placeholder.info(f"{event.stage}: {event.message}")
                if event.agent_output:
                    preview = event.agent_output if len(event.agent_output) <= 2000 else (event.agent_output[:2000] + "\n... [truncated]")
                    st.session_state.agent_traces.append(
                        f"### {event.agent_id} [{event.stage}]\n{preview}"
                    )
                    trace_placeholder.markdown("\n\n".join(st.session_state.agent_traces[-6:]))
                if event.partial_answer:
                    answer_placeholder.markdown(event.partial_answer)

                with viz_col:
                    step_key = row["step"]
                    main_plot_placeholder.plotly_chart(
                        build_timeseries_fig(st.session_state.events),
                        use_container_width=True,
                        key=f"timeseries_live_{step_key}",
                    )
                    graph_placeholder.plotly_chart(
                        build_graph_fig(st.session_state.events),
                        use_container_width=True,
                        key=f"graph_live_{step_key}",
                    )
                    latest = st.session_state.events[-1]
                    metric_cols[0].metric("Cumulative Tokens", latest.get("cumulative_tokens", 0))
                    metric_cols[1].metric("Tool Calls", latest.get("tool_calls", 0))
                    metric_cols[2].metric("Hot Tokens", latest.get("hot_tokens", 0))
                    radar_placeholder.plotly_chart(
                        build_radar_fig(latest),
                        use_container_width=True,
                        key=f"radar_live_{step_key}",
                    )
                    heatmap_placeholder.plotly_chart(
                        build_heatmap_fig(st.session_state.events),
                        use_container_width=True,
                        key=f"heatmap_live_{step_key}",
                    )
                    timeline_placeholder.plotly_chart(
                        build_timeline_fig(st.session_state.events),
                        use_container_width=True,
                        key=f"timeline_live_{step_key}",
                    )
                    sunburst_placeholder.plotly_chart(
                        build_sunburst_fig(latest),
                        use_container_width=True,
                        key=f"sunburst_live_{step_key}",
                    )
                    ledger_placeholder.plotly_chart(
                        build_ledger_area_fig(st.session_state.events),
                        use_container_width=True,
                        key=f"ledger_live_{step_key}",
                    )
                    graph_progress_placeholder.plotly_chart(
                        build_graph_progress_fig(st.session_state.events),
                        use_container_width=True,
                        key=f"graph_progress_live_{step_key}",
                    )
                    with pillar_meta_placeholder.container():
                        st.markdown("### NEXUS Pillar Telemetry")
                        st.json(latest.get("pillars", {}))

            final_answer = st.session_state.events[-1]["message"] if st.session_state.events else "Completed."
            if st.session_state.events:
                final_answer = workflow.mesh.get("synthesizer", "final_answer") or final_answer
            st.session_state.messages.append({"role": "assistant", "content": str(final_answer)})

        asyncio.run(consume())
        st.session_state.running = False
        st.rerun()


if __name__ == "__main__":
    run_app()
