support/utils.py

import io
import markdown

import matplotlib
matplotlib.use('Agg')
import matplotlib.pyplot as plt
import matplotlib.patches as mpatches
import os
import regex
import tempfile
import uuid

from PIL import Image
from support.log_manager import logger


def display_model_name(model_name):
    """Format model name for display"""
    model_map = {
        "gpt-5": "GPT-5",
        "gpt-5-mini": "GPT-5 Mini",
        "gpt-5-nano": "GPT-5 Nano",
        "gpt-4.1-nano": "GPT-4.1 Nano",
        "gemini-2.5-pro": "Gemini 2.5 Pro",
        "gemini-2.5-flash": "Gemini 2.5 Flash",
        "gemini-2.0-flash-001": "Gemini 2.0 Flash",
        "gemini-2.0-flash-lite-001": "Gemini 2.0 Flash Lite",
        "claude-sonnet-4-5-20250929": "Claude Sonnet 4.5",
        "claude-3-7-sonnet-20250219": "Claude 3.7 Sonnet",
        "claude-3-5-haiku-20241022": "Claude 3.5 Haiku",
        "claude-3-haiku-20240307": "Claude 3 Haiku",
        "deepseek-ai/DeepSeek-V3.1": "DeepSeek V3.1",
        "deepseek-ai/DeepSeek-R1": "DeepSeek R1",
        "Qwen/Qwen3-235B-A22B-Instruct-2507": "Qwen3 235B",
        "openai/gpt-oss-120b": "GPT-OSS 120B",
        "openai/gpt-oss-20b": "GPT-OSS 20B",
        "moonshotai/Kimi-K2-Thinking": "Kimi K2 Thinking",
        "human brain": "Human Brain"
    }
    return model_map.get(model_name, model_name)


def format_messages_as_feed(messages, players, winner_and_score):
    """Convert messages to Instagram-like feed HTML"""

    if winner_and_score:
        winner_team = winner_and_score[0]
        fireworks_animation = f"""
            <div class='firework message_{winner_team}'></div>
        """
    else:
        fireworks_animation = ""

    emoji_pattern = regex.compile(r'^\p{Emoji_Presentation}|\p{Emoji}\uFE0F')
    html = '<div id="message_feed">'

    # Initialize markdown with proper extensions
    md = markdown.Markdown(extensions=[
        'fenced_code',
        'tables',
        'nl2br',
        'sane_lists'  # Better list handling
    ])

    # Create a mapping of sender identifiers to players
    player_map = {}
    for p in players:
        if p.team and p.role:
            key = f"{p.team}_{p.role}"
            if key not in player_map:
                player_map[key] = []
            player_map[key].append(p)

            if p.role == "player":
                agent_num = len(player_map[key])
                player_map[f"{p.team}_agent_{agent_num}"] = p

    for msg in messages:
        content = getattr(msg, 'content', '')

        if not content or not str(content).strip():
            continue

        metadata = getattr(msg, 'metadata', {}) or {}
        title = metadata.get('title', '') if isinstance(metadata, dict) else ''
        sender = metadata.get('sender', '') if isinstance(metadata, dict) else ''

        # Find the player for this sender
        player = None
        team = ''
        avatar = 'assets/robot.png'
        display_sender = ''

        if sender:
            sender_lower = sender.lower()

            if sender_lower in player_map:
                if isinstance(player_map[sender_lower], list):
                    player = player_map[sender_lower][0]
                else:
                    player = player_map[sender_lower]

            if player:
                avatar = f"/gradio_api/file={player.avatar}"
                display_sender = f"{player.name} ({player.role.title()})"
                team = player.team
            else:
                display_sender = sender.replace('_', ' ').title()

                if 'red' in sender_lower:
                    team = 'red'
                elif 'blue' in sender_lower:
                    team = 'blue'
                elif 'judge' in sender_lower:
                    avatar = "/gradio_api/file=assets/avatars/judge.png"
                    team = 'judge'

        # Determine card class based on team and role
        card_class = 'message_card'

        if team == 'red':
            card_class += ' message_red'
        elif team == 'blue':
            card_class += ' message_blue'
        elif team == 'judge' or 'judge' in sender.lower():
            card_class += ' message_judge'

        if player and player.role:
            if player.role == 'boss':
                card_class += ' message_boss'
            elif player.role == 'captain':
                card_class += ' message_captain'
            elif player.role == 'player':
                card_class += ' message_agent'

        if 'Thinking' in title:
            card_class += ' message_thinking'
        elif '🛠️' in title or 'tool' in title.lower() or 'Using' in content:
            card_class += ' message_tool'

        # Extract icon from title if present
        icon = ''
        clean_title = title
        if title:
            parts = title.split(' ', 1)
            if len(parts) > 1 and regex.match(emoji_pattern, parts[0]):
                icon = parts[0]
                clean_title = parts[1] if len(parts) > 1 else ''

        if not clean_title:
            icon = "✉️"

        # IMPORTANT: Reset markdown parser and convert content
        md.reset()
        content_html = md.convert(str(content))

        html += f'''
        <div class="{card_class}">
            <div class="message_header">
                <img src="{avatar}" class="message_avatar" onerror="this.style.display='none'"/>
                <div class="message_header_text">
                    <span class="message_sender">{display_sender}</span>
                    <span class="message_title">{clean_title or 'Message'}</span>
                </div>
                <span class="message_icon">{icon}</span>
            </div>
            <div class="message_content">{content_html}</div>
        </div>
        '''

    html += '</div>'
    if fireworks_animation:
        html = f'{html}{fireworks_animation}'
    return html


