from smolagents import (
    CodeAgent, 
    InferenceClientModel, 
    OpenAIServerModel, 
    VisitWebpageTool, 
    WebSearchTool, 
    WikipediaSearchTool, 
    PythonInterpreterTool,
    FinalAnswerTool
)
import os
from vision_tool import image_reasoning_tool

#token=os.getenv("HF_API_TOKEN")

OPENROUTER_API_KEY = os.getenv("OPENROUTER_API_KEY")
if not OPENROUTER_API_KEY:
    raise EnvironmentError("OPENROUTER_API_KEY environment variable not set")

common = dict(
    api_base="https://openrouter.ai/api/v1",
    api_key=OPENROUTER_API_KEY,
    #extra_body={"usage": {"include": True}}
)

#common = dict(
    #api_base = "https://generativelanguage.googleapis.com/v1beta/openai/",
    #api_key = os.getenv("GEMINI_API_KEY"),
#)
class GaiaAgent:
    def __init__(self):
        #import os
        #token = os.getenv("HF_API_TOKEN")
        #print("DEBUG: HF_API_TOKEN is", token[:6]+ "...")
        self.model = OpenAIServerModel(
            #max_tokens = 8096,
            temperature = 0.5,
            #model_id = "gemini-2.0-flash", 
            #model_id = 'qwen/qwen-2.5-coder-32b-instruct:free',
            model_id = "deepseek/deepseek-chat-v3-0324:free",
            #custom_role_conversions=None,      
            #provider="hf-inference",
            #token=token,
            **common
        )

#Creating the agent
        self.agent = CodeAgent(
            model=self.model,
            tools=[
                VisitWebpageTool(),
                WebSearchTool(),
                WikipediaSearchTool(),
                PythonInterpreterTool(),
                FinalAnswerTool(),
                image_reasoning_tool,
            ], 
            max_steps=6,
            verbosity_level = 1,        #or other parameter
            grammar = None,             
            planning_interval=3,
            additional_authorized_imports=[
                "markdownify",
                "json",
                "requests",
                "urllib.request",
                "urllib.parse",
                "wikipedia-api",
                "numpy",
                "math",
                "pytesseract",
                "PIL",
                "chess",
                "pandas",
                "os",
                "bs4",
                "openpyxl",
                "lxml"
            ]
            #name="web_agent",            
            #description="Brows the web to find information",      
            #additional_autorized_imports = ["pandas"],
        )
    def __call__(self, question: str) -> str:
        try:
            return self.agent.run(question)
        except Exception as e:
            return f"ERROR: {e}"