# Libs
import os
import yaml
import importlib

from smolagents import (
    CodeAgent,
    DuckDuckGoSearchTool,
    VisitWebpageTool,
    WikipediaSearchTool,
    OpenAIServerModel,
    SpeechToTextTool,
    FinalAnswerTool,
)

# Local
from tools import GetTaskFileTool, LoadXlsxFileTool, LoadTextFileTool, AnalyzeImageTool

sys_instruction = (
    "You are a general-purpose AI assistant. For each question, first reason through your answer, "
    "then use the 'final_answer' tool to submit ONLY the final output.\n"
    "Your final answer MUST be one of the following:\n"
    "- A number (without commas or units, unless explicitly instructed),\n"
    "- A short phrase,\n"
    "- A comma-separated list of numbers or strings.\n"
    "For strings, avoid articles and abbreviations, and write digits fully unless told otherwise. "
    "Apply these formatting rules consistently to each item in a list."
)

prompts = yaml.safe_load(
    importlib.resources.files("smolagents.prompts").joinpath("code_agent.yaml").read_text()
)
prompts["system_prompt"] = sys_instruction + prompts["system_prompt"]

# req_instruction = (
req_instruction = (
    "You are a highly capable and autonomous agent named {{name}}, designed to solve complex tasks efficiently.\n"
    "A valued client has assigned you the following task:\n"
    "---\n"
    "Task:\n"
    "{{task}}\n"
    "---\n"
    "To complete this task successfully, follow these steps carefully:\n"
    "    1. Comprehend the task and identify the intended goal.\n"
    "    2. Break the task into clear, logical steps.\n"
    "    3. Select and prepare the tools or resources you need.\n"
    "    4. Set up the required environment or context.\n"
    "    5. Execute each step methodically.\n"
    "    6. Monitor outcomes and identify any deviations.\n"
    "    7. Revise your plan if necessary based on feedback.\n"
    "    8. Maintain internal state and track progress.\n"
    "    9. Verify that the goal has been fully achieved.\n"
    "   10. Present the final result clearly and concisely.\n"
    "If you succeed, you will be rewarded with a significant bonus.\n\n"
    "Your final_answer MUST be:\n"
    "- a number (retain its original type; do not include units),\n"
    "- a concise phrase,\n"
    "- or a comma-separated list of numbers or strings (no articles, no abbreviations).\n\n"
    "Only the content passed to the final_answer tool will be preserved. Any other content will be discarded."
)

prompts['managed_agent']['task'] = req_instruction
prompts['managed_agent']['report'] = "{{final_answer}}"


def get_model(
        model_id: str = "gpt-4.1-mini",
        model_temperature: float = 0.7,
    ):
    """
    Create and return an OpenAIServerModel instance with the specified model ID and temperature.
    """
    # Load the model
    if "gpt" in model_id:
        model = OpenAIServerModel(
                model_id=model_id,
                api_key=os.getenv("OPENAI_API_KEY"),
                temperature=model_temperature
        )
    elif "gemini" in model_id:
        model = OpenAIServerModel(
                model_id=model_id,
                api_key=os.getenv("GOOGLEAI_API_KEY"),
                api_base="https://generativelanguage.googleapis.com/v1beta/openai/",
                temperature=model_temperature
        )
    else:
        raise ValueError(f"Unknown model_id: {model_id}. Supported models are gpt and gemini.")
    return model

def get_agent(
        model_id: str = "gpt-4.1-mini",
        model_temperature: float = 0.7,
        agent_max_steps: int = 15,
    ):
    """
    Create and return a CodeAgent instance with the specified model and tools.
    """

    # Defne the agent with the tools
    agent = CodeAgent(
        tools=[
            DuckDuckGoSearchTool(),
            VisitWebpageTool(),
            WikipediaSearchTool(),
            GetTaskFileTool(),
            SpeechToTextTool(),
            LoadXlsxFileTool(),
            LoadTextFileTool(),
            AnalyzeImageTool(model_id=model_id),
            FinalAnswerTool(),
        ],
        model=get_model(
            model_id=model_id,
            model_temperature=model_temperature,
        ),
        prompt_templates=prompts,
        max_steps=agent_max_steps,
        additional_authorized_imports = ["pandas"],
        name="GAIAAgent",
    )
    return agent

def test():
    """test tests something. """
    # Test the agent with a simple task
    return True