ui/dashboard.py

import gradio as gr
from .logic import interact_with_env, reset_ui

with gr.Blocks(title="Voice Agent Orchestrator") as demo:
    gr.Markdown("# 🏦 System Orchestrator Dashboard")
    gr.Markdown("Visualizing **Layer 1** (Voice Agent Policy) vs **Layer 2** (Customer Agent) within the Tool Sandbox.")
    
    # --- Hidden State Trackers ---
    chat_history = gr.State([])
    cum_reward = gr.State(0.0)
    step_count = gr.State(0)
    customer_state = gr.State({"Intent": "Waiting...", "Satisfaction": "N/A", "Cooperation": "N/A"})
    db_state = gr.State({"Status": "Waiting for episode start"})

    with gr.Row():
        # ==========================================
        # LEFT COLUMN: The Main Arena (Layer 1 & 2)
        # ==========================================
        with gr.Column(scale=3):
            gr.Markdown("### 🏟️ The Main Arena (Dialogue)")
            conversation_box = gr.Chatbot(label="Agent Policy vs Customer Environment", height=450)
            
            agent_input = gr.Textbox(
                label="Layer 1 Action Space Input", 
                placeholder='Type text or JSON e.g. {"action_type": "tool_call", "content": "lookup_account", "thinking": "User seems angry."}'
            )
            
            with gr.Row():
                send_btn = gr.Button("Send Action", variant="primary")
                clear_btn = gr.Button("Reset Episode / Clear")

            gr.Markdown("### 🧠 Layer 1 Internals (Belief State & Tools)")
            with gr.Row():
                thinking_box = gr.Textbox(label="Agent Belief State", interactive=False, lines=2)
                tool_box = gr.Textbox(label="Tool Calls Triggered", interactive=False, lines=2)
            
            accuracy_box = gr.Textbox(label="💰 Immediate Step Reward", interactive=False)

        # ==========================================
        # RIGHT COLUMN: The Sandbox & Hidden States
        # ==========================================
        with gr.Column(scale=1):
            gr.Markdown("### 🎭 Layer 2 Customer State")
            customer_state_box = gr.Markdown("**Hidden Intent:** Waiting...\n\n**Satisfaction Tracker:** N/A\n\n**Cooperation Level:** N/A")
            
            gr.Markdown("---")
            gr.Markdown("### 🗄️ Tool/API Sandbox DB")
            # This JSON box will update live if your tools modify the database!
            db_state_box = gr.JSON(value={"Status": "Waiting for episode start"}, label="Account & Txn Engine")

    # ==========================================
    # BOTTOM ACCORDION: The Reward Module
    # ==========================================
    with gr.Accordion("📈 Reward Module & Scoring Metrics", open=True):
        scorecard_box = gr.Markdown("### 🏆 Reward Module Output\nWaiting for episode...")

    # --- Wiring ---
    action_inputs = [agent_input, chat_history, cum_reward, step_count, customer_state, db_state]
    action_outputs = [
        conversation_box, chat_history, thinking_box, tool_box, 
        accuracy_box, customer_state_box, db_state_box, scorecard_box, 
        cum_reward, step_count, customer_state, db_state
    ]

    # Submit via Enter key or Button click
    agent_input.submit(interact_with_env, inputs=action_inputs, outputs=action_outputs).then(lambda: "", None, [agent_input])
    send_btn.click(interact_with_env, inputs=action_inputs, outputs=action_outputs).then(lambda: "", None, [agent_input])
    
    # Clear button routing
    clear_btn.click(
        reset_ui, 
        inputs=None, 
        outputs=[conversation_box, chat_history, cum_reward, step_count, thinking_box, tool_box, accuracy_box, customer_state_box, db_state_box, scorecard_box, customer_state, db_state]
    )