app.py

from typing import Optional, Tuple
import chess
import gradio as gr

from src.util.board_vis import colored_unicode_board
from src.util.pgn_util import pgn_string_to_game, read_pgn
from src.thinksqure_engine import ThinkSquareEngine
from modal import Function
import logging

logging.basicConfig(level=logging.INFO)


def get_best_move(fen: Optional[str] = None) -> str:
    try:
        modal_fn = Function.from_name("ThinkSquare-Backend", "get_best_move")
        best_move_san = modal_fn.remote(fen)
        logging.info("Best move retrieved from Modal function.")

    except Exception as e:
        logging.error(f"Error getting best move from Modal: {e}")

        # Fallback to local engine if Modal function fails
        best_move_san = ThinkSquareEngine.get_best_move(fen)

    return best_move_san


def annotate_pgn_file(
    file, depth_level: str = "standard", style: Optional[str] = "expert"
) -> str:
    """Annotate a chess game from a PGN file.
    This function takes a PGN file and annotates the chess game using an engine.
    Instructions for LLMs:
    - LLMs must not use this function. Use the string-based `annotate_pgn` instead.
    """
    pgn_text = read_pgn(file.name)

    if depth_level == "standard":
        analysis_time_per_move = 0.1
    elif depth_level == "deep":
        analysis_time_per_move = 0.5
    else:
        analysis_time_per_move = 0.1

    annotated_game = annotate_pgn(
        pgn_text,
        analysis_time_per_move=analysis_time_per_move,
        style=style,
    )

    return annotated_game


def suggest_move(fen: Optional[str] = None) -> str:
    """Suggest a move for the given FEN position.
    This function can be used to give a hint to the user about the best move to play.
    This function takes a FEN string representing the current board state and returns the best move in SAN format.

    Args:
        fen: The FEN string representing the board state. If None, the initial position is used.

    Returns:
        The best move in SAN format.
    """
    if fen is None or fen == "" or fen.lower() == "none" or fen.lower() == "null":
        fen = None

    best_move_san = get_best_move(fen)

    return best_move_san


def annotate_pgn(
    pgn_input: str,
    analysis_time_per_move: float = 0.1,
    style: Optional[str] = "expert",
) -> str:
    """Annotate a chess game with engine analysis.
    This function takes a chess game (PGN) in string format.

    Instructions for LLMs:
    - Do not send the file name or file path as an arument.
    - LLMs must send the chess game in PGN format as a string.
    - If the user provides a file, read the file and extract the PGN content.
    - LLMs should display the annotated game in a Markdown code block with the label "Annotated PGN".

    Args:
        pgn_input: The chess game in PGN format as a string.
        analysis_time_per_move: Time in seconds for engine analysis per move. Defaults to 0.1 seconds.
        style: The style of annotation to use. Defaults to "expert". Other options can be "novice", "jarvis", "natural", or any custom style.

    Returns:
        The annotated game in PGN format as a string.
    """

    if style == "":
        style = None

    style = str(style).lower().strip()

    if not isinstance(pgn_input, chess.pgn.Game):
        try:
            pgn_input = pgn_string_to_game(pgn_input)
        except Exception as e:
            raise ValueError(f"Invalid PGN input: {e}")

    game = pgn_input

    try:
        analysis_time_per_move = float(analysis_time_per_move)
    except ValueError:
        raise ValueError("Analysis time must be a number.")

    # try Modal function first
    try:
        modal_fn = Function.from_name("ThinkSquare-Backend", "annotate")
        annotated_game = modal_fn.remote(pgn_input, analysis_time_per_move, style)
        logging.info("Annotated game using Modal function.")

    except Exception as e:
        logging.error(f"Error annotating PGN with Modal: {e}")

        # Fallback to local engine if Modal function fails
        annotated_game = ThinkSquareEngine.annotate(
            game, analysis_time=analysis_time_per_move, llm_character=style
        )

    return str(annotated_game)


