from agents import ( Agent , Runner , AsyncOpenAI , 
                   OpenAIChatCompletionsModel , function_tool , RunHooks  , RunContextWrapper , TResponseInputItem , ModelResponse , Tool )
import requests
import os
from dotenv import load_dotenv
import os
import asyncio

load_dotenv()

api_key = os.getenv("GEMINI_API_KEY")

model = 'gpt-4o-mini'

#! Runner Hooks

class CustomRunnerHook(RunHooks):
    def __init__(self , display_name):
        self.counter = 0
        self.display_name = display_name     

    async def on_llm_start(self, context : RunContextWrapper, agent : Agent, system_prompt : str, input_items : list[TResponseInputItem]):
        self.counter += 1
        print("LLM Called")
        print(f"Agent Display Name : {self.display_name}")
        print(f"counter : {self.counter}")
        print(f"Agent Name : {agent.name}")     
        print(f"System prompt to LMM : {system_prompt}")        
        print(f"Input Items : {input_items}")        
    

    async def on_llm_end(self, context : RunContextWrapper, agent : Agent , response : ModelResponse):
        self.counter += 1
        print("LLM Return ")
        print(f"Agent Display Name : {self.display_name}")
        print(f"counter : {self.counter}")
        print(f"Agent Name : {agent.name}")     
        print(f"LLM response : {response.output}")        

    

@function_tool
def weather(city:str):
    url = 'http://api.weatherapi.com/v1/current.json?key=8e3aca2b91dc4342a1162608252604&q'
    params = {'q' : city}
    response = requests.get(url=url , params=params)
    data = response.json()
    return data


weatherDetails = Agent(
     name="weather Details Agent",
     instructions="""
     You will got the weather result of any city , country from Weather Agent analyz the output and based on that tell user do it is save
     to go outside in this weather also provide the weather output you have received from weather Agent
     """,
     model=model,
)

agent = Agent(
    name="Weather Agent",
    instructions=f"""
    You are a helpfull Assistant You can solve user query , call weather tool to get weather details once you got the weather result handoff
    it to the weather Details Agent do not return any output  pass the result to the Weather Details Agent if user explicity mentioned
    to you that only tell me about weather then no hand off """,
    model=model,
    tools=[weather],
    handoffs=[weatherDetails],
)

async def main():
    result = await Runner.run(starting_agent=agent , input="what is 2 + 2" , hooks=CustomRunnerHook(display_name="Runner Hooks tracing all Agents"))
    print("Final output : " , result.final_output)

asyncio.run(main())