import chess, chess.svg, math
from autogen import ConversableAgent, register_function
from typing_extensions import Annotated

board = None
board_svgs = None
made_move = None

def get_legal_moves() -> Annotated[str, "A list of legal moves in UCI format"]:
    return "Possible moves are: " + ",".join(
        [str(move) for move in board.legal_moves]
    )

def make_move(move: Annotated[str, "A move in UCI format."]) -> Annotated[str, "Result of the move."]:
    move = chess.Move.from_uci(move)
    board.push_uci(str(move))
    global made_move
    made_move = True

    board_svgs.append(chess.svg.board(
        board,
        arrows=[(move.from_square, move.to_square)],
        fill={move.from_square: "gray"},
        size=250
    ))

    piece = board.piece_at(move.to_square)
    piece_symbol = piece.unicode_symbol()
    piece_name = (
        chess.piece_name(piece.piece_type).capitalize()
        if piece_symbol.isupper()
        else chess.piece_name(piece.piece_type)
    )
    
    return f"Moved {piece_name} ({piece_symbol}) from "\
           f"{chess.SQUARE_NAMES[move.from_square]} to "\
           f"{chess.SQUARE_NAMES[move.to_square]}."

def check_made_move(msg):
    global made_move
    
    if made_move:
        made_move = False
        return True
    else:
        return False

def get_num_turns(num_moves):
    # Each turn includes two moves (one by each player)
    # The first move by player black kicks off the chat
    # The first move by player white starts the game 

    num_turns = math.ceil(num_moves / 2)
    
    if num_moves % 2 == 0:
        num_turns += 1
        
    return num_turns

def initialize():
    global board, board_svgs, made_move
    board = chess.Board()
    board_svgs = []
    made_move = False

def run_multi_agent(llm, task):
    #initialize()
    
    llm_config = {"model": llm}
    
    user_proxy = autogen.ConversableAgent(
        name="Admin",
        system_message="Give the task, and send "
        "instructions to writer to refine the blog post.",
        code_execution_config=False,
        llm_config=llm_config,
        human_input_mode="ALWAYS",
    )

    planner = autogen.ConversableAgent(
        name="Planner",
        system_message="Given a task, please determine "
        "what information is needed to complete the task. "
        "Please note that the information will all be retrieved using"
        " Python code. Please only suggest information that can be "
        "retrieved using Python code. "
        "After each step is done by others, check the progress and "
        "instruct the remaining steps. If a step fails, try to "
        "workaround",
        description="Planner. Given a task, determine what "
        "information is needed to complete the task. "
        "After each step is done by others, check the progress and "
        "instruct the remaining steps",
        llm_config=llm_config,
    )

    engineer = autogen.AssistantAgent(
        name="Engineer",
        llm_config=llm_config,
        description="An engineer that writes code based on the plan "
        "provided by the planner.",
    )

    executor = autogen.ConversableAgent(
        name="Executor",
        system_message="Execute the code written by the "
        "engineer and report the result.",
        human_input_mode="NEVER",
        code_execution_config={
            "last_n_messages": 3,
            "work_dir": "coding",
            "use_docker": False,
        },
    )

    writer = autogen.ConversableAgent(
        name="Writer",
        llm_config=llm_config,
        system_message="Writer."
        "Please write blogs in markdown format (with relevant titles)"
        " and put the content in pseudo ```md``` code block. "
        "You take feedback from the admin and refine your blog.",
        description="Writer."
        "Write blogs based on the code execution results and take "
        "feedback from the admin to refine the blog."
    )

    groupchat = autogen.GroupChat(
        agents=[user_proxy, engineer, writer, executor, planner],
        messages=[],
        max_round=10,
    )

    manager = autogen.GroupChatManager(
        groupchat=groupchat, llm_config=llm_config
    )

    groupchat_result = user_proxy.initiate_chat(
        manager,
        message=task,
    )
    
    return groupchat_result