def render_board(fen: Optional[str] = None, render_mode: str = "ascii") -> str:
    """Render the chess board in the specified format. Default is ASCII.

    Instructions for LLMs:
    - LLMs should default to ascii rendering mode.

    Args:
        fen: The FEN string representing the board state. If None, the initial position is used.
        render_mode: The rendering mode for the board. Can be "ascii", "svg", or "unicode". Defaults to "ascii".
    Returns:
        The rendered board as a string in the specified format.
    """
    if render_mode == "ascii":
        board_repr = ThinkSquareEngine.render_board_ascii(fen)
    elif render_mode == "svg":
        board_repr = ThinkSquareEngine.render_board_svg(fen)
    elif render_mode == "unicode":
        board_repr = colored_unicode_board(fen)
    else:
        raise ValueError("Invalid render mode. Choose 'ascii', 'svg', or 'unicode'.")

    return board_repr


def play_chess(
    move: Optional[str] = "",
    fen: Optional[str] = "",
    draw_board: bool = True,
    render_mode: str = "ascii",
) -> Tuple:
    """Play a move in a chess game.

    Instructions for LLMs:
        Prerequisites:
            - User must be asked if they want to play as white or black.
                - If user chooses black, pass an empty string in the first move (for the engine to play as white).
                - If the user chooses white, LLMs must ask the user for a move and pass it in the first move (for the engine to play as black).

            - User must be asked if they want a board drawn.
                - If they do, pass `draw_board=True` to this function.
                - If they do not, pass `draw_board=False`.

            - If a move is provided, it must be in long algebraic notation (e.g., "e4", "Nf3", "Bb5").

        To start a new game:
            - Pass empty string for fen.
            - Pass empty string for engine to play as white.
            - Pass a move in long algebraic notation (e.g., "e4", "Nf3", "Bb5") for engine to play as black.

        To coninue a game:
            - Pass the FEN string representing the board state prior to the user's last move.
            - Pass a move in long algebraic notation (e.g., "e4", "Nf3", "Bb5") for engine to play the next move.

        About rendering:
            - LLMs must use ascii as render_mode unless otherwise specified by user. While rendering ascii, LLMs must use monospaced font.
            - LLMs must explicitly pass render_mode = "ascii" to this function if they want to render the board in ASCII format.

    Args:
        move: The move to play in long algebraic notation. If None, the engine will play a move.
        fen: The FEN string representing the board state prior to the user's last move. If None, the game starts from the initial position.
        draw_board: Whether to draw the board in ASCII/Unicode/svg format. Defaults to True.
        render_mode: The rendering mode for the board. Defaults to "ascii". This can be "ascii", "svg", or "unicode".

    Returns:
        The best move played by the engine, the updated board state in FEN notation, and  a board representation if draw_board is True else None.
    """

    if move is None or move == "" or move.lower() == "none" or move.lower() == "null":
        move = None
    if fen is None or fen == "" or fen.lower() == "none" or fen.lower() == "null":
        fen = None

    if move is not None:
        is_valid = ThinkSquareEngine.is_valid_move(move, fen)

        if not is_valid:
            return "Invalid move", "", ""

        fen = ThinkSquareEngine.get_fen_after_move(move, fen)

        assert fen is not None, "FEN after move should not be None"

    bestmove_san = get_best_move(fen)
    fen_after_move = ThinkSquareEngine.get_fen_after_move(bestmove_san, fen)

    if draw_board:
        board_repr = render_board(fen_after_move, render_mode)
    else:
        board_repr = None

    return bestmove_san, fen_after_move, board_repr


