import asyncio

from agents import function_tool
from agents.extensions.handoff_prompt import RECOMMENDED_PROMPT_PREFIX
from agents.realtime import RealtimeAgent, realtime_handoff

"""
When running the UI example locally, you can edit this file to change the setup. THe server
will use the agent returned from get_starting_agent() as the starting agent."""

### TOOLS


@function_tool(
    name_override="faq_lookup_tool", description_override="Lookup frequently asked questions."
)
async def faq_lookup_tool(question: str) -> str:
    # Simulate a slow API call
    await asyncio.sleep(3)

    q = question.lower()
    if "wifi" in q or "wi-fi" in q:
        return "We have free wifi on the plane, join Airline-Wifi"
    elif "bag" in q or "baggage" in q:
        return (
            "You are allowed to bring one bag on the plane. "
            "It must be under 50 pounds and 22 inches x 14 inches x 9 inches."
        )
    elif "seats" in q or "plane" in q:
        return (
            "There are 120 seats on the plane. "
            "There are 22 business class seats and 98 economy seats. "
            "Exit rows are rows 4 and 16. "
            "Rows 5-8 are Economy Plus, with extra legroom. "
        )
    return "I'm sorry, I don't know the answer to that question."


@function_tool(needs_approval=True)
async def update_seat(confirmation_number: str, new_seat: str) -> str:
    """
    Update the seat for a given confirmation number.

    Args:
        confirmation_number: The confirmation number for the flight.
        new_seat: The new seat to update to.
    """
    return f"Updated seat to {new_seat} for confirmation number {confirmation_number}"


@function_tool
def get_weather(city: str) -> str:
    """Get the weather in a city."""
    return f"The weather in {city} is sunny."


faq_agent = RealtimeAgent(
    name="FAQ Agent",
    handoff_description="A helpful agent that can answer questions about the airline.",
    instructions=f"""{RECOMMENDED_PROMPT_PREFIX}
    You are an FAQ agent. If you are speaking to a customer, you probably were transferred to from the triage agent.
    Use the following routine to support the customer.
    # Routine
    1. Identify the last question asked by the customer.
    2. Use the faq lookup tool to answer the question. Do not rely on your own knowledge.
    3. If you cannot answer the question, transfer back to the triage agent.""",
    tools=[faq_lookup_tool],
)

seat_booking_agent = RealtimeAgent(
    name="Seat Booking Agent",
    handoff_description="A helpful agent that can update a seat on a flight.",
    instructions=f"""{RECOMMENDED_PROMPT_PREFIX}
    You are a seat booking agent. If you are speaking to a customer, you probably were transferred to from the triage agent.
    Use the following routine to support the customer.
    # Routine
    1. Ask for their confirmation number.
    2. Ask the customer what their desired seat number is.
    3. Use the update seat tool to update the seat on the flight.
    If the customer asks a question that is not related to the routine, transfer back to the triage agent. """,
    tools=[update_seat],
)

triage_agent = RealtimeAgent(
    name="Triage Agent",
    handoff_description="A triage agent that can delegate a customer's request to the appropriate agent.",
    instructions=(
        f"{RECOMMENDED_PROMPT_PREFIX} "
        "You are a helpful triaging agent. You can use your tools to delegate questions to other appropriate agents."
    ),
    tools=[get_weather],
    handoffs=[
        realtime_handoff(faq_agent, tool_name_override="transfer_to_faq_agent"),
        realtime_handoff(seat_booking_agent, tool_name_override="transfer_to_seat_booking_agent"),
    ],
)

faq_agent.handoffs.append(
    realtime_handoff(triage_agent, tool_name_override="transfer_to_triage_agent")
)
seat_booking_agent.handoffs.append(
    realtime_handoff(triage_agent, tool_name_override="transfer_to_triage_agent")
)


def get_starting_agent() -> RealtimeAgent:
    return triage_agent
