from smolagents import CodeAgent, tool, InferenceClientModel, WebSearchTool, load_tool, PromptTemplates, Tool, FinalAnswerTool
from smolagents import PromptTemplates, PlanningPromptTemplate, FinalAnswerPromptTemplate, ManagedAgentPromptTemplate

from my_tool_reverse_string import ReverseStringTool
from my_tool_image_load import ImageLoadTool
from my_tool_chess_board import ChessBoard
from my_tool_fen import FENTool
from my_tool_chess_analysis import ChessAnalysisTool


task_id = "cca530fc-4052-43b2-b130-b30968d8aa44"

MODEL_REASONING = "Qwen/Qwen2.5-Coder-32B-Instruct"
#MODEL_REASONING = "deepseek-ai/deepseek-coder-33b-instruct"     # not available
#MODEL_REASONING = "deepseek-ai/DeepSeek-Coder-V2-Instruct"      # not available
#MODEL_CHESS     = "jayasuryajsk/chess-reasoner-qwen"
#MODEL_REASONING = "Qwen/Qwen2.5-72B-Instruct"                  # not good
#"meta-llama/Meta-Llama-3-70B-Instruct"
# jayasuryajsk/chess-reasoner-qwen
# https://huggingface.co/jayasuryajsk/chess-reasoner-qwen

PROMPT_TEMPLATES = PromptTemplates(
    system_prompt="""
            You are a general AI assistant. 
            
            Answer the following questions as best you can.
            
            Describe your initial plan as a set of bullet points. 
            Each bullet point should describe in one sentence an action which is to be taken in this step.
            Do NOT provide hypothetical examples for the final answer.
            
            Use the tools provided, and only those which are necessary to answer the question. 
            If you are going to use a tool, describe in detail how you are going
            to use that particular tool and explain parameters used to invoke the tool.
            
            Tools provided :
            * _my_image_load : load an image for given task_id, available arguments: task_id ,
            
            Use <code></code> examples as here
            https://github.com/huggingface/smolagents/blob/main/src/smolagents/vision_web_browser.py
            
            * _my_chess_board : process an image and extract list of chess pieces, available arguments: img ,
            * _my_fen_tool : convert list of chess pieces into FEN notation, available arguments: chest_pieces , 
            * _my_chess_analysis : analyze chess position provided in FEN notation and provide best move, 
                available arguments: fen, player_color
            * _my_reverse_string : reverse a string, available arguments: input_str ,
            
            YOUR FINAL ANSWER should be a number OR as few words as possible OR a comma separated list of
            numbers and/or strings.
            If you are asked for a number, don’t use comma to write your number neither use units such as $ or percent
            sign unless specified otherwise.
            If you are asked for a string, don’t use articles, neither abbreviations (e.g. for cities), and write the digits in
            plain text unless specified otherwise.
            If you are asked for a comma separated list, apply the above rules depending of whether the element to be put
            in the list is a number or a string.
            
            Report your answer with the following template: FINAL ANSWER: [YOUR FINAL ANSWER].
            
        """,

    planning=PlanningPromptTemplate(
        initial_plan="""
                
            """,
        update_plan_pre_messages="""
                
            """,
        update_plan_post_messages="""
                
            """,
    ),
    managed_agent=ManagedAgentPromptTemplate(task="", report=""),
    final_answer=FinalAnswerPromptTemplate(
        pre_messages="",
        post_messages="""
                
            """
    ),
)

#question = f"Load an image for task id {task_id} and describe the chess position shown on the image. "
#question = f"Load an image for task id {task_id} and display it using matplotlib "
#question = (f"Load an image for task id {task_id}, extyract checc pieces position into FEN notation, and provide "
#            f"best move for black ")
question = ("How many studio albums were published by Mercedes Sosa between 2000 and 2009 (included)? You can use the "
       "latest "
 "2022 version of english wikipedia.")

chess_board_model_name = "my_chess_pieces_recognition.pth"
chess_board_model_dir = "/mnt/c/Users/krzsa/IdeaProjects/Agents-Course-Assignment/saved_models"

reasoning_agent = CodeAgent(
    name="CourseAssistant",
    description="General AI Assistant",
    tools=[
        ImageLoadTool(),
        FinalAnswerTool(),
        ReverseStringTool(),
        ChessBoard(chess_board_model_name, chess_board_model_dir),
        FENTool(),
        ChessAnalysisTool(),
        WebSearchTool()
    ],
    model=InferenceClientModel(model_id=MODEL_REASONING),
    planning_interval=3, # This is where you activate planning!,
    prompt_templates=PROMPT_TEMPLATES,
    managed_agents=[],
    additional_authorized_imports=[
        "PIL",
        "chess",
        "matplotlib",
        "matplotlib.pyplot",
        "stockfish",
        "my_chess_board_tool",
        "my_fen_tool",
        "my_image_load",
        "my_reverse_string",
        "my_chess_analysis_tool"
    ],
)

reasoning_agent.run(question)