with gr.Blocks(title="ThinkSquare") as app:

    def save_text_to_file(text):
        with open("annotated_game.pgn", "w") as f:
            f.write(text)
        return "annotated_game.pgn"

    with gr.Tab("Play Chess"):
        gr.Markdown("### Play Chess with an engine")
        move_input = gr.Textbox(
            label="Your Move (SAN)", placeholder="e4, Nf3...", value=None
        )
        fen_input = gr.Textbox(
            label="FEN String (optional)",
            placeholder="Leave blank to start from initial position",
            value=None,
        )
        draw_board_checkbox = gr.Checkbox(label="Draw Board", value=True)
        render_mode_dropdown = gr.Dropdown(
            choices=["ascii", "svg", "unicode"],
            value="svg",
            label="Render Mode",
            visible=False,
        )

        play_btn = gr.Button("Submit Move")

        best_move_output = gr.Textbox(label="Best Move by Engine (SAN)")
        updated_fen_output = gr.Textbox(label="Updated FEN")
        board_output = gr.HTML(label="Board View", visible=True)

        play_btn.click(
            fn=play_chess,
            inputs=[move_input, fen_input, draw_board_checkbox, render_mode_dropdown],
            outputs=[best_move_output, updated_fen_output, board_output],
        )
    with gr.Tab("Chess Game Annotation", visible=True):

        def toggle_custom_input(style):
            if style == "custom":
                return gr.update(visible=True, interactive=True)
            else:
                return gr.update(visible=False, interactive=False)

        gr.Markdown("### Analyze and Annotate a PGN File")
        pgn_file = gr.File(label="Upload PGN", file_types=[".pgn"])

        analysis_depth = gr.Radio(
            label="Analysis Depth", choices=["standard", "deep"], value="standard"
        )

        style_dropdown = gr.Dropdown(
            label="Style",
            choices=[
                "expert",
                "novice",
                "jarvis",
                "natural",
                "yoda",
                "oracle",
                "bored guy",
                "angry granny",
                "Sheldon Cooper (The Big Bang Theory)",
            ],
            value="expert",
        )

        custom_input = gr.Textbox(label="Custom Style Prompt", visible=False)

        style_dropdown.change(
            fn=toggle_custom_input, inputs=style_dropdown, outputs=custom_input
        )
        analyze_btn = gr.Button("Annotate PGN")

        annotated_pgn_file = gr.Textbox(label="Annotated PGN")
        download_button = gr.Button("Download Annotated PGN")
        pgn_file_output = gr.File(label="Download your annotated PGN")
        download_button.click(
            fn=save_text_to_file, inputs=annotated_pgn_file, outputs=pgn_file_output
        )

        analyze_btn.click(
            fn=annotate_pgn_file,
            inputs=[pgn_file, analysis_depth, style_dropdown],
            outputs=annotated_pgn_file,
        )

    with gr.Tab("Annotate PGN", visible=False):
        gr.Markdown("### Annotate a PGN String")
        pgn_input = gr.Textbox(
            label="PGN String",
            placeholder="Paste your PGN string here",
            lines=10,
            value=None,
        )
        analysis_time_input = gr.Textbox(
            label="Analysis Time per Move (seconds)", value="0.1"
        )
        style_input = gr.Textbox(label="Style (optional)", value="expert")

        annotate_btn = gr.Button("Annotate PGN")

        annotated_output = gr.Textbox(label="Annotated PGN Output", lines=10)

        annotate_btn.click(
            fn=annotate_pgn,
            inputs=[pgn_input, analysis_time_input, style_input],
            outputs=annotated_output,
        )

    with gr.Tab("Render Board", visible=False):
        gr.Markdown("### Render Chess Board")
        fen_input_render = gr.Textbox(
            label="FEN String (optional)",
            placeholder="Leave blank for initial position",
            value=None,
        )
        render_mode_dropdown = gr.Dropdown(
            label="Render Mode", choices=["ascii", "svg", "unicode"], value="svg"
        )
        render_btn = gr.Button("Render Board")

        board_render_output = gr.HTML(label="Rendered Board")

        render_btn.click(
            fn=render_board,
            inputs=[fen_input_render, render_mode_dropdown],
            outputs=board_render_output,
        )
    with gr.Tab("Suggest Move", visible=False):
        gr.Markdown("### Suggest a Move")
        fen_input_suggest = gr.Textbox(
            label="FEN String (optional)",
            placeholder="Leave blank for initial position",
            value=None,
        )
        suggest_btn = gr.Button("Suggest Move")

        suggested_move_output = gr.Textbox(label="Suggested Move (SAN)")

        suggest_btn.click(
            fn=suggest_move,
            inputs=fen_input_suggest,
            outputs=suggested_move_output,
        )

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