"""WalkXR Role-Play Agent — Streamlit Demo."""

from dotenv import load_dotenv

load_dotenv()  # Load .env first

import streamlit as st
from langchain_core.messages import HumanMessage

from scenarios.bus_stop import BUS_STOP_SCENARIO
from src.graph import app

st.set_page_config(page_title="WalkXR Role-Play", page_icon="🚶")
st.title("WalkXR Role-Play Agent")
st.caption("Stateful Role-Play with Reflective Memory")

# ── Hardcoded Intro / Outro ─────────────────────────────
INTRO_TEXT = (
    "**[Scene]** It's 6 PM at a quiet bus stop. The evening air is cool. "
    "A man in a worn military jacket stands slightly apart from the other commuters, "
    "staring at the ground. His name is James — though you don't know that yet.\n\n"
    "He looks like he's carrying something heavy, and it's not his backpack.\n\n"
    "---\n"
    "*You have up to 10 turns to talk to him. "
    "There's no right or wrong thing to say — just be present.*"
)

OUTRO_TEXT = (
    "---\n\n"
    "**[End of Session]**\n\n"
    "The bus arrives. James gives a small nod — barely noticeable, "
    "but different from when you first approached. "
    "He steps onto the bus without a word.\n\n"
    "Maybe that was enough. Maybe it wasn't. "
    "But you showed up, and that matters.\n\n"
    "---"
)

# ── Session Initialization ───────────────────────────────
if "walk_state" not in st.session_state:
    st.session_state.walk_state = {
        "messages": [],
        "scenario": BUS_STOP_SCENARIO,
        "turn_index": 0,
        "memory": {
            "confidence_level": 0.5,
            "rapport_score": 0.0,
            "emotion_trajectory": [],
            "dominant_emotion": "neutral",
            "current_conversation_stage": "opening",
            "stage_history": ["opening"],
            "key_observations": [],
            "empathy_hits": 0,
            "empathy_misses": 0,
            "empathy_ratio": 0.5,
            "avg_response_length": 0.0,
            "silence_count": 0,
            "npc_openness": 0.3,
            "npc_mood": "guarded",
        },
        "agent_output": {},
        "reflection_result": "",
        "error": None,
    }
    st.session_state.chat_history = []  # For UI display
    st.session_state.ended = False
    st.session_state.intro_shown = False

scenario = BUS_STOP_SCENARIO

# ── Sidebar: Scenario Info ───────────────────────────────
with st.sidebar:
    st.header("Scenario")
    st.write(f"**{scenario['title']}**")
    st.write(scenario["description"])
    st.divider()
    st.write(f"**NPC:** {scenario['npc_name']}")
    st.write(f"**Setting:** {scenario['setting']}")
    st.write(f"**Goal:** {scenario['goal']}")
    st.divider()

    memory = st.session_state.walk_state.get("memory", {})
    col1, col2 = st.columns(2)
    col1.metric("Trust", f"{memory.get('confidence_level', 0.5):.0%}")
    col2.metric("Rapport", f"{memory.get('rapport_score', 0.0):.0%}")

    col3, col4 = st.columns(2)
    col3.metric("Turn", f"{st.session_state.walk_state.get('turn_index', 0)}/{scenario['max_turns']}")
    col4.metric("Empathy", f"{memory.get('empathy_ratio', 0.5):.0%}")

    st.write(f"**Stage:** {memory.get('current_conversation_stage', 'opening')}")
    st.write(f"**James's mood:** {memory.get('npc_mood', 'guarded')}")
    st.write(f"**Openness:** {memory.get('npc_openness', 0.3):.0%}")

    st.divider()
    stage_history = memory.get("stage_history", [])
    if stage_history:
        st.caption("Stage progression: " + " → ".join(stage_history))

    if st.button("Start Over"):
        for key in list(st.session_state.keys()):
            del st.session_state[key]
        st.rerun()

# ── Hardcoded Intro ──────────────────────────────────────
if not st.session_state.intro_shown:
    st.markdown(INTRO_TEXT)
    st.session_state.intro_shown = True
else:
    st.markdown(INTRO_TEXT)

# ── Render Chat History ──────────────────────────────────
for msg in st.session_state.chat_history:
    with st.chat_message(msg["role"]):
        st.write(msg["content"])

# ── Display Outro + Reflection ───────────────────────────
if st.session_state.ended:
    # Hardcoded outro
    st.markdown(OUTRO_TEXT)

    # LLM-generated reflection
    reflection_text = st.session_state.walk_state.get("reflection_result", "")
    if reflection_text:
        st.subheader("Coach's Reflection")
        st.info(reflection_text)

    # Emotion trajectory chart
    trajectory = memory.get("emotion_trajectory", [])
    if trajectory:
        import pandas as pd

        df = pd.DataFrame(trajectory)
        if not df.empty and "valence" in df.columns:
            st.subheader("Emotion Trajectory")
            st.line_chart(df.set_index("turn")[["valence", "arousal"]])

    st.stop()

# ── User Input ───────────────────────────────────────────
user_input = st.chat_input("Say something to James...")

if user_input:
    # Display user message
    st.session_state.chat_history.append({"role": "user", "content": user_input})
    with st.chat_message("user"):
        st.write(user_input)

    # Add user message to state
    walk_state = st.session_state.walk_state
    walk_state["messages"] = list(walk_state.get("messages", [])) + [
        HumanMessage(content=user_input)
    ]
    walk_state["user_input"] = user_input

    # Run graph
    with st.spinner("James is thinking..."):
        result = app.invoke(walk_state)

    # Update state
    st.session_state.walk_state = result

    # Error check
    if result.get("error"):
        st.error(result["error"])
    else:
        # Display NPC response
        agent_output = result.get("agent_output", {})
        npc_response = agent_output.get("response", "")

        st.session_state.chat_history.append(
            {"role": "assistant", "content": npc_response}
        )
        with st.chat_message("assistant"):
            st.write(npc_response)

        # Check if ended
        if result.get("reflection_result"):
            st.session_state.ended = True
            st.rerun()