import os
import time
from smolagents import (
    ToolCallingAgent,
    CodeAgent,
    MCPClient,
    WikipediaSearchTool,
    InferenceClientModel,
)

max_steps_internet_researcher = 10
max_steps_manager = 10
max_retries = 3
cooldown_on_error = 10


qwuen_72b = "Qwen/Qwen2.5-72B-Instruct"
model_235b_thinking = "qwen/qwen3-235b-a22b-thinking-2507"
model_235b_instruct = "qwen/Qwen3-235B-A22B-Instruct-2507"
model_235b = "qwen/qwen3-235b-a22b-fp8"
model_480b = "qwen/qwen3-coder-480b-a35b-instruct"


def get_hf_model():
    hf_token = os.getenv("HF_TOKEN")
    return InferenceClientModel(model_235b_instruct, token=hf_token, timeout=300)


def get_tavily_mcp_client():
    tavily_token = os.getenv("TAVILY_TOKEN")

    # context manager + Streamable HTTP transport:
    return MCPClient(
        {
            "url": f"https://mcp.tavily.com/mcp/?tavilyApiKey={tavily_token}",
            "transport": "streamable-http",
        }
    )


class Agent:
    def __init__(self):

        model = get_hf_model()
        mcpClient = get_tavily_mcp_client()

        tools = mcpClient.get_tools()
        tools.append(
            WikipediaSearchTool(
                user_agent="Research (Elias.Rosendahl.Jensen-CIC@ibm.com)",
                language="en",
                content_type="summary",
                extract_format="WIKI",
            )
        )

        internet_reseacher = ToolCallingAgent(
            tools=tools,
            model=model,
            max_steps=max_steps_internet_researcher,
            name="web_search_agent",
            description="Runs web searches for you.",
        )

        agent = CodeAgent(
            tools=[],
            managed_agents=[internet_reseacher],
            max_steps=max_steps_manager,
            model=model,
        )

        self.agent = agent

        print("BasicAgent initialized.")

    def __call__(self, question: str) -> str:
        prompt = f"You are a general AI assistant. I will ask you a question. Report your thoughts, and finish your answer with ONLY the 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. 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), 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. the question is: {question}"
        print(f"Agent received question: {prompt}...")

        for _ in range(0, max_retries):
            try:
                answer = self.agent.run(prompt)

                if isinstance(answer, str):
                    if "AGENT ERROR" in answer:
                        print("agent error: ", answer)
                        time.sleep(cooldown_on_error)
                        continue
                print(f"Agent returning fixed answer: {answer}")

                return answer
            except Exception as e:
                print("Other error: ", e)
                time.sleep(cooldown_on_error)

        return "ERROR: no more retries"