def generate_team_html(players, starting_team=None, is_human_playing=False):
    """Generate improved team cards with working avatars and starting team indicator"""

    role_order = {"boss": 0, "captain": 1, "player": 2}

    red_players = sorted(
        [p for p in players if p.team == "red"],
        key=lambda p: role_order.get(p.role, 99)
    )

    blue_players = sorted(
        [p for p in players if p.team == "blue"],
        key=lambda p: role_order.get(p.role, 99)
    )

    def create_team_card(team_players, team_color, is_starting):
        team_name = team_color.upper()
        emoji = "🔴" if team_color == "red" else "🔵"

        # Add starting team badge if applicable
        starting_badge = ""
        if is_starting:
            starting_badge = "<div class='starting-badge'>🎯 STARTING TEAM</div>"
        else:
            starting_badge = "<div class='starting-badge'>⚔️ OPPONENT TEAM</div>"

        players_html = ""
        for p in team_players:
            role_badge = f"<span class='role-badge role-{p.role}'>{p.role.upper()}</span>"
            # Use Gradio's file serving path
            avatar_path = f"/gradio_api/file={p.avatar}"

            human_indicator = ""
            if not is_human_playing:
                if p.role == "boss":
                    human_indicator = f"""
                    <button class='play-as-boss-btn' onclick='showBossNameInput("{team_color}")'>
                        🎮 Play as Boss
                    </button>
                    """
            if p.role == "boss" and p.model_name == "Human brain":
                human_indicator = "<span class='human-indicator'>👤 HUMAN PLAYER</span>"

            players_html += f"""
            <div class='player-row'>
                <img src='{avatar_path}' class='player-avatar' onerror="this.style.display='none'">
                <div class='player-details'>
                    <div class='player-name'>{p.name}</div>
                    <div class='player-meta'>
                        {role_badge}
                        <span class='player-model'>{display_model_name(p.model_name)}</span>
                        {human_indicator}
                    </div>
                </div>
            </div>
            """

        return f"""
        <div class='team-card team-{team_color}'>
            <div class='team-header'>
                <h3>{emoji} {team_name} TEAM</h3>
                {starting_badge}
            </div>
            <div class='team-body'>
                {players_html}
            </div>
        </div>
        """

    red_card = create_team_card(red_players, "red", starting_team == "red")
    blue_card = create_team_card(blue_players, "blue", starting_team == "blue")

    return f"""
    <div class='teams-container'>
        {red_card}
        <div class='vs-divider'>
            <div class='vs-text'>⚔️</div>
        </div>
        {blue_card}
    </div>
    """


