# 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 Setup and Authentication
import json
import os
import time
import uuid
from google.auth import default
from google.auth.transport.requests import Request
from google.colab import auth
import requests

# 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"
LOCATION = "global"  # @param {type:"string"}
PROJECT_ID = "[your-project-id]"  # @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}"
BASE_URL = f"{ENDPOINT}/v1beta/a2a/{TENANT}/v1"
HEADERS = {
    "Authorization": f"Bearer {access_token}",
    "Content-Type": "application/json",
}
print(f"Target Tenant: {TENANT}")
print(f"Access Token Length: {len(access_token) if access_token else 0}")

url = f"{BASE_URL}/card"

try:
  response = requests.get(url, headers=HEADERS, timeout=30)
  response.raise_for_status()
  print("Agent Card:")
  print(json.dumps(response.json(), indent=2))
except Exception as e:
  print(f"Error fetching agent card: {e}")

# @title 2. Send Message and Extract Outputs
import json
import requests
import uuid

USER_QUERY = "What were the top 5 most popular start stations for bike trips?"  # @param {type:"string"}
SHOW_THINKING = False  # @param {type:"boolean"}


def send_message(query, show_thinking=False):
  url = f"{BASE_URL}/message:send"
  payload = {
      "tenant": TENANT,
      "message": {
          "message_id": f"msg-{uuid.uuid4()}",
          "role": "ROLE_USER",
          "content": [{"text": query}],
      },
      "configuration": {"blocking": True},
  }

  print(f"Sending query: '{query}'...")
  response = requests.post(url, headers=HEADERS, json=payload, timeout=120)
  response.raise_for_status()

  res_json = response.json()
  task = res_json.get("task")

  if not task:
    print("Received direct Message:")
    print(json.dumps(res_json, indent=2))
    return

  print(f"\nTask ID: {task.get('id')}")
  print(f"Status: {task.get('status', {}).get('state')}\n")

  # Optional: Display intermediate agent scratchpad / planning steps
  if show_thinking:
    print("--- Intermediate Thoughts ---")
    for msg in task.get("history", []):
      if msg.get("role") != "ROLE_USER":
        for part in msg.get("content", []):
          if "text" in part:
            print(f"> {part['text']}\n")

  # Display structured artifacts (Final Response, SQL, Tables, Job details)
  artifacts = task.get("artifacts", [])
  print(f"--- Artifacts ({len(artifacts)}) ---\n")
  for art in artifacts:
    name = art.get("name", "Artifact")
    desc = art.get("description", "")
    print(f"### {name} {f'({desc})' if desc else ''}")

    for part in art.get("parts", []):
      sub_type = (
          part.get("metadata", {}).get("gda_message", {}).get("subType", "")
      )
      sub_label = f"[{sub_type}] " if sub_type else ""

      if "text" in part:
        print(f"{sub_label}{part['text']}")
      elif "data" in part:
        print(f"{sub_label}\n{json.dumps(part['data'], indent=2)}")
    print("-" * 40)


try:
  send_message(USER_QUERY, show_thinking=SHOW_THINKING)
except Exception as e:
  print(f"Error during execution: {e}")

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