from smolagents import CodeAgent, ToolCallingAgent, TransformersModel
from tools import available_tools
import torch
import json
import os


def instantiate_agent(executor_type : str="local", agent_type: str ="tool_calling", tools = available_tools) -> CodeAgent | ToolCallingAgent:
    
    # Локальный агент с моделью кодером
    if executor_type == "local" and agent_type == "code":

        code_system_prompt = os.getenv("CODE_AGENT_SYSTEM_PROMPT")

        print(code_system_prompt)

        hf_model = "Qwen/Qwen2.5-Coder-32B-Instruct"
        # hf_model = "Qwen/Qwen2.5-Coder-7B-Instruct"
        # hf_model = "Qwen/Qwen2.5-Coder-3B-Instruct"   # For debug purpose

        llm = TransformersModel(model_id=hf_model,
                                device_map="cuda",
                                torch_dtype=torch.bfloat16,
                                )
        
        agent = CodeAgent(tools=available_tools,
                          model=llm,
                          additional_authorized_imports=['pandas','numpy', 'numpy.*', 'csv', 'markdownify', 'requests'],
                          prompt_templates=({'system_prompt': code_system_prompt}),
                          max_steps=20,
                          )
        return agent
    
    elif executor_type == "local" and agent_type == "tool_calling":

        tool_call_system_prompt = os.getenv("TOOL_CALLING_SYSTEM_PROMPT")

        hf_model = "Qwen/Qwen2.5-7B-Instruct"
        # hf_model = "Qwen/Qwen2.5-3B-Instruct"         # For debug purpose
        # hf_model = "google/gemma-2-2b-it"
        # hf_model = "Qwen/Qwen2.5-7B-Instruct-1M"

        llm = TransformersModel(model_id=hf_model, device_map="cuda", torch_dtype=torch.bfloat16)

        agent = ToolCallingAgent(tools=available_tools, model=llm, max_steps=3, planning_interval=1)
        agent.prompt_templates["system_prompt"] = agent.prompt_templates["system_prompt"] + tool_call_system_prompt



        return agent 
    
    else:
        raise ValueError(f"Unsupported executor type: {executor_type} or agent type: {agent_type}")


if __name__ == "__main__":

    agent = instantiate_agent()

    question = """
Question: Where were the Vietnamese specimens described by Kuznetzov in Nedoshivina's 2010 paper eventually deposited? Just give me the city name without abbreviations.

Tools required:
1. search engine

Approximately, the problem can be solved as follows::
1. Search "Kuznetzov Nedoshivina 2010"
2. Find the 2010 paper "A catalogue of type specimens of the Tortricidae described by V. I. Kuznetzov from Vietnam and deposited in the Zoological Institute, St. Petersburg"
"""
    agent.run(question)