import os
from PIL import Image
from smolagents import CodeAgent, DuckDuckGoSearchTool, InferenceClientModel, VisitWebpageTool

from find_batman_mobile_agent import calculate_cargo_travel_time, check_reasoning_and_plot



def define_multi_agent():
    example_model = InferenceClientModel(model_id="Qwen/Qwen2.5-Coder-32B-Instruct", provider="together")

    task = """Find all Batman filming locations in the world, calculate the time to transfer via cargo plane to here (we're in Gotham, 40.7128, 74.0060), and return them to me as a pandas dataframe.
    Also give me some supercar factories with the same cargo plane transfer time.
    """

    example_multi_agent = CodeAgent(
        model=example_model,
        tools=[DuckDuckGoSearchTool(), VisitWebpageTool(), calculate_cargo_travel_time],
        additional_authorized_imports=["pandas"],
        max_steps=20
    )

    result = example_multi_agent.run(task)

    example_multi_agent.planning_interval = 4

    detailed_report = example_multi_agent.run(f"""
        You're an expert analyst. You make comprehensive reports after visiting many websites.
        Don't hesitate to search for many queries at once in a for loop.
        For each data point that you find, visit the source url to confirm numbers.
    """)

    print(detailed_report)

    example_web_model = InferenceClientModel(
        "Qwen/Qwen2.5-Coder-32B-Instruct", provider="together", max_tokens=8096
    )

    example_web_agent = CodeAgent(
        model=example_web_model,
        tools=[
            DuckDuckGoSearchTool(),
            VisitWebpageTool(),
            calculate_cargo_travel_time,
        ],
        name="web_agent",
        description="Browses the web to find information",
        verbosity_level=0,
        max_steps=10
    )
    manager_agent = CodeAgent(
        model=InferenceClientModel("deepseek-ai/DeepSeek-R1", provider="together", max_tokens=8096),
        tools=[calculate_cargo_travel_time],
        managed_agents=[example_web_agent],
        additional_authorized_imports=[
            "geopandas",
            "plotly",
            "shapely",
            "json",
            "pandas",
            "numpy"
        ],
        planning_interval=5,
        verbosity_level=2,
        final_answer_checks=[check_reasoning_and_plot],
        max_steps=15
    )
    return manager_agent