from smolagents import (
    ToolCallingAgent,
    CodeAgent,
    DuckDuckGoSearchTool,
    VisitWebpageTool,
    InferenceClientModel,
)
from dotenv import load_dotenv
from tracing import setup_tracing

load_dotenv()

# Initialize tracing when module is imported
trace_provider = None

MANAGER_PROMPT = """You are a helpful assistant tasked with answering questions using a set of tools.
Now, I will ask you a question. Report your thoughts, and finish your answer with the following template:
FINAL ANSWER: [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.
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.
Your answer should only start with "FINAL ANSWER: ", then follows with the answer. """


def initialize_tracing(enabled=True, provider="langfuse"):
    """
    Initialize tracing for the agent module

    Args:
        enabled: Whether tracing should be active
        provider: Which provider to use - "langfuse" or "phoenix"
    """
    global trace_provider
    if trace_provider is None:
        trace_provider = setup_tracing(
            service_name="smolagent", enabled=enabled, provider=provider
        )
    return trace_provider


def get_agent():
    # Ensure tracing is initialized
    initialize_tracing()

    # SmolagentsInstrumentor will automatically trace agent operations

    llm_qwen = InferenceClientModel(
        model_id="Qwen/Qwen2.5-Coder-32B-Instruct", provider="together"
    )
    llm_deepseek = InferenceClientModel(
        "deepseek-ai/DeepSeek-R1",
        provider="together",
        max_tokens=8096,
        # "Qwen/Qwen3-235B-A22B-FP8",
        # provider="together",
        # max_tokens=8096,
    )

    # Create web agent
    web_agent = ToolCallingAgent(
        tools=[DuckDuckGoSearchTool(), VisitWebpageTool()],
        model=llm_qwen,
        max_steps=3,
        name="Web_Agent",
        description="A web agent that can search the web and visit webpages.",
        verbosity_level=1,
    )

    # Create manager agent
    manager_agent = CodeAgent(
        tools=[],
        managed_agents=[web_agent],
        model=llm_deepseek,
        max_steps=5,
        planning_interval=10,
        additional_authorized_imports=["pandas", "numpy"],
        verbosity_level=1,
        description=MANAGER_PROMPT,
    )
    return manager_agent


if __name__ == "__main__":
    # Initialize tracing when run directly
    # Choose one provider: "langfuse" (default) or "phoenix"
    initialize_tracing(enabled=True, provider="phoenix")

    # Get agent with tracing already configured
    agent = get_agent()

    # Run agent - SmolagentsInstrumentor will automatically trace the execution
    print("Running agent with tracing enabled...")
    result = agent.run(
        "What is the latest news about AI? Please search the web and summarize the results."
    )
    print(f"Result: {result}")
    print(
        "If using Phoenix: run 'python -m phoenix.server.main serve' and view at http://localhost:6006"
    )
    print("If using Langfuse: view traces at https://cloud.langfuse.com")