import os

from smolagents import CodeAgent, DuckDuckGoSearchTool, VisitWebpageTool

from cargo_tool import calculate_cargo_travel_time
from model import model

# TASK for single agent:

# 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° N, 74.0060° W), and return them to me as a pandas dataframe.
# Also give me some supercar factories with the same cargo plane transfer time."""

# def run_single_agent() -> None:
#     agent = CodeAgent(
#         model=model,
#         tools=[DuckDuckGoSearchTool(), VisitWebpageTool(), calculate_cargo_travel_time],
#         additional_authorized_imports=["pandas"],
#         max_steps=20,
#     )
#     print(agent.run(TASK))


# if __name__ == "__main__":
#     run_single_agent()

# TASK for multi-agent:

def build_web_agent() -> CodeAgent:
    return CodeAgent(
        model=model,
        tools=[DuckDuckGoSearchTool(), VisitWebpageTool()],
        name="web_agent",
        description="Browses the web to find places. Always returns a list of: name, latitude, longitude. Does not calculate anything.",
        verbosity_level=0,
        max_steps=10,
    )


MAP_FILE = "saved_map.png"


def check_map_saved(final_answer, memory, agent=None) -> bool:
    # Rejects the final answer until the map exists; the agent sees this message and keeps working
    if not os.path.exists(MAP_FILE):
        raise Exception(f'No map found. Plot the points with px.scatter_map and save with fig.write_image("{MAP_FILE}") before final_answer.')
    # A real world map is tens of KB; a tiny file means a blank or fake placeholder
    if os.path.getsize(MAP_FILE) < 10_000:
        raise Exception(f"{MAP_FILE} is too small to be a real map. Do not write placeholder images; plot the points with px.scatter_map.")
    return True


def build_manager_agent() -> CodeAgent:
    return CodeAgent(
        model=model,
        tools=[calculate_cargo_travel_time],
        managed_agents=[build_web_agent()],
        # "plotly.*" (not "plotly") so the agent may import plotly.express
        additional_authorized_imports=["geopandas.*", "plotly.*", "shapely.*", "json", "pandas", "numpy"],
        planning_interval=5,
        verbosity_level=2,
        max_steps=15,
        final_answer_checks=[check_map_saved],
    )

MAP_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° N, 74.0060° W).
Also give me some supercar factories with the same cargo plane transfer time. You need at least 6 points in total.
Represent this as spatial map of the world, with the locations represented as scatter points with a color that depends on the travel time, and save it to saved_map.png!

Here's an example of how to plot and return a map:
import plotly.express as px
df = px.data.carshare()
fig = px.scatter_map(df, lat="centroid_lat", lon="centroid_lon", text="name", color="peak_hour", size=100,
     color_continuous_scale=px.colors.sequential.Magma, size_max=15, zoom=1)
fig.write_image("saved_map.png")
final_answer(fig)

Never try to process strings using code: when you have a string to read, just print it and you'll see it.

"""

if __name__ == "__main__":
    # Remove the map from an earlier run so the check only passes for a map made in this run
    if os.path.exists(MAP_FILE):
        os.remove(MAP_FILE)
    manager = build_manager_agent()
    manager.run(MAP_TASK)