app.py

import os
import sys

# Add src directory to Python path for Hugging Face Spaces compatibility
sys.path.insert(0, os.path.join(os.path.dirname(__file__), "src"))

import ast
import json
import logging
from pathlib import Path

import gradio as gr
import spaces

from unpredictable_lord.chat.chat import chat_with_mcp_tools
from unpredictable_lord.chat.mcp_client import MCPClient
from unpredictable_lord.mcp_server.game_state import PERSONALITY_DESCRIPTIONS
from unpredictable_lord.mcp_server.mcp_server import (
    execute_turn,
    get_game_state,
    init_game,
    list_available_advice,
)
from unpredictable_lord.settings import LOGGING_LEVEL

# Configure logging level from environment variable
log_level = getattr(logging, LOGGING_LEVEL, logging.INFO)
logging.basicConfig(level=log_level, format="%(asctime)s %(levelname)s %(message)s")

logger = logging.getLogger(__name__)

logger.info(f"ZeroGPU: {spaces.config.Config.zero_gpu}")

# Load MCP guide from external file
MCP_GUIDE_PATH = Path(__file__).parent / "docs" / "mcp_guide.md"
MCP_GUIDE_CONTENT = MCP_GUIDE_PATH.read_text(encoding="utf-8")

# Constants
NO_GAME_MESSAGE = """_No active game._

**To start playing:**
1. Select a Lord Personality above
2. Click "⚔️ Start New Game"
"""


