from smolagents import Tool
from stockfish import Stockfish
import shutil

class ChessPredictorTool(Tool):
    name = "chess_predictor"
    description = "Analyzes a chess state (FEN) and predicts the best move."
    inputs = {
        "fen": {
            "type": "string",
            "description": "FEN (Forsyth-Edwards notation) value",
        }
    }
    output_type = "string"

    def __init__(self):
        super().__init__()

        stockfish_path = shutil.which("stockfish") or "/usr/games/stockfish"

        if not stockfish_path:
            raise RuntimeError(f"Can't find stockfish on PATH or at /usr/games/stockfish.")

        print(f"Stockfish path: {stockfish_path}")

        self.stockfish = Stockfish(
            path=stockfish_path,
            depth=10,
            parameters={"Threads": 2, "Minimum Thinking Time": 30}
        )

    def forward(self, fen: str) -> str:
        if not self.stockfish.is_fen_valid(fen):
            return f"Error: invalid FEN notation: {fen}"

        self.stockfish.set_fen_position(fen)

        print(f"Current chessboard:\n {self.stockfish.get_board_visual()}")

        best_move = self.stockfish.get_best_move()
        return f"{best_move}" if best_move else "No valid move found on chessboard."