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


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
    )

    # 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.",
    )

    # Create manager agent
    manager_agent = CodeAgent(
        tools=[],
        managed_agents=[web_agent],
        model=llm_deepseek,
        max_steps=3,
        planning_interval=10,
        additional_authorized_imports=["pandas", "numpy"],
    )
    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")