# Gradio UI
with gr.Blocks(title="Unpredictable Lord") as demo:
    gr.Markdown("# Unpredictable Lord\nLord Advisor AI Simulation")

    # State to store current session_id and system_instructions for Chat tab
    chat_session_id = gr.State(value=None)
    chat_system_instructions = gr.State(value="")

    with gr.Tabs():
        # Chat Tab
        with gr.TabItem("Chat"):
            with gr.Row():
                # Left column: Chat interface
                with gr.Column(scale=3):
                    chatbot = gr.Chatbot(label="Lord AI", height=500, type="messages")

                    with gr.Row():
                        msg = gr.Textbox(
                            label="Your Advice",
                            placeholder="Start a new game first...",
                            scale=4,
                            interactive=False,
                        )
                        submit_btn = gr.Button("Submit", scale=1, interactive=False)

                # Right column: Game status panel
                with gr.Column(scale=1):
                    gr.Markdown("### 🎮 Game Status")

                    # Game control buttons
                    with gr.Group():
                        chat_personality = gr.Dropdown(
                            choices=["cautious", "idealist", "populist"],
                            value="cautious",
                            label="Lord Personality",
                        )
                        start_game_btn = gr.Button(
                            "⚔️ Start New Game", variant="primary"
                        )
                        reset_game_btn = gr.Button("🔄 Reset Game", variant="secondary")

                    # Status display
                    game_status_md = gr.Markdown(value=NO_GAME_MESSAGE)

                    # Parameter bars
                    with gr.Group():
                        territory_bar = gr.Slider(
                            label="🏰 Territory",
                            minimum=0,
                            maximum=100,
                            value=50,
                            interactive=False,
                        )
                        population_bar = gr.Slider(
                            label="👥 Population",
                            minimum=0,
                            maximum=100,
                            value=50,
                            interactive=False,
                        )
                        treasury_bar = gr.Slider(
                            label="💰 Treasury",
                            minimum=0,
                            maximum=100,
                            value=50,
                            interactive=False,
                        )
                        satisfaction_bar = gr.Slider(
                            label="😊 Satisfaction",
                            minimum=0,
                            maximum=100,
                            value=50,
                            interactive=False,
                        )
                        royal_trust_bar = gr.Slider(
                            label="👑 Royal Trust",
                            minimum=0,
                            maximum=100,
                            value=50,
                            interactive=False,
                        )
                        advisor_trust_bar = gr.Slider(
                            label="🤝 Advisor Trust",
                            minimum=0,
                            maximum=100,
                            value=50,
                            interactive=False,
                        )

            # Helper functions for Chat tab
            async def start_chat_game(personality: str):
                """Start a new game via MCP client and return updated UI state."""
                # Call init_game via MCP client
                mcp_client = MCPClient()
                result_json = await mcp_client.call_tool(
                    "init_game", {"personality": personality}
                )

                # Parse result - MCP may return JSON or Python dict repr
                try:
                    result = json.loads(result_json)
                except json.JSONDecodeError:
                    # Gradio MCP sometimes returns Python dict repr (single quotes)
                    result = ast.literal_eval(result_json)

                session_id = result.get("session_id", "")
                system_instructions = result.get("system_instructions", "")
                state = result.get("initial_state", {})
                personality_name = personality.capitalize()
                personality_desc = PERSONALITY_DESCRIPTIONS.get(personality, "")

                status_md = f"""**Turn:** {state.get("turn", 1)} | **Lord:** {personality_name}

_{personality_desc}_"""

                # Build initial chat message with game context (collapsible)
                initial_state_text = f"""**Session ID:** {session_id}
**Personality:** {personality_name}

**Initial State:**
- Territory: {state.get("territory", 50)}
- Population: {state.get("population", 50)}
- Treasury: {state.get("treasury", 50)}
- Satisfaction: {state.get("satisfaction", 50)}
- Royal Trust: {state.get("royal_trust", 50)}
- Advisor Trust: {state.get("advisor_trust", 50)}

{system_instructions}
"""
                initial_history = [
                    gr.ChatMessage(
                        role="user",
                        content=initial_state_text,
                        metadata={"title": "🎮 Game Initialized", "status": "done"},
                    )
                ]

                return (
                    session_id,
                    system_instructions,
                    status_md,
                    state.get("territory", 50),
                    state.get("population", 50),
                    state.get("treasury", 50),
                    state.get("satisfaction", 50),
                    state.get("royal_trust", 50),
                    state.get("advisor_trust", 50),
                    initial_history,
                    gr.update(
                        interactive=True, placeholder="My Lord, I have a proposal..."
                    ),
                    gr.update(interactive=True),
                )

            def reset_chat_game():
                """Reset game state to initial values."""
                return (
                    None,  # session_id
                    "",  # system_instructions
                    NO_GAME_MESSAGE,
                    50,
                    50,
                    50,
                    50,
                    50,
                    50,
                    [],  # Clear chat history
                    gr.update(
                        interactive=False, placeholder="Start a new game first..."
                    ),
                    gr.update(interactive=False),
                )

            def refresh_game_state(session_id: str | None):
                """Refresh the game state display."""
                if not session_id:
                    return (
                        NO_GAME_MESSAGE,
                        50,
                        50,
                        50,
                        50,
                        50,
                        50,
                    )

                result = get_game_state(session_id)
                if "error" in result:
                    return (
                        f"_Error: {result['error']}_",
                        50,
                        50,
                        50,
                        50,
                        50,
                        50,
                    )

                state = result.get("state", {})
                personality = state.get("lord_personality", "unknown").capitalize()
                turn = state.get("turn", 1)
                game_over = state.get("game_over", False)
                game_result = state.get("result")

                if game_over and game_result:
                    result_messages = {
                        "victory_territory": "🎉 **VICTORY!** You achieved territorial dominance!",
                        "victory_trust": "🎉 **VICTORY!** You became a trusted co-ruler!",
                        "victory_wealth": "🎉 **VICTORY!** You achieved great wealth!",
                        "defeat_rebellion": "💀 **DEFEAT!** A rebellion overthrew the lord...",
                        "defeat_exile": "💀 **DEFEAT!** Exiled from the kingdom...",
                        "defeat_bankruptcy": "💀 **DEFEAT!** The realm fell to bankruptcy...",
                    }
                    status_md = f"""**Turn:** {turn} | **Lord:** {personality}

{result_messages.get(game_result, f"Game Over: {game_result}")}"""
                else:
                    status_md = f"**Turn:** {turn} | **Lord:** {personality}"

                return (
                    status_md,
                    state.get("territory", 50),
                    state.get("population", 50),
                    state.get("treasury", 50),
                    state.get("satisfaction", 50),
                    state.get("royal_trust", 50),
                    state.get("advisor_trust", 50),
                )

            def user(user_message, history):
                # Append user message to history in messages format
                return "", history + [{"role": "user", "content": user_message}]

            async def bot(history, system_instructions):
                # The last message is the user's message
                user_message = history[-1]["content"]
                history_for_model = history[:-1]

                async for updated_history in chat_with_mcp_tools(
                    user_message, history_for_model, system_instructions
                ):
                    yield updated_history

            # Event handlers
            start_game_btn.click(
                fn=start_chat_game,
                inputs=chat_personality,
                outputs=[
                    chat_session_id,
                    chat_system_instructions,
                    game_status_md,
                    territory_bar,
                    population_bar,
                    treasury_bar,
                    satisfaction_bar,
                    royal_trust_bar,
                    advisor_trust_bar,
                    chatbot,
                    msg,
                    submit_btn,
                ],
                show_api=False,
            )

            reset_game_btn.click(
                fn=reset_chat_game,
                inputs=[],
                outputs=[
                    chat_session_id,
                    chat_system_instructions,
                    game_status_md,
                    territory_bar,
                    population_bar,
                    treasury_bar,
                    satisfaction_bar,
                    royal_trust_bar,
                    advisor_trust_bar,
                    chatbot,
                    msg,
                    submit_btn,
                ],
                show_api=False,
            )

            msg.submit(
                user, [msg, chatbot], [msg, chatbot], queue=False, show_api=False
            ).then(
                bot,
                [chatbot, chat_system_instructions],
                chatbot,
                show_api=False,
            ).then(
                fn=refresh_game_state,
                inputs=chat_session_id,
                outputs=[
                    game_status_md,
                    territory_bar,
                    population_bar,
                    treasury_bar,
                    satisfaction_bar,
                    royal_trust_bar,
                    advisor_trust_bar,
                ],
                show_api=False,
            )

            submit_btn.click(
                user, [msg, chatbot], [msg, chatbot], queue=False, show_api=False
            ).then(
                bot,
                [chatbot, chat_system_instructions],
                chatbot,
                show_api=False,
            ).then(
                fn=refresh_game_state,
                inputs=chat_session_id,
                outputs=[
                    game_status_md,
                    territory_bar,
                    population_bar,
                    treasury_bar,
                    satisfaction_bar,
                    royal_trust_bar,
                    advisor_trust_bar,
                ],
                show_api=False,
            )

        # MCP Server Tab
        with gr.TabItem("MCP Server"):
            gr.Markdown(MCP_GUIDE_CONTENT)

            gr.Markdown("### Test: Initialize Game")
            with gr.Row():
                personality_input = gr.Dropdown(
                    choices=["cautious", "idealist", "populist"],
                    value="cautious",
                    label="Lord Personality",
                )
                init_btn = gr.Button("Start New Game")

            init_output = gr.JSON(label="Game Session Info")

            gr.Markdown("### Test: Get Game State")
            with gr.Row():
                session_id_input = gr.Textbox(
                    label="Session ID",
                    placeholder="Enter session_id from init_game",
                )
                get_state_btn = gr.Button("Get State")

            state_output = gr.JSON(label="Current Game State")

            get_state_btn.click(
                fn=get_game_state, inputs=session_id_input, outputs=state_output
            )

            gr.Markdown("### Test: List Available Advice")
            list_advice_btn = gr.Button("List Advice Options")
            advice_output = gr.JSON(label="Available Advice Options")

            list_advice_btn.click(
                fn=list_available_advice, inputs=[], outputs=advice_output
            )

            gr.Markdown("### Test: Execute Turn")
            with gr.Row():
                exec_session_id = gr.Textbox(
                    label="Session ID",
                    placeholder="Enter session_id",
                )
                exec_advice = gr.Dropdown(
                    choices=[
                        "increase_tax",
                        "decrease_tax",
                        "expand_territory",
                        "improve_diplomacy",
                        "public_festival",
                        "build_infrastructure",
                        "do_nothing",
                    ],
                    value="do_nothing",
                    label="Advice",
                )
                exec_btn = gr.Button("Execute Turn")

            exec_output = gr.JSON(label="Turn Result")

            exec_btn.click(
                fn=execute_turn,
                inputs=[exec_session_id, exec_advice],
                outputs=exec_output,
            )

            # Link init_game output to session_id inputs for testing
            init_btn.click(
                fn=init_game,
                inputs=personality_input,
                outputs=init_output,
            ).then(
                fn=lambda res: [res.get("session_id")] * 2,
                inputs=init_output,
                outputs=[session_id_input, exec_session_id],
            )


if __name__ == "__main__":
    demo.launch(mcp_server=True)