# Copyright 2026 Google LLC
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     https://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

# @title 1. Environment Setup & Authentication
# Install core dependencies
%pip install google-auth requests httpx nest_asyncio a2a-sdk

import asyncio
import json
import os
import uuid
from google.auth import default
from google.auth.transport.requests import Request
from google.colab import auth
import httpx
import nest_asyncio

from a2a.client import ClientConfig, ClientFactory
from a2a.types import (
    AgentCapabilities,
    AgentCard,
    AgentInterface,
    Message,
    Part,
    Role,
    SendMessageRequest,
)
from a2a.utils.constants import TransportProtocol

nest_asyncio.apply()

# Authenticate the user
auth.authenticate_user()

# Get credentials and project ID
creds, _ = default()
creds.refresh(Request())
access_token = creds.token

ENDPOINT = "https://geminidataanalytics.googleapis.com"
PROJECT_ID = "[your-project-id]"  # @param {type:"string"}
LOCATION = "global"  # @param {type:"string"}
# AGENT_ID can be found from the Cloud URL, e.g.
# https://console.cloud.google.com/bigquery/agents_hub/<your-agent-id>?project=<your-project-id>
AGENT_ID = "[your-agent-id]"  # @param {type:"string"}

if not PROJECT_ID or PROJECT_ID == "[your-project-id]":
  PROJECT_ID = str(os.environ.get("GOOGLE_CLOUD_PROJECT"))
if not LOCATION:
  LOCATION = os.environ.get("GOOGLE_CLOUD_REGION")

TENANT = f"projects/{PROJECT_ID}/locations/{LOCATION}/dataAgents/{AGENT_ID}"
print(f"Target Tenant: {TENANT}")

# @title 2. Initialize A2A SDK Client


# A2A SDK Transport Adapter for Gemini Data Analytics API
class GDATransport(httpx.AsyncBaseTransport):

  def __init__(self, transport):
    self._transport = transport

  async def handle_async_request(self, request):
    # Route /dataAgents/.../message to /dataAgents/.../v1/message
    url_str = str(request.url)
    if "/dataAgents/" in url_str and "/v1/" not in url_str:
      request.url = httpx.URL(
          url_str.replace("/message:send", "/v1/message:send").replace(
              "/message:stream", "/v1/message:stream"
          )
      )

    # Adapt request payload for Gemini Data Analytics backend
    req_body = await request.aread()
    if req_body:
      try:
        data = json.loads(req_body)
        if "message" in data:
          msg = data["message"]
          if "parts" in msg and "content" not in msg:
            msg["content"] = msg.pop("parts")
          if "messageId" in msg and "message_id" not in msg:
            msg["message_id"] = msg.pop("messageId")
        data["configuration"] = {"blocking": True}

        new_req_body = json.dumps(data).encode("utf-8")
        headers = httpx.Headers(request.headers)
        headers["content-length"] = str(len(new_req_body))

        request = httpx.Request(
            method=request.method,
            url=request.url,
            headers=headers,
            content=new_req_body,
        )
      except Exception as e:
        print(f"[Debug] Request adapt error: {e}")

    response = await self._transport.handle_async_request(request)

    # Adapt response payload ("content" -> "parts") for A2A v1.0 SDK proto parser
    if response.status_code == 200:
      try:
        res_body = await response.aread()
        data_str = res_body.decode("utf-8").replace('"content":', '"parts":')
        new_res_body = data_str.encode("utf-8")

        headers = httpx.Headers(response.headers)
        headers.pop("content-encoding", None)
        headers["content-length"] = str(len(new_res_body))

        return httpx.Response(
            status_code=response.status_code,
            headers=headers,
            content=new_res_body,
            request=request,
        )
      except Exception as e:
        print(f"[Debug] Response adapt error: {e}")

    return response


# Initialize A2A SDK Client
httpx_client = httpx.AsyncClient(
    transport=GDATransport(httpx.AsyncHTTPTransport()),
    headers={"Authorization": f"Bearer {access_token}"},
    timeout=120.0,
)

client_config = ClientConfig(
    streaming=True,
    httpx_client=httpx_client,
    supported_protocol_bindings=[
        TransportProtocol.HTTP_JSON,
        TransportProtocol.JSONRPC,
    ],
)

card = AgentCard(
    name="TargetAgent",
    description="Gemini Data Analytics Agent",
    capabilities=AgentCapabilities(),
    supported_interfaces=[
        AgentInterface(
            protocol_binding=TransportProtocol.HTTP_JSON,
            url=f"{ENDPOINT}/v1beta/a2a",
        )
    ],
)

client = ClientFactory(config=client_config).create(card)
print("Client created successfully using A2A SDK!")

# @title 3. Send Message and Stream Results


async def run_query(query_text):
  print(f"Sending query: '{query_text}'...\n")

  msg = Message(
      message_id=f"msg-{uuid.uuid4()}",
      role=Role.ROLE_USER,
      parts=[Part(text=query_text)],
  )
  request = SendMessageRequest(tenant=TENANT, message=msg)

  async for response in client.send_message(request):
    task = getattr(response, "task", None)
    if task and hasattr(task, "id") and task.id:
      print(f"Task ID: {task.id}")
      status = getattr(task, "status", None)
      if status:
        print(f"Task Status: {getattr(status, 'state', status)}\n")

      for art in getattr(task, "artifacts", []):
        desc = getattr(art, "description", "")
        print(f"### [Artifact] {art.name} {f'({desc})' if desc else ''}")
        for part in getattr(art, "parts", []):
          if hasattr(part, "text") and part.text:
            print(part.text)

    msg_resp = getattr(response, "message", None)
    if msg_resp:
      for part in getattr(msg_resp, "parts", []):
        if hasattr(part, "text") and part.text:
          print(f"[Agent]: {part.text}")


QUERY = "What were the top 5 most popular start stations for bike trips?"  # @param {type:"string"}
asyncio.run(run_query(QUERY))

# @title Resource Cleanup
print(
    "No specific cloud resources were created in this demo that require manual"
    " deletion."
)
print(
    "However, you can use this section to reset any local session state if"
    " needed."
)