from se_agents.agent import Agent
from config import Config
from se_agents.tools import *
from tools import *
from se_agents.runner import Runner
from pydantic import BaseModel


class Task(BaseModel):
    task_id: str
    question: str
    Level: str
    file_name: str


async def run_agent(query):
    agent = Agent(
        api_key=Config.get_openrouter_api_key(),
        base_url=Config.get_openrouter_api_url(),
        model=Config.get_model(),
        tools=[
            ExaSearch(Config.get_exa_api_key()),
            ExaCrawl(Config.get_exa_api_key()),
            ExaSearchBase(Config.get_exa_api_key()),
            ExaSearchContent(Config.get_exa_api_key()),
            ExaSearchHighlights(Config.get_exa_api_key()),
            OpenAIVisionTool(),
            ThinkTool(),
            FinalOutput(),
            YoutubeVideoInterpreter(),
            TaskFileDownloader(),
            RetriveInfoTaskFile(),
            SpeechToText(),
            CodeInterpreter(),
        ],
        description="You are a general AI assistant. You are direct and concise.",
        rules="""Answer with the following template: [YOUR FINAL ANSWER]. 
    YOUR FINAL ANSWER should be a number OR as few words as possible OR a comma separated list of numbers and/or strings, you do not explain your answer, you just provided it in the least amount of characters as possible.
    For every task that is given to you first think about it using the 'think' tool.
    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, do not write the number as a string instead used the number notation. Example: Not Sixty three, Do 63. 
    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.
    If the task at hand has any file associated with it you never give a direct answer, first use the think tool, to process the request, once the thinking is complete you must first download using the task_file_downloader tool, and the use the proper tool to analyze the file.
    If the task as any audio file associated with it, you must use the speech_to_text tool to convert it to a .txt file.
    If the task involves the interpretation of code, you must use the code_interpreter tool to get the content of the code file.""",
        add_default_rules=False,
        add_think_instructions=True,
        add_final_output_instructions=True,
    )
    runner = Runner(agent, enforce_final=True)
    final_response = ""
    async for response in runner.run(query):
        if response.type == "response":
            # This will typically contain the final output when enforce_final=True,
            # or intermediate text otherwise.
            final_response += response.content
            return final_response
        # Note: 'thinking' is now handled via ThinkTool, resulting in 'tool_call' and 'tool_response' events.
        elif response.type == "tool_call":
            # Includes calls to ThinkTool, FinalOutput, and others.
            print(f"\n\n\033[92m🟡 Tool Call:\n{response.content}\033[0m\n")
        elif response.type == "tool_response":
            print(f"\n\n\033[94m🟢 Tool Response:\n{response.content}\033[0m\n")
        elif response.type == "tool_error":
            print(f"\n\n\033[91m🔴 Tool Error:\n{response.content}\033[0m\n")