def format_game_history_html(game_history):
    """Format game history as HTML with team cards similar to play page"""
    if not game_history:
        return "<div style='text-align: center; padding: 40px;'><h3>No games played yet.</h3></div>"

    games_html = []

    for game in game_history:
        # Sort players by role
        role_order = {"boss": 0, "captain": 1, "player": 2}
        red_players = sorted(
            [p for p in game['players'] if p['team'] == 'red'],
            key=lambda p: role_order.get(p['role'], 99)
        )
        blue_players = sorted(
            [p for p in game['players'] if p['team'] == 'blue'],
            key=lambda p: role_order.get(p['role'], 99)
        )

        # Determine winner and create result message
        winner = game['winner']
        red_score = game['red_score']
        blue_score = game['blue_score']
        reason = game['reason']

        # Create result message
        if reason == 'killer':
            loser = 'red' if winner == 'blue' else 'blue'
            result_emoji = "💀"
            result_text = f"{loser.upper()} team hit the killer word!"
            result_class = "result-killer"
        else:
            result_emoji = "🏆"
            result_text = f"{winner.upper()} team won!"
            result_class = f"result-{winner}"

        # Score display
        # red_score_class = "score-winner" if winner == "red" else "score-loser"
        # blue_score_class = "score-winner" if winner == "blue" else "score-loser"

        def create_team_card_history(team_players, team_color, is_winner):
            team_name = team_color.upper()
            emoji = "🔴" if team_color == "red" else "🔵"
            score = red_score if team_color == "red" else blue_score

            winner_badge = ""
            if is_winner:
                winner_badge = "<div class='winner-badge'>🏆 WINNER</div>"
            else:
                winner_badge = "<div class='winner-badge'>😭 LOSER</div>"

            players_html = ""
            for p in team_players:
                role_badge = f"<span class='role-badge role-{p['role']}'>{p['role'].upper()}</span>"

                human_indicator = ""
                if p['is_human']:
                    human_indicator = "<span class='human-indicator'>👤 HUMAN</span>"

                players_html += f"""
                <div class='player-row'>
                    <div class='player-details'>
                        <div class='player-name'>{p['name']}</div>
                        <div class='player-meta'>
                            {role_badge}
                            <span class='player-model'>{display_model_name(p['model'])}</span>
                            {human_indicator}
                        </div>
                    </div>
                </div>
                """

            score_class = "score-winner" if is_winner else "score-loser"

            return f"""
            <div class='team-card team-{team_color}'>
                <div class='team-header'>
                    <h3>{emoji} {team_name} TEAM</h3>
                    {winner_badge}
                </div>
                <div class='team-score {score_class}'>
                    <span class='score-label'>Remaining Words:</span>
                    <span class='score-value'>{score}</span>
                </div>
                <div class='team-body'>
                    {players_html}
                </div>
            </div>
            """

        red_card = create_team_card_history(red_players, "red", winner == "red")
        blue_card = create_team_card_history(blue_players, "blue", winner == "blue")

        game_html = f"""
        <div class='game-history-card'>
            <div class='game-header'>
                <div class='game-info'>
                    <span class='game-number'>Game #{game['id']}</span>
                    <span class='game-date'>📅 {game['timestamp']}</span>
                </div>
                <div class='game-result {result_class}'>
                    {result_emoji} {result_text}
                </div>
            </div>
            <div class='teams-container'>
                {red_card}
                <div class='vs-divider'>
                    <div class='vs-text'>VS</div>
                </div>
                {blue_card}
            </div>
        </div>
        """

        games_html.append(game_html)

    return "".join(games_html)


def plot_game_board_with_guesses(board, guessed_words=None):
    """
    Plot the game board with crosses or marks over guessed words.
    """

    logger.info(f"Board: {board}")
    logger.info(f"Guessed words: {guessed_words}")

    color_map = {
        'red': '#dc3545',
        'blue': '#0d6efd',
        'neutral': '#d4c5b9',
        'killer': '#212529'
    }

    guessed_words = set(guessed_words or [])
    word_color_pairs = board.get('word_color_pairs', [])
    if len(word_color_pairs) < 25:
        word_color_pairs += [("???", 'neutral')] * (25 - len(word_color_pairs))

    # Create a new figure explicitly (don't rely on global state)
    fig = plt.figure(figsize=(12, 12))
    ax = fig.add_subplot(111)
    ax.set_xlim(-0.5, 4.5)
    ax.set_ylim(-0.5, 4.5)
    ax.axis('off')
    ax.set_aspect('equal')

    idx = 0
    for i in range(5):
        for j in range(5):
            if idx >= len(word_color_pairs):
                break

            word, color = word_color_pairs[idx]
            bg_color = color_map.get(color, '#e9ecef')

            rect = mpatches.FancyBboxPatch(
                (j - 0.48, i - 0.48), 0.96, 0.96,
                boxstyle="round,pad=0.02",
                facecolor=bg_color,
                edgecolor='#495057',
                linewidth=2.5,
                zorder=1
            )
            ax.add_patch(rect)

            ax.text(
                j, i, word,
                ha="center", va="center",
                fontsize=15 if len(word) <= 8 else 13,
                fontweight='bold',
                color='white',
                zorder=2,
                wrap=True
            )

            if word in guessed_words:
                ax.plot(
                    [j - 0.4, j + 0.4], [i - 0.4, i + 0.4],
                    color='black', linewidth=3, alpha=0.8, zorder=3
                )
                ax.plot(
                    [j - 0.4, j + 0.4], [i + 0.4, i - 0.4],
                    color='black', linewidth=3, alpha=0.8, zorder=3
                )

            idx += 1

    ax.invert_yaxis()
    starting_team = board.get('starting_team', 'red')
    title_color = color_map[starting_team]

    fig.suptitle(
        f'{starting_team.upper()} TEAM STARTS',
        fontsize=18,
        fontweight='bold',
        color=title_color,
        y=0.98
    )

    fig.tight_layout()

    # Create a unique temporary file for this specific request
    temp_file = tempfile.NamedTemporaryFile(delete=False, suffix='.png', prefix=f'board_{uuid.uuid4().hex}_')
    temp_filename = temp_file.name
    temp_file.close()

    try:
        # Save to the unique temporary file
        fig.savefig(temp_filename, format="png", bbox_inches="tight", dpi=120, facecolor='white')
        plt.close(fig)  # Close the figure immediately

        # Load the image from the file
        img = Image.open(temp_filename)
        img.load()  # Load into memory

        return img
    finally:
        # Clean up the temporary file
        try:
            os.unlink(temp_filename)
        except:
            pass