# Instalar pacotes necessários
!pip install google-adk

import google.generativeai as genai
from google.adk.agents import LlmAgent
from google.adk.runners import Runner
from google.adk.sessions import InMemorySessionService
from bcb import sgs
import pandas as pd

import getpass
import os

if not os.environ.get("GOOGLE_API_KEY"):
  os.environ["GOOGLE_API_KEY"] = getpass.getpass("Enter API key for Google Gemini: ")

# Definir o modelo Gemini 2.0 Flash
MODEL = "gemini-2.0-flash"

# Função que consulta o IPCA acumulado em 12 meses
def get_ipca_12m(series_id: int, start: str) -> dict:
    """
    Consulta o IPCA no SGS/BCB e calcula a variação acumulada em 12 meses.

    Parâmetros
    ----------
    series_id : int
        ID da série do SGS (default=433 para IPCA).
    start : str
        Data de início no formato YYYY-MM-DD.

    Retorno
    -------
    dict
        Exemplo:
        {
            "series_id": 433,
            "date": "2025-08-31",
            "ipca_12m": 3.45
        }
    """
    try:
      # Baixar série do SGS
      ipca = sgs.get({ "ipca": series_id }, start=start)

      if ipca.empty:
            raise ValueError("Nenhum dado retornado pela API do SGS.")

      # Calcular inflação acumulada em 12 meses
      dfr = ipca.rolling(12)
      i12 = dfr.apply(lambda x: (1 + x/100).prod() - 1).dropna() * 100

      # Selecionar último valor disponível
      ultimo = i12.dropna().tail(1)

      if ultimo.empty:
        raise ValueError("Não foi possível calcular a inflação em 12 meses.")

      resultado = {
              "series_id": series_id,
              "date": ultimo.index[0].strftime("%Y-%m-%d"),
              "ipca_12m": round(float(ultimo["ipca"].values[0]), 2)
          }

      return resultado

    except Exception as e:
      return { "error": str(e) }

# Definir o agente de indicadores econômicos com descrição e instruções
economic_agent = LlmAgent(
    name="ipca_collector",
    model=MODEL,
    tools=[get_ipca_12m],

    description="""Agente especializado em indicadores econômicos do Brasil,
    com foco na coleta e interpretação da inflação medida pelo IPCA.""",

    instruction="""
                Responda exclusivamente a perguntas relacionadas à inflação do Brasil (IPCA).

                Sempre que o usuário solicitar a taxa de inflação acumulada em 12 meses, utilize a ferramenta `get_ipca_12m` com os seguintes parâmetros:
                - series_id = 433
                - start = data mais recente possível no formato "yyyy-mm-dd"

                Se o usuário não informar um período ou data, faça automaticamente a requisição usando "2023-01-01" como data inicial.
                Nunca pergunte ao usuário qual período utilizar: colete os dados diretamente, independentemente da pergunta.

                Apresente a resposta em linguagem clara e objetiva, incluindo a data de referência do dado.

                Caso o usuário faça perguntas fora do escopo (ex.: PIB, câmbio, juros), responda informando que você só pode fornecer informações sobre inflação com base no IPCA.
"""
)

# Define o nome do aplicativo que usará o ADK
APP_NAME = "my-adk-app"

# Define o identificador do usuário
USER_ID = "user123"

# Define o identificador da sessão
SESSION_ID = "session-abc"

# Cria um serviço de gerenciamento de sessões em memória
session_service = InMemorySessionService()

# Cria uma nova sessão associada ao app, usuário e ID da sessão
session = await session_service.create_session(app_name=APP_NAME, user_id=USER_ID, session_id=SESSION_ID)

# Inicializa o Runner, responsável por executar o agente dentro do contexto da sessão
runner = Runner(agent=economic_agent, app_name=APP_NAME, session_service=session_service)

# Importa biblioteca para lidar com tarefas assíncronas
import asyncio

# Importa tipos do Google GenAI para estruturar mensagens em formato ADK
from google.genai import types  # Para criar Content/Parts da mensagem

# Define uma função assíncrona que envia uma consulta ao agente
async def call_agent_async(query: str):
  """Envia uma consulta ao agente e imprime a resposta final."""
  # Exibe a consulta feita pelo usuário
  print(f"\n>>> User Query: {query}")

  # Prepara a mensagem do usuário no formato aceito pelo ADK
  content = types.Content(role='user', parts=[types.Part(text=query)])

  # Define a resposta padrão caso o agente não retorne nada
  final_response_text = "Agent did not produce a final response."

  # Executa a lógica do agente de forma assíncrona, gerando eventos
  async for event in runner.run_async(user_id=USER_ID, session_id=SESSION_ID, new_message=content):
      # (Opcional) imprime todos os eventos recebidos durante a execução
      # print(f"  [Event] Author: {event.author}, Type: {type(event).__name__}, Final: {event.is_final_response()}, Content: {event.content}")

      # Verifica se o evento é a resposta final do agente
      if event.is_final_response():
          # Caso o agente tenha retornado conteúdo textual
          if event.content and event.content.parts:
             final_response_text = event.content.parts[0].text
          # Caso o agente tenha escalado a execução para erro ou exceção
          elif event.actions and event.actions.escalate:
             final_response_text = f"Agent escalated: {event.error_message or 'No specific message.'}"
          # Interrompe o loop assim que a resposta final for encontrada
          break

  # Exibe a resposta final processada pelo agente
  print(f"<<< Agent Response: {final_response_text}")

# Define uma função assíncrona para rodar a conversa inicial
async def run_conversation():
     # Faz a primeira pergunta ao agente sobre a taxa de inflação no Brasil
     await call_agent_async("Qual a taxa de inflação do brasil?")

     # Faz a segunda pergunta, esperando que o agente simule um erro ou escalonamento
     await call_agent_async("A taxa de inflação acima da meta de 3% significa o que?")  # Espera mensagem de erro da ferramenta

# Executa a função assíncrona de conversa
await run_conversation()