from smolagents import CodeAgent, GoogleSearchTool, VisitWebpageTool, HfApiModel, ToolCallingAgent

from src.party_planner.tools.travel_time import calculate_cargo_travel_time


def create_agent(
        agent_type: str,
        name: str,
        model: HfApiModel,
        tools: list,
        max_steps: int = 10,
        additional_imports: list[str] = None,
        description: str = "",
        interval: int = 0,
        verbosity: int = 0,
        **kwargs
):
    """
    **kwargs can be: managed_agents, final_answer_checks, ...
    """
    if agent_type == "code_agent":
        return CodeAgent(
            name=name,
            model=model,
            tools=tools,
            additional_authorized_imports=additional_imports,
            max_steps=max_steps,
            description=description,
            planning_interval=interval,
            verbosity_level=verbosity,
            **kwargs
        )
    elif agent_type == "tool_calling_agent":
        return ToolCallingAgent(
            name=name,
            model=model,
            tools=tools,
            additional_authorized_imports=additional_imports,
            max_steps=max_steps,
            description=description,
            planning_interval=interval,
            verbosity_level=verbosity,
            **kwargs
        )
    return None


if __name__ == "__main__":
    import os
    from dotenv import load_dotenv
    from src.model import get_model

    load_dotenv()
    # os.environ["SERPER_API_KEY"] = os.getenv("SERPER_API_KEY")
    os.environ["SERPAPI_API_KEY"] = os.getenv("SERPAPI_API_KEY")

    AgentType = "code_agent"
    Name = "web_agent"
    Model = get_model(
        model_id="Qwen/Qwen2.5-Coder-32B-Instruct",
        provider="hf-inference"  # "together"
    )
    ToolNames = [
        GoogleSearchTool(provider="serpapi"),
        VisitWebpageTool(),
        calculate_cargo_travel_time
    ]
    AdditionalImports = ["pandas"]
    MaxSteps = 3
    Description = "Browses the web to find information"
    Interval = 4

    # Simple agent served as a baseline for the multi-agent system
    Agent = create_agent(
        agent_type=AgentType,
        name=Name,
        model=Model,
        tools=ToolNames,
        additional_imports=AdditionalImports,
        max_steps=MaxSteps,
        description=Description,
        interval=Interval
    )

    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."""
    Prompt = 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.

    {Task}
    """
    result = Agent.run(Prompt)
    print('\n' * 2, result)