from smolagents import Tool
import numpy as np
import textwrap


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

# https://chessvision.ai/docs/tools/fen2image/
# https://en.wikipedia.org/wiki/Forsyth%E2%80%93Edwards_Notation


class FENTool(Tool):
    name = "_my_tool_fen"
    description = """
        Convert a list of chess pieces into FEN notation
        To invoke the tool use code as below
        <code>
        loaded_fen = _my_tool_fen(chest_pieces="dummy")
        </code>
    """

    inputs = {
        "chest_pieces": {
            "type": "string",
            "description": "string with chest pieces",
        }
    }

    output_type = "string"

    def __shortenFEN(self, fen):
        """Reduce FEN to shortest form (ex. '111p11Q' becomes '3p2Q')"""
        return fen.replace('11111111','8').replace('1111111','7') \
            .replace('111111','6').replace('11111','5') \
            .replace('1111','4').replace('111','3').replace('11','2')

    def forward(self, chest_pieces: str) -> str:
        arr = textwrap.wrap(chest_pieces, 8)
        #print(f"Arr: \n{arr}")

        arr = arr[::-1]
        #print(f"Arr: \n{arr}")

        arr = [row[::-1] for row in arr]

        fen_long = "/".join(arr)
        #print(f"FEN long: \n{fen_long}")

        fen_short = self.__shortenFEN(fen_long)
        #print(f"FEN short: \n{fen_short}")

        fen_enhanced = f"{fen_short} b - - 0 1"
        print(f"FEN: {fen_enhanced}")

        _out = fen_enhanced
        return _out