from chessimg2pos import predict_fen
from stockfish import Stockfish
from langchain_core.tools import tool
from logging_config import logger  # Import the shared logger

# Adapted from Alberto Formaggio's GAIA agents
# https://github.com/AlbertoFormaggio1/gaia-ai-agents/blob/main/tools/chess_tool.py
# This tool requires the Stockfish chess engine to be installed and accessible:
# not possible in HuggingFace Spaces due to system-level installation requirements.
# You can run it locally or in an environment where you can install Stockfish.

@tool
def chess_tool(task_id: str, color_to_move: str) -> str:
    """
    Given an image of a chessboard, and the color to move, 
    predict the FEN notation and suggest the best move.
    https://en.wikipedia.org/wiki/Forsyth%E2%80%93Edwards_Notation
    Important: the output of this tool is not to be modified in any way.
    Args:
        task_id (str): The identifier for the chessboard image.
        color_to_move (str): 'w' or 'b', the color to move ('white' or 'black').
        
    Returns:
        str: Best move suggestion.
    """
    logger.info(f"Invoking chess tool with task_id: {task_id!r} and color_to_move: {color_to_move!r}")

    if color_to_move not in ['w', 'b']:
        return "Error: color_to_move must be 'w' or 'b'."
    
    predicted_fen = predict_fen(f'./files/{task_id}.png', output_type='simple')

    if color_to_move == 'b':
        # if the color to move is black, we need to flip the board
        rows = reversed(predicted_fen.split('/'))#type: ignore
        board = '/'.join([row[::-1] for row in rows])
    elif color_to_move == 'w':
        board = predicted_fen

    fen = f"{board} {color_to_move} -- 0 1" #type: ignore

    # Initialize Stockfish engine (ensure the path to executable is correct)
    stockfish = Stockfish(path="/home/jmlv/engines/stockfish/stockfish")
    stockfish.set_fen_position(fen)
    best_move = stockfish.get_best_move()
    logger.info(f"Best move determined: {best_move!r}")
    piece = stockfish.get_what_is_on_square(best_move[:2])  # type:ignore
    next_move_fen = piece.value.upper() + best_move[2:] # type:ignore

    return next_move_fen #type: ignore