from smolagents import CodeAgent,  LiteLLMModel
import os
from cloudpathlib import AnyPath
from google.genai import types

from tools import initialize_all_available_tools


class MyAgent:
    def __init__(self, verbosity_level=0):
        model = LiteLLMModel(
            model_id="gemini/gemini-2.0-flash-lite",
            api_key=os.environ.get('GOOGLE_API_KEY')
        )
        tools = initialize_all_available_tools()

        self.agent = CodeAgent(
            tools=tools, 
            model=model, 
            max_steps=10,
            verbosity_level=verbosity_level
            )
        
        with AnyPath("gaia_prompt.txt").open("r", encoding='utf-8') as f:
            self.gaia_prompt = f.read()

    def get_answer(self, question, additional_data: dict | None = None):

        prompt = self.make_prompt(question)
        additional_args = self.make_additional_args(additional_data)

        answer = self.agent.run(prompt, additional_args=additional_args)

        answer = str(answer)

        fa_tmlp = "FINAL ANSWER:"
        if answer.startswith(fa_tmlp):
            answer = answer[len(fa_tmlp):]
        return answer.strip() 
    
    def make_prompt(self, question):

        custom_prompt = """ 
        I will ask you a question.  
        Your 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).
        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.

        Examples:
        Q: What is the capital of France?
        A: Paris

        Q: What is 2+2 equal to?
        A: 4
        """

        return custom_prompt + '\n' + question
    
    def make_additional_args(self, additional_data: dict | None = None):
        if additional_data is None:
            return None
        
        return {"additional_file": types.Part.from_bytes(
            data=additional_data['content'],
            mime_type=additional_data['Content-Type'],
        )}