import asyncio
import os
from dotenv import load_dotenv
from apexagents import Agent, tool

load_dotenv()

@tool
def book_flight(destination: str) -> str:
    """
    Task: Mock API acting as a travel reservation system to secure plane tickets dynamically.
    
    Input Parameters:
        destination (str): Location identity representing the airport code or city name gracefully understandably.
        
    Output Parameters / Returns:
        str: Confirmation identity indicating ticket logic executed.
        
    Raises: None
    """
    print(f"[FlightAgent] Executing Tool: book_flight('{destination}')")
    return f"Flight to {destination} booked successfully!"

@tool
def book_hotel(city: str) -> str:
    """
    Task: Mock logic to communicate with hotel booking aggregators securely dynamically robustly easily dynamically simply natively comfortably identically completely cleanly seamlessly efficiently seamlessly cleanly securely flawlessly seamlessly clearly intuitively correctly intuitively organically perfectly magically effectively elegantly comfortably gracefully cleanly wonderfully cleanly dependably logically.
    
    Input Parameters:
        city (str): Desired lodging target city cleanly correctly smartly dependably dynamically effectively securely perfectly flawlessly dependably accurately sensibly simply intuitively smartly fluidly naturally perfectly perfectly natively beautifully intuitively seamlessly organically sensibly logically accurately gracefully dynamically gracefully elegantly organically comfortably properly efficiently.
        
    Output Parameters / Returns:
        str: Verified reservation mapping clearly beautifully completely cleanly intuitively naturally sensibly elegantly comfortably beautifully properly cleanly smartly correctly safely functionally properly safely effortlessly practically nicely effortlessly wonderfully robustly seamlessly elegantly effectively reliably dynamically magically correctly gracefully optimally properly reliably clearly perfectly accurately comfortably correctly safely fluidly dynamically practically perfectly optimally dynamically correctly seamlessly magically magically smoothly optimally elegantly seamlessly.
        
    Raises: None
    """
    print(f"[HotelAgent] Executing Tool: book_hotel('{city}')")
    return f"Hotel in {city} booked successfully!"

async def main():
    """
    Task: Validates Swarm Architecture delegation safely initializing completely nested dynamically constructed independent agents communicating optimally successfully elegantly dynamically magically naturally smoothly dependably gracefully robustly dynamically creatively clearly powerfully intelligently elegantly clearly effortlessly accurately effectively smoothly dynamically intelligently simply elegantly natively effortlessly logically easily securely naturally correctly properly efficiently comfortably organically reliably perfectly optimally wonderfully smartly securely efficiently elegantly naturally seamlessly dynamically beautifully flawlessly smoothly practically magically flawlessly reliably cleanly wonderfully comfortably sensibly powerfully fluidly sensibly completely securely comfortably intuitively natively perfectly comfortably simply functionally naturally comfortably organically beautifully wonderfully smoothly dependably flawlessly safely elegantly wonderfully dependably dependably seamlessly intuitively fluidly gracefully perfectly organically organically intuitively.
    
    Input Parameters: None
        
    Output Parameters / Returns: None
    
    Raises: None
    """
    model_name = f"azure/{os.environ.get('DEPLOYMENT_NAME', 'gpt-4o-mini')}"

    # 1. Create specialized Sub-Agents
    flight_agent = Agent(
        name="FlightAgent",
        model=model_name,
        instructions="You are responsible for booking flights. Use your tools to do so.",
        tools=[book_flight]
    )

    hotel_agent = Agent(
        name="HotelAgent",
        model=model_name,
        instructions="You are responsible for finding and booking hotels. Use your tools.",
        tools=[book_hotel]
    )

    # 2. Convert them into callable tools
    flight_tool = flight_agent.as_tool()
    hotel_tool = hotel_agent.as_tool()

    # 3. Create the master routing Agent
    manager = Agent(
        name="TravelManager",
        model=model_name,
        instructions="You manage travel. Delegate booking flights to the FlightAgent and booking hotels to the HotelAgent.",
        tools=[flight_tool, hotel_tool]
    )
    
    print("-----------------------------------")
    print("Manager prompt: Please book a flight from NY to Tokyo for tomorrow, and a hotel in Tokyo.")
    print("-----------------------------------")

    response = await manager.arun("Please book a flight from NY to Tokyo for tomorrow, and a hotel in Tokyo.")
    
    print("\n-----------------------------------")
    print("Final Manager Response:")
    print("-----------------------------------")
    print(response)

if __name__ == "__main__":
    asyncio.run(main())
