import asyncio
from pydantic import BaseModel
from agents import (
    Agent,
    Runner,
    ModelSettings,
    function_tool,
    input_guardrail,
    GuardrailFunctionOutput,
    RunContextWrapper,
    TResponseInputItem,
)

class InputGuardrailCheck(BaseModel):
    is_query_valid: bool
    reasoning: str

guardrail_agent = Agent(
    name="Guardrail check",
    instructions="""You have to Check if the user's query is valid or not. An invalid query is related to harmful,abusive or sexual content""",
    output_type=InputGuardrailCheck,
)

class AgentOutput(BaseModel):
    answer: str
    reasoning: str

@input_guardrail
async def input_guardrail_fn(
    ctx: RunContextWrapper[None], agent: Agent, input: str | list[TResponseInputItem]
) -> GuardrailFunctionOutput:
    result = await Runner.run(guardrail_agent, input, context=ctx.context)
    # print(result.final_output.is_query_valid)
    reason = result.final_output.reasoning
    print(reason)
    return GuardrailFunctionOutput(
        output_info=result.final_output,
        # tripwire_triggered=False #result.final_output.is_math_homework,
        tripwire_triggered=not result.final_output.is_query_valid,
    )

@function_tool
def get_weather(city: str) -> str:
    """
    Get the current weather for a given city.
    """
    # Simulate a weather API call
    return f"The current weather in {city} is sunny with a temperature of 25°C."

@function_tool
def give_support_number() -> str:
    """
    Provide support information of PIAIC.
    """
    return "For support, please contact our helpdesk at +1234567890."

agent_2 = Agent(
    name="Customer_Support_Agent",
    instructions="You are a helpful Customer Support Agent with tools to provide info regarding PIAIC",
    tools=[give_support_number],
    model="gpt-4.1",
    model_settings=ModelSettings(parallel_tool_calls=True),
)
agent = Agent(
    name="Helper agent",
    instructions="""You are a helpful Assistant. Make a plan before answering the user's question. 
    If the user's input contains a general query and a tool call, make sure the answer contains both rather than just the tool call.""",
    tools=[
        get_weather,
        agent_2.as_tool(
            tool_name="Customer_Support_Agent",
            tool_description="This is a Customer Support Agent, Call this for support related queries.",
        ),
    ],
    model="gpt-4.1",
    model_settings=ModelSettings(parallel_tool_calls=True),
    input_guardrails=[input_guardrail_fn],
    output_type=AgentOutput,
)

async def main():
    result = await Runner.run(
        agent,
        input="""
        What is the capital of France?
        What is the weather like in Karachi?
        What is support number of PIAIC?
        What is the stock price of Tesla?
        """,
    )
    print(result.final_output)

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