from smolagents import Tool
from stockfish import Stockfish
import traceback


#AUTHORIZED_TYPES = [
#    "string",
#    "boolean",
#    "integer",
#    "number",
#    "image",
#    "audio",
#    "array",
#    "object",
#    "any",
#    "null",
#]

# https://pypi.org/project/stockfish/


class ChessAnalysisTool(Tool):
    name = "_my_tool_chess_analysis"
    description = """
        Analyze chess position provided as FEN notation and answer with best move for player color
        To invoke the tool use code as below
        <code>
        best_answer = _my_tool_chess_analysis(fen=loaded_fen, player_color="dummy")
        </code>
    """

    inputs = {
        "fen": {
            "type": "string",
            "description": "board position in FEN notation",
        },
        "player_color": {
            "type": "string",
            "description": "black or white to make next move",
            "nullable": True
        }
    }

    output_type = "string"

    is_initialized = True

    def __init__(self):
        print(f"***KS*** Chess analysis tool initializing ...")
        self.engine = Stockfish()
        print(f"***KS*** Chess analysis tool initialized with engine: {self.engine}")

    def initialize(self):
        print(f"***KS*** Chess analysis tool initialized (2) with engine: {self.engine}")

    def forward(self, fen: str, player_color = "") -> str:
        # TODO use player_color correctly and set it on FEN in proper position
        print(f"***KS*** Chess analysis tool getting best move for player: {player_color} with position: {fen}")
        best_move = ""
        try:
            self.engine.set_fen_position(fen)
            board = self.engine.get_board_visual()
            print(f"{board}")
            top_3_moves = self.engine.get_top_moves(3)
            print(f"Top 3 moves:\n{top_3_moves}")
            best_move = self.engine.get_best_move()
        except Exception as ex:
            print(traceback.format_exc())
            print(f"***KS*** Exception invoking ChessAnalysisTool: {ex}")

        return f"FINAL ANSWER: {best_move}"