pages/play.py

import gradio as gr

from support.api_keys_manager import get_required_teams
from support.build_graph import MyGraph
from support.database import save_game_to_db
from support.log_manager import logger
from support.manage_game import Game
from support.game_settings import TEAM_MODEL_PRESETS, JS_BTN
from support.utils import format_messages_as_feed, generate_team_html, plot_game_board_with_guesses


# graph = MyGraph()
TEAM_OPTIONS = list(TEAM_MODEL_PRESETS.keys()) + ["random"]


def validate_api_keys(red_team, blue_team, openai_key, google_key, anthropic_key, hf_key, current_api_keys):
    """Validate that all required API keys are provided."""
    required_teams = get_required_teams(red_team, blue_team)

    # Build the API keys dictionary
    api_keys = {}
    missing_keys = []

    key_mapping = {
        "openai": (openai_key, "OPENAI_API_KEY", "OpenAI"),
        "google": (google_key, "GOOGLE_API_KEY", "Google"),
        "anthropic": (anthropic_key, "ANTHROPIC_API_KEY", "Anthropic"),
        "opensource": (hf_key, "HUGGINGFACEHUB_API_TOKEN", "HuggingFace")
    }

    for team in required_teams:
        key_value, key_name, display_name = key_mapping[team]
        if key_value and key_value.strip():
            api_keys[key_name] = key_value.strip()
        else:
            missing_keys.append(display_name)

    if missing_keys:
        error_msg = f"⚠️ Missing API keys for: {', '.join(missing_keys)}"
        return api_keys, False, error_msg

    return api_keys, True, "✅ All API keys validated!"


with gr.Blocks(fill_width=True) as demo:
    game_state = gr.State()
    players_state = gr.State(value=[])
    board_state = gr.State(value={})
    game_started = gr.State(value=False)
    guessed_words_state = gr.State(value=[])
    messages_state = gr.State(value=[])
    chat_history_state = gr.State(value=[])
    is_human_playing_state = gr.State(value=False)
    turn_state = gr.State(value=0)
    winners_state = gr.State(value=None)
    api_keys_state = gr.State(value={})
    graph_instance = gr.State() 

    # Team selection (visible initially)
    with gr.Row(elem_id="game_setup_row"):
        with gr.Row(elem_id="team_selection_row") as team_selection:
            red_team_dropdown = gr.Dropdown(
                choices=TEAM_OPTIONS,
                value="google",
                label="🔴 Red Team",
                info="Select the model team for the Red side.",
                elem_id="red_team_dropdown",
            )
            blue_team_dropdown = gr.Dropdown(
                choices=TEAM_OPTIONS,
                value="openai",
                label="🔵 Blue Team",
                info="Select the model team for the Blue side.",
                elem_id="blue_team_dropdown"
            )

        with gr.Group(visible=True, elem_id="api_keys_section") as api_keys_section:
            gr.Markdown("### 🔑 API Keys")
            gr.Markdown("*Enter API keys for the selected teams. Keys are only required for active teams.*")

            google_api_input = gr.Textbox(
                label="Google API Key",
                placeholder="AIza...",
                type="password",
                visible=True,  # Default visible since google is selected
                elem_id="google_api_key",
                interactive=True
            )

            openai_api_input = gr.Textbox(
                label="OpenAI API Key",
                placeholder="sk-...",
                type="password",
                visible=True,
                elem_id="openai_api_key",
                interactive=True
            )

            anthropic_api_input = gr.Textbox(
                label="Anthropic API Key",
                placeholder="sk-ant-...",
                type="password",
                visible=False,
                elem_id="anthropic_api_key",
                interactive=True
            )

            hf_api_input = gr.Textbox(
                label="HuggingFace API Token",
                placeholder="hf_...",
                type="password",
                visible=False,
                elem_id="hf_api_key"
            )

            validation_status = gr.Markdown("", visible=False)

            new_game_btn = gr.Button("🔄 Generate New Game", variant="secondary", size="sm", elem_id="new_game_btn", visible=False)

    start_btn = gr.Button(
        "🎲 Generate Teams and Start Playing!",
        variant="primary",
        size="lg",
        elem_id="start_game_btn",
        visible=True
    )

    # Game content (hidden initially)
    with gr.Column(visible=False, elem_id="game_content") as game_content:
        with gr.Accordion("👥 Teams Overview", open=True, elem_id="teams_accordion") as teams_section:
            team_html = gr.HTML("")
            player_info_md = gr.Markdown(value="", elem_id="player_info")

            # Boss control section
            with gr.Row(elem_id="boss_control_row") as play_as_boss_section:
                with gr.Column(scale=1):
                    gr.Markdown("### 🎮 Play as Boss")
                    with gr.Row():
                        red_boss_btn = gr.Button("🔴 Play as Red Boss", size="sm", elem_id="red_boss_btn")
                        blue_boss_btn = gr.Button("🔵 Play as Blue Boss", size="sm", elem_id="blue_boss_btn")

            # Hidden input sections for boss name
            with gr.Column(visible=False, elem_id="red_boss_input") as red_boss_input_section:
                gr.Markdown("#### 👤 Enter Your Name (Red Team)")
                red_boss_name_input = gr.Textbox(
                    placeholder="Your name...",
                    label="Name",
                    show_label=False,
                    elem_id="red_boss_name_input"
                )
                with gr.Row(elem_id="red_boss_actions"):
                    save_red_boss_btn = gr.Button("💾 Save", variant="primary", size="sm", elem_id="save_red_boss_btn")
                    cancel_red_boss_btn = gr.Button("❌ Cancel", size="sm", elem_id="cancel_red_boss_btn")

                with gr.Row(elem_id="red_error"):
                    empty_red = gr.Button("❌ You have to insert a valid name, or cancel", variant="primary", size="sm", visible=False, interactive=False, elem_id="error_red_display")

            with gr.Column(visible=False, elem_id="blue_boss_input") as blue_boss_input_section:
                gr.Markdown("#### 👤 Enter Your Name (Blue Team)")
                blue_boss_name_input = gr.Textbox(
                    placeholder="Your name...",
                    label="Name",
                    show_label=False,
                    elem_id="blue_boss_name_input"
                )
                with gr.Row(elem_id="blue_boss_actions"):
                    save_blue_boss_btn = gr.Button("💾 Save", variant="primary", size="sm", elem_id="save_blue_boss_btn")
                    cancel_blue_boss_btn = gr.Button("❌ Cancel", size="sm", elem_id="cancel_blue_boss_btn")

                with gr.Row(elem_id="blue_error"):
                    empty_blue = gr.Button("❌ You have to insert a valid name, or cancel", variant="primary", size="sm", visible=False, interactive=False, elem_id="error_blue_display")

        with gr.Row(elem_id="game_row", equal_height=True):
            with gr.Column(elem_id="board_column"):
                gr.Markdown("### 🎯 Game Board", elem_id="board_header")
                board_plot = gr.Image(
                    label="Game Board",
                    show_label=False,
                    elem_id="game_board_img",
                )

            with gr.Column(elem_id="chat_column"):
                gr.Markdown("### 💬 Game Feed", elem_id="chat_header")

                # Replace ChatInterface with custom HTML feed
                message_feed = gr.HTML(value="", elem_id="message_feed_container", autoscroll=True)

                with gr.Row(elem_id="input_game_section"):
                    with gr.Row(visible=False) as input_game_values:
                        user_input = gr.Textbox(
                            placeholder="Type your message...",
                            show_label=True,
                            scale=9,
                            visible=True,
                            label="Your Clue"
                        )
                        dropdown = gr.Dropdown(
                            choices=[i for i in range(1, 10)],
                            label="Clue number",
                            value=1,
                            interactive=True
                        )
                    send_btn = gr.Button("Send", elem_id="send_message_btn")

    # Helper functions
    def start_game(red_team_choice, blue_team_choice, openai_key, google_key, anthropic_key, hf_key):
        # Validate API keys first
        api_keys, is_valid, message = validate_api_keys(
            red_team_choice, blue_team_choice, openai_key, google_key, anthropic_key, hf_key, {}
        )

        if not is_valid:
            # Return error state without starting game
            return {
                validation_status: gr.update(visible=True, value=message),
                api_keys_state: api_keys,
                game_state: None,
                new_game_btn: gr.update(visible=False),
                start_btn: gr.update(visible=True),
                game_content: gr.update(visible=False),
                team_html: "",
                board_plot: None,
                game_started: False,
                players_state: [],
                board_state: {},
                messages_state: [],
                graph_instance: None  # ADD THIS
            }

        # Create game with validated API keys
        game = Game(red_team=red_team_choice, blue_team=blue_team_choice, api_keys=api_keys)
        starting_team = game.board.get('starting_team', 'red')
        html = generate_team_html(game.players, starting_team)
        board_img = plot_game_board_with_guesses(game.board)

        # CREATE GRAPH INSTANCE FOR THIS SESSION
        graph = MyGraph()  # ADD THIS

        return {
            game_state: game,
            new_game_btn: gr.update(visible=True),
            start_btn: gr.update(visible=False),
            game_content: gr.update(visible=True),
            team_html: html,
            board_plot: board_img,
            game_started: True,
            players_state: game.players,
            board_state: game.board,
            messages_state: [],
            api_keys_state: api_keys,
            validation_status: gr.update(visible=True, value="✅ Game started successfully!"),
            graph_instance: graph  # ADD THIS
        }

    def generate_new_game(red_team_choice, blue_team_choice, openai_key, google_key, anthropic_key, hf_key, current_api_keys):
        # Validate API keys first
        api_keys, is_valid, message = validate_api_keys(
            red_team_choice, blue_team_choice, openai_key, google_key, anthropic_key, hf_key, current_api_keys
        )

        if not is_valid:
            # Raise error to show in UI
            raise gr.Error(message)

        # Create game with validated API keys
        game = Game(red_team=red_team_choice, blue_team=blue_team_choice, api_keys=api_keys)
        starting_team = game.board.get('starting_team', 'red')
        html = generate_team_html(game.players, starting_team)
        board_img = plot_game_board_with_guesses(game.board)

        # CREATE NEW GRAPH INSTANCE
        graph = MyGraph()  # ADD THIS

        return game, html, board_img, game.players, game.board, [], gr.update(visible=True), api_keys, graph  # ADD graph to return

    def show_red_boss_input():
        return gr.Column(visible=True)

    def show_blue_boss_input():
        return gr.Column(visible=True)

    def hide_red_boss_input():
        return gr.Column(visible=False)

    def hide_blue_boss_input():
        return gr.Column(visible=False)

    def update_api_key_inputs(red_team, blue_team):
        """Update visibility and state of API key inputs based on team selection."""

        required_teams = get_required_teams(red_team, blue_team)

        # Return visibility states for each team's API input
        updates = {}
        for team in ["openai", "google", "anthropic", "opensource"]:
            updates[team] = gr.update(visible=team in required_teams)

        return (
            updates["openai"],
            updates["google"],
            updates["anthropic"],
            updates["opensource"]
        )

    def update_boss_player(team_color, name, players, board, is_human_playing):
        """Update the boss player with human name and model"""
        if not name or not name.strip():
            return generate_team_html(players, board.get('starting_team')), players, gr.Column(visible=True), gr.Row(visible=True), gr.Row(visible=False), is_human_playing, gr.Button(visible=True)

        for player in players:
            if player.team == team_color and player.role == "boss":
                player.name = name.strip()
                player.model_name = "Human brain"
                break

        # Regenerate HTML with updated player info
        starting_team = board.get('starting_team')
        html = generate_team_html(players, starting_team, True)
        is_human_playing = True
        return html, players, gr.Column(visible=False), gr.Row(visible=False), gr.Row(visible=True), is_human_playing, gr.Button(visible=False)

    async def process_message(user_msg, messages, players, board, dropdown, guessed, chat_history, is_human_playing, turn, graph):  # ADD graph parameter
        """Process user message and stream responses"""

        accumulated_messages = list(messages) if messages else []
        async for new_msg, guessed, updated_board, updated_chat_history, winners in graph.stream_graph(user_msg, messages, players, board, dropdown, guessed, chat_history, is_human_playing, turn):

            if len(new_msg) > len(accumulated_messages):
                # Add only the new messages that weren't there before
                accumulated_messages = list(new_msg)
            else:
                # Update existing messages (for streaming updates)
                accumulated_messages = list(new_msg)

            messages = accumulated_messages
            feed_html = format_messages_as_feed(accumulated_messages, players, winners)
            yield feed_html, guessed, messages, updated_board, updated_chat_history, winners

    def update_plot(guessed_words, board, game_state):
        # Don't update if no game is active
        if game_state is None:
            return gr.update()  # No update

        if not guessed_words or not board:
            return gr.update()  # No update

        logger.info(f"Calling update_plot: {guessed_words}")
        board_img = plot_game_board_with_guesses(board, guessed_words)
        return board_img

    def deactivate_send(game, winner_and_score):
        if winner_and_score is None or game is None:
            return gr.Button(visible=True)
        else:
            return gr.Button(visible=False)

    def send_to_database(game, winner_and_score):

        """Save game results to database"""
        if winner_and_score is None or game is None:
            return

        winner_team = winner_and_score[0]
        scores = winner_and_score[1]
        reason = winner_and_score[2] if len(winner_and_score) > 2 else None

        # Log to console
        logger.info(f"\n{'='*50}")
        logger.info("GAME COMPLETED - Saving to Database")
        logger.info(f"{'='*50}")
        logger.info(f"Winner: {winner_team.upper()} team")
        logger.info(f"Scores - Red: {scores[0]}, Blue: {scores[1]}")
        if reason:
            logger.info(f"Reason: {reason}")
        logger.info("\nTeam Compositions:")
        logger.info(f"Red Team ({game.red_team_choice}):")
        for player in [p for p in game.players if p.team == 'red']:
            logger.info(f"  - {player.role}: {player.name} ({player.model_name})")
        logger.info(f"Blue Team ({game.blue_team_choice}):")
        for player in [p for p in game.players if p.team == 'blue']:
            logger.info(f"  - {player.role}: {player.name} ({player.model_name})")

        # Save to database
        try:
            game_id = save_game_to_db(game, winner_and_score)
            logger.info(f"\n✅ Game saved successfully! (ID: {game_id})")
            logger.info(f"{'='*50}\n")
        except Exception as e:
            logger.info(f"\n❌ Error saving game to database: {e}")
            logger.info(f"{'='*50}\n")


    # Event handlers
    start_btn.click(
        fn=start_game,
        inputs=[red_team_dropdown, blue_team_dropdown, openai_api_input, google_api_input, anthropic_api_input, hf_api_input],
        outputs=[game_state, new_game_btn, start_btn, game_content, team_html, board_plot, game_started, players_state, board_state, messages_state, api_keys_state, validation_status, graph_instance]  # ADD graph_instance
    )

    # Update new_game_btn.click
    new_game_btn.click(
        fn=generate_new_game,
        inputs=[
            red_team_dropdown,
            blue_team_dropdown,
            openai_api_input,
            google_api_input,
            anthropic_api_input,
            hf_api_input,
            api_keys_state
        ],
        outputs=[
            game_state,
            team_html,
            board_plot,
            players_state,
            board_state,
            messages_state,
            play_as_boss_section,
            api_keys_state,
            graph_instance  # ADD THIS
        ]
    )

    red_team_dropdown.change(
        fn=update_api_key_inputs,
        inputs=[red_team_dropdown, blue_team_dropdown],
        outputs=[openai_api_input, google_api_input, anthropic_api_input, hf_api_input]
    )

    blue_team_dropdown.change(
        fn=update_api_key_inputs,
        inputs=[red_team_dropdown, blue_team_dropdown],
        outputs=[openai_api_input, google_api_input, anthropic_api_input, hf_api_input]
    )

    # Show/hide boss input sections
    red_boss_btn.click(
        fn=show_red_boss_input,
        outputs=[red_boss_input_section]
    )

    blue_boss_btn.click(
        fn=show_blue_boss_input,
        outputs=[blue_boss_input_section]
    )

    cancel_red_boss_btn.click(
        fn=hide_red_boss_input,
        outputs=[red_boss_input_section]
    )

    cancel_blue_boss_btn.click(
        fn=hide_blue_boss_input,
        outputs=[blue_boss_input_section]
    )

    # Save boss name handlers
    save_red_boss_btn.click(
        fn=lambda name, players, board, is_human_playing: update_boss_player("red", name, players, board, is_human_playing),
        inputs=[red_boss_name_input, players_state, board_state, is_human_playing_state],
        outputs=[team_html, players_state, red_boss_input_section, play_as_boss_section, input_game_values, is_human_playing_state, empty_red]
    ).then(
        fn=lambda: "",
        outputs=[red_boss_name_input]
    )

    save_blue_boss_btn.click(
        fn=lambda name, players, board, is_human_playing: update_boss_player("blue", name, players, board, is_human_playing),
        inputs=[blue_boss_name_input, players_state, board_state, is_human_playing_state],
        outputs=[team_html, players_state, blue_boss_input_section, play_as_boss_section, input_game_values, is_human_playing_state, empty_blue]
    ).then(
        fn=lambda: "",
        outputs=[blue_boss_name_input]
    )

    send_btn.click(
        fn=lambda: gr.Button(interactive=False),
        outputs=[send_btn],
        js=JS_BTN
    ).then(
        fn=process_message,
        inputs=[user_input, messages_state, players_state, board_state, dropdown, guessed_words_state, chat_history_state, is_human_playing_state, turn_state, graph_instance],  # ADD graph_instance
        outputs=[message_feed, guessed_words_state, messages_state, board_state, chat_history_state, winners_state]
    ).then(
        fn=lambda: "",
        outputs=[user_input]
    ).then(
        fn=lambda turn: turn + 1,
        inputs=[turn_state],
        outputs=[turn_state]
    ).then(
        fn=lambda: gr.Button(interactive=True),
        outputs=[send_btn]
    )

    guessed_words_state.change(
        fn=update_plot,
        inputs=[guessed_words_state, board_state, game_state],  # Remove board_plot from inputs
        outputs=[board_plot]
    )

    winners_state.change(
        fn=send_to_database,
        inputs=[game_state, winners_state]
    ).then(
        fn=deactivate_send,
        inputs=[game_state, winners_state],
        outputs=[send_btn]
    )

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