from dotenv import load_dotenv
from google.adk.sessions import InMemorySessionService
from google.adk.agents.callback_context import CallbackContext
from google.adk.agents.invocation_context import InvocationContext
from google.genai import types
from google.adk.agents import LlmAgent, BaseAgent
from google.adk.events import Event
from google.adk.models.lite_llm import LiteLlm
from google.adk.runners import Runner
from typing import Optional, AsyncGenerator

import pprint
import json

load_dotenv(dotenv_path=".env")


APP_NAME 		= "adk-summit"
USER_ID 		= "user1"
SESSION_ID	= "session001"
validator_model_id = "azure/o_series/o3-mini"


session_service = InMemorySessionService()
initial_state = {"idea" : {}}

session = session_service.create_session(
	app_name=APP_NAME,
	user_id=USER_ID,
	session_id=SESSION_ID,
	state=initial_state
)

def evaluate_decision(
	callback_context: CallbackContext) -> Optional[types.Content] :
	"""
	"""
	print(callback_context.state.get("idea"))
	try :
		response = json.loads(callback_context.state.get("idea"))

	except json.JSONDecodeError :
		pass

	except KeyError :
		pass

	else :
		callback_context.state["idea"] = response["type"]

	print("After mod :::")
	print(callback_context.state.get("idea"))
	return None

class GiftAgent(BaseAgent) :
	main_agent: LlmAgent

	def __init__(self, name: str, main_agent: LlmAgent) :
		"""
		"""
		super().__init__(name=name, main_agent=main_agent)

	async def _run_async_impl(
		self, ctx: InvocationContext) -> AsyncGenerator[Event, None] :
		"""
		"""
		async for event in self.main_agent.run_async(ctx) :
			pprint.pp(event.dict(), indent=2)

			yield event

main_agent = LlmAgent(
	name="main_agent",
	model=LiteLlm(model=validator_model_id),
	sub_agents=[],
	description="""You're a helpful assistant.""",
	instruction="""Come up with a nice gift idea. Return a JSON that contains
	the following properties: type, price_estimate, rarity.
	""",
	input_schema=None,
	output_key="idea",
	after_agent_callback=evaluate_decision,
)

gift_agent = GiftAgent(name="gift_agent", main_agent=main_agent)

def call_agent(runner_arg: Runner) :
	"""
	"""
	current_session = session_service.get_session(app_name=APP_NAME, user_id=USER_ID, session_id=SESSION_ID)
	if not current_session :
		print("Session not found...")
		return

	events = runner_arg.run(
		user_id=USER_ID, session_id=SESSION_ID, new_message=types.Content(
			role="user", parts=[types.Part(text="Execute")]))

	final_response = "No final response captured."

	for event in events :
		if event.is_final_response() and event.content and event.content.parts :
			final_response = event.content.parts[0].text

	final_session = session_service.get_session(app_name=APP_NAME, user_id=USER_ID, session_id=SESSION_ID)

if __name__ == "__main__" :

	runner = Runner(
		agent=gift_agent,
		app_name=APP_NAME,
		session_service=session_service
	)
	call_agent(runner)