# 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.

%pip install google-cloud-geminidataanalytics

# @title Authenticate
from google.colab import auth

auth.authenticate_user()

# @title Imports
import json
import textwrap
import time

import altair as alt
from google.iam.v1 import iam_policy_pb2, policy_pb2
from google.protobuf import field_mask_pb2
from google.protobuf.json_format import MessageToDict
from IPython.display import HTML, display
import pandas as pd
import proto

# @title Import the Gemini Data Analytics client library


from google.cloud import geminidataanalytics_v1 as geminidataanalytics

data_agent_client = geminidataanalytics.DataAgentServiceClient()
data_chat_client = geminidataanalytics.DataChatServiceClient()

# @title Billing Project and sample Prompt (System Instruction)

# fmt: off
billing_project = "[your-project-id]"  # @param {type:"string"}

location = "global"

# provide critical context for your Conversational Analytics Agent here
system_instruction = "Think like an Analyst"  # @param {type:"string"}
# fmt: on

# @title BigQuery Datasource

# List datasource(s) you'd like to ask questions of.
# This example includes multiple (two) sources.
# Make sure you have permission to query any tables listed here.

# fmt: off
bq_project_id_1 = "bigquery-public-data"  # @param {type:"string"}
bq_dataset_id_1 = "faa"  # @param {type:"string"}
bq_table_id_1 = "us_airports"  # @param {type:"string"}

bq_project_id_2 = "bigquery-public-data"  # @param {type:"string"}
bq_dataset_id_2 = "san_francisco"  # @param {type:"string"}
bq_table_id_2 = "street_trees"  # @param {type:"string"}
# fmt: on


# Setup
bigquery_table_reference_1 = geminidataanalytics.BigQueryTableReference(
    project_id=bq_project_id_1,
    dataset_id=bq_dataset_id_1,
    table_id=bq_table_id_1,
)
bigquery_table_reference_2 = geminidataanalytics.BigQueryTableReference(
    project_id=bq_project_id_2,
    dataset_id=bq_dataset_id_2,
    table_id=bq_table_id_2,
)

# optional example queries (only leveraged for BigQuery datasources currently)
example_queries = [
    geminidataanalytics.ExampleQuery(
        natural_language_question="How many airports are there?",
        sql_query="SELECT COUNT(*) FROM `bigquery-public-data.faa.us_airports`",
    ),
    geminidataanalytics.ExampleQuery(
        natural_language_question=(
            "How many airports are in {state} with elevation greater than"
            " {elevation}?"
        ),
        sql_query=(
            "SELECT COUNT(*) FROM `bigquery-public-data.faa.us_airports` WHERE"
            " LOWER(state_abbreviation) = @state AND elevation > @elevation"
        ),
        parameters=[
            geminidataanalytics.QueryParameter(
                name="state",
                description="The state abbreviation in lower case.",
                data_type="STRING",
            ),
            geminidataanalytics.QueryParameter(
                name="elevation",
                description="The elevation as a float value.",
                data_type="FLOAT64",
            ),
        ],
    ),
]

use_example_queries = True  # @param {type:"boolean"}

glossary_terms = [
    geminidataanalytics.GlossaryTerm(
        display_name="Airport Code",
        description=(
            "A unique identifier for an airport, represented by a 3 character"
            " long identifier (the IATA code). E.g. 'JFK' for New York."
        ),
        labels=["code", "IATA", "identifier"],
    )
]

use_glossary_terms = True  # @param {type: "boolean"}

bq = geminidataanalytics.BigQueryTableReferences(
    table_references=[
        bigquery_table_reference_1,
        bigquery_table_reference_2,
    ]
)

# @title Looker Datasource

# List up to 5 Looker datasources you'd like to ask questions of.
# This example includes two datasources.

# Looker datasource 1
lookml_model_1 = "<add lookml_model here>"  # @param {type:"string"}
explore_1 = "<add explore here>"  # @param {type:"string"}

# Looker datasource 2 (optional)
# lookml_model_2 = "<add lookml_model here>"  # @param {type:"string"}
# explore_2 = "<add explore here>"  # @param {type:"string"}


# User Credentials (apply to all Looker datasources)
looker_client_id = "<add client_id here>"  # @param {type:"string"}
looker_client_secret = "<add client_secret here>"  # @param {type:"string"}
# OR
# looker_access_token = "<add access_token here>"  # @param {type:"string"}


""" When connecting to the Looker datasource, you can authenticate using either:
1. client_id and client_secret (supported with only PUBLIC instance)
2. access_token (supported with both PUBLIC and PRIVATE instance) """

credentials = geminidataanalytics.Credentials(
    oauth=geminidataanalytics.OAuthCredentials(
        secret=geminidataanalytics.OAuthCredentials.SecretBased(
            client_id=looker_client_id, client_secret=looker_client_secret
        )
    )
)
# Uncomment this and comment out the above one if you are using access_token for authentication
# credentials = geminidataanalytics.Credentials(
#     oauth=geminidataanalytics.OAuthCredentials(
#         token=geminidataanalytics.OAuthCredentials.TokenBased(
#             access_token = looker_access_token
#         )
#     )
# )


""" Looker datasource setup for PUBLIC Instance """

# Required only for Looker PUBLIC instance
looker_instance_uri_1 = "https://my_company.looker.com"  # @param {type:"string"}
# (Optional)
# looker_instance_uri_2 = "https://my_company.looker.com"  # @param {type:"string"}

looker_explore_reference_1 = geminidataanalytics.LookerExploreReference(
    looker_instance_uri=looker_instance_uri_1,
    lookml_model=lookml_model_1,
    explore=explore_1,
)
# looker_explore_reference_2 = geminidataanalytics.LookerExploreReference(
#     looker_instance_uri=looker_instance_uri_2, lookml_model=lookml_model_2, explore=explore_2
# )


""" Looker datasource setup for PRIVATE Instance """

# Required only for Looker PRIVATE instance
# looker_instance_id_1 = "<add looker_instance_id here>"  # @param {type:"string"}
# looker_service_directory_name_1 = "<add looker_service_directory_name here>"  # @param {type:"string"}
# (Optional)
# looker_instance_id_2 = "<add looker_instance_id here>"  # @param {type:"string"}
# looker_service_directory_name_2 = "<add looker_service_directory_name here>"  # @param {type:"string"}

# looker_explore_reference_1 = geminidataanalytics.LookerExploreReference(
#     private_looker_instance_info=geminidataanalytics.PrivateLookerInstanceInfo(
#         looker_instance_id=looker_instance_id_1,
#         service_directory_name=looker_service_directory_name_1,
#     ),
#     lookml_model=lookml_model_1,
#     explore=explore_1
# )
# (Optional)
# looker_explore_reference_2 = geminidataanalytics.LookerExploreReference(
#     private_looker_instance_info=geminidataanalytics.PrivateLookerInstanceInfo(
#         looker_instance_id=looker_instance_id_2,
#         service_directory_name=looker_service_directory_name_2,
#     ),
#     lookml_model=lookml_model_2,
#     explore=explore_2
# )


looker = geminidataanalytics.LookerExploreReferences(
    explore_references=[
        looker_explore_reference_1,
        # looker_explore_reference_2
    ]
)

looker_golden_queries = [
    geminidataanalytics.LookerGoldenQuery(
        natural_language_questions=[
            "What is the highest observed positive longitude?"
        ],
        looker_query=geminidataanalytics.LookerQuery(
            model="airports",
            explore="airports",
            fields=["airports.longitude"],
            filters=[
                geminidataanalytics.LookerQuery.Filter(
                    field="airports.longitude", value=">0"
                )
            ],
            sorts=["airports.longitude desc"],
            limit="1",
        ),
    ),
    geminidataanalytics.LookerGoldenQuery(
        natural_language_questions=[
            "What are the major airport codes and cities in CA?",
            "Can you list the cities and airport codes of airports in CA?",
        ],
        looker_query=geminidataanalytics.LookerQuery(
            model="airports",
            explore="airports",
            fields=["airports.city", "airports.code"],
            filters=[
                geminidataanalytics.LookerQuery.Filter(
                    field="airports.major", value="Y"
                ),
                geminidataanalytics.LookerQuery.Filter(
                    field="airports.state", value="CA"
                ),
            ],
        ),
    ),
]

use_looker_golden_queries = False  # @param {type:"boolean"}

# @title Looker Studio Datasource

# fmt: off
studio_datasource_id = "<add studio_datasource_id here>"  # @param {type:"string"}
# fmt: on

# Setup
studio_references = geminidataanalytics.StudioDatasourceReference(
    datasource_id=studio_datasource_id
)

studio = geminidataanalytics.StudioDatasourceReferences(
    studio_references=[studio_references]
)

# @title Select Data Source

selected_datasource = "bigquery_data_sources"  # @param ["bigquery_data_sources", "looker_data_sources", "looker_studio_data_sources"]

bq_selected = False
looker_selected = False
looker_studio_selected = False
datasource_references = geminidataanalytics.DatasourceReferences()

if selected_datasource == "looker_data_sources":
  looker_selected = True
  datasource_references.looker = looker
elif selected_datasource == "bigquery_data_sources":
  bq_selected = True
  datasource_references.bq = bq
elif selected_datasource == "looker_studio_data_sources":
  looker_studio_selected = True
  datasource_references.studio = studio


# Context set-up
published_context = geminidataanalytics.Context(
    system_instruction=system_instruction,
    datasource_references=datasource_references,
    options=geminidataanalytics.ConversationOptions(
        analysis=geminidataanalytics.AnalysisOptions(
            python=geminidataanalytics.AnalysisOptions.Python(
                enabled=False
            )  # if wanting to use advanced analysis with python. Default is False.
        )
    ),
)

# Context set-up for chat using inline context
inline_context = geminidataanalytics.Context(
    system_instruction=system_instruction,
    datasource_references=datasource_references,
    options=geminidataanalytics.ConversationOptions(
        analysis=geminidataanalytics.AnalysisOptions(
            python=geminidataanalytics.AnalysisOptions.Python(
                enabled=False
            )  # if wanting to use advanced analysis with python. Default is False.
        )
    ),
)

# @title Parser for List methods

def display_resource_list(resource_list):
    for i, resource in enumerate(resource_list):
      # Convert the proto message to a dictionary for field access and JSON serialization
      # 'resource._pb' accesses the underlying protobuf message required by MessageToDict
      resource_dict = MessageToDict(resource._pb)

      resource_name = resource_dict.get("name").split("/")[-1]
      resource_display_name = resource_dict.get("displayName")
      formatted_json = json.dumps(resource_dict, indent=2)
      div_id = f"resource_details_{i}"

      resource_heading = f"""{resource_name} ( {resource_display_name} )""" if resource_display_name else resource_name

      display(HTML(f'''
      <div style="margin: 4px 0;">
          <span onclick="
              var x = document.getElementById('{div_id}');
              var t = this.innerText;
              if (x.style.display === 'none') {{
                  x.style.display = 'block';
                  x.style.maxHeight = '500px';
                  x.style.overflowY = 'auto';
              }} else {{
                  x.style.display = 'none';
              }}
          " style="color: #1a73e8; cursor: pointer; margin-left: 8px; user-select: none;">{resource_heading}</span>
          <pre id="{div_id}" style="display: none; background: #f8f9fa; border: 1px solid #e0e0e0; padding: 12px; margin: 8px 0; border-radius: 4px; overflow-x: auto; font-size: 12px; color: #3c4043;">{formatted_json}</pre>
      </div>
      '''))

# @title Define utilities for viewing responses


def handle_text_response(resp):
  parts = resp.parts
  full_text = "".join(parts)
  if "\n" not in full_text and len(full_text) > 80:
    wrapped_text = textwrap.fill(full_text, width=80)
    print(wrapped_text)
  else:
    print(full_text)


def display_schema(data):
  fields = data.fields
  df = pd.DataFrame({
      "Column": map(lambda field: field.name, fields),
      "Type": map(lambda field: field.type, fields),
      "Description": map(
          lambda field: getattr(field, "description", "-"), fields
      ),
      "Mode": map(lambda field: field.mode, fields),
  })
  display(df)


def display_section_title(text):
  display(HTML(f"<h2>{text}</h2>"))


def format_looker_table_ref(table_ref):
  return f"lookmlModel: {table_ref.lookml_model}, explore: {table_ref.explore}"


def format_bq_table_ref(table_ref):
  return f"{table_ref.project_id}.{table_ref.dataset_id}.{table_ref.table_id}"


def display_datasource(datasource):
  source_name = ""
  if "studio_datasource_id" in datasource:
    source_name = datasource.studio_datasource_id
  elif "looker_explore_reference" in datasource:
    source_name = format_looker_table_ref(datasource.looker_explore_reference)
  else:
    source_name = format_bq_table_ref(datasource.bigquery_table_reference)

  print(source_name)
  if "schema" in datasource:
    display_schema(datasource.schema)


def handle_schema_response(resp):
  if "query" in resp and "question" in resp.query:
    print(resp.query.question)
  elif "result" in resp:
    display_section_title("Schema resolved")
    print("Data sources:")
    for datasource in resp.result.datasources:
      display_datasource(datasource)


def handle_data_response(resp):
  if "query" in resp:
    query = resp.query
    display_section_title("Retrieval query")
    if "name" in query:
      print(f"Query name: {query.name}")
    if "question" in query:
      print(f"Question: {query.question}")
    if "datasources" in query:
      print("Data sources:")
      for datasource in query.datasources:
        display_datasource(datasource)
  elif "generated_sql" in resp:
    display_section_title("SQL generated")
    print(resp.generated_sql)
  elif "result" in resp:
    display_section_title("Data retrieved")

    if "schema" in resp.result:
      fields = [field.name for field in resp.result.schema.fields]
      d = {}
      for el in resp.result.data:
        for field in fields:
          if field in d:
            d[field].append(el[field])
          else:
            d[field] = [el[field]]

      display(pd.DataFrame(d))


def handle_chart_response(resp):
  def _value_to_dict(v):
    if isinstance(v, proto.marshal.collections.maps.MapComposite):
      return _map_to_dict(v)
    if isinstance(v, proto.marshal.collections.RepeatedComposite):
      return [_value_to_dict(el) for el in v]
    if isinstance(v, (int, float, str, bool)):
      return v
    return MessageToDict(v)

  def _map_to_dict(d):
    out = {}
    for k in d:
      if isinstance(d[k], proto.marshal.collections.maps.MapComposite):
        out[k] = _map_to_dict(d[k])
      else:
        out[k] = _value_to_dict(d[k])
    return out

  if "query" in resp and "instructions" in resp.query:
    print(resp.query.instructions)
  elif "result" in resp:
    vegaConfig = resp.result.vega_config
    vegaConfig_dict = _map_to_dict(vegaConfig)
    alt.Chart.from_json(json.dumps(vegaConfig_dict)).display()


def show_message(msg):
  m = msg.system_message
  if "text" in m:
    handle_text_response(m.text)
  elif "schema" in m:
    handle_schema_response(m.schema)
  elif "data" in m:
    handle_data_response(m.data)
  elif "chart" in m:
    handle_chart_response(m.chart)
  print("\n")

###############################################################################
#######         Define prompt / system instruction for the question.             #######
#######    It's optional and can guide the agent's response strategy.   #######
###############################################################################

###############################################################################
##### Example template for system instruction for a sales analyst agent:  #####
###############################################################################

# """
# - system_instruction: >-
#     You are an expert sales analyst and understand how to answer questions about
#     the sales data for a fictitious e-commerce store
# - tables:
#     - table:
#         - name: bigquery-public-data.thelook_ecommerce.orders
#         - description: orders for The Look fictitious e-commerce store.
#         - synonyms: sales
#         - tags: 'sale, order, sales_order'
#         - fields:
#             - field:
#                 - name: order_id
#                 - description: unique identifier for each order
#             - field:
#                 - name: user_id
#                 - description: unique identifier for each user
#             - field:
#                 - name: status
#                 - description: status of the order
#                 - sample_values:
#                     - complete
#                     - shipped
#                     - returned
#             - field:
#                 - name: gender
#                 - description: gender of the user
#                 - sample_values:
#                     - male
#                     - female
#             - field:
#                 - name: created_at
#                 - description: date and time when the order was created in timestamp format
#             - field:
#                 - name: returned_at
#                 - description: >-
#                     date and time when the order was returned in timestamp
#                     format
#             - field:
#                 - name: num_of_item
#                 - description: number of items in the order
#                 - aggregations: 'sum, avg'
#             - field:
#                 - name: earnings
#                 - description: total sales from the order
#                 - aggregations: 'sum, avg'
#             - field:
#                 - name: cost
#                 - description: total cost to get the items for the order
#                 - aggregations: 'sum, avg'
#         - measures:
#             - measure:
#                 - name: profit
#                 - description: raw profit
#                 - exp: cost - earnings
#                 - synonyms: gains
#         - golden_queries:
#             - golden_query:
#                 - natural_language_query: How many orders are there?
#                 - sql_query: SELECT COUNT(*) FROM sqlgen-testing.thelook_ecommerce.orders
#             - golden_query:
#                 - natural_language_query: How many orders were shipped?
#                 - sql_query: >-
#                     SELECT COUNT(*) FROM sqlgen-testing.thelook_ecommerce.orders
#                     WHERE status = 'shipped'
#         - golden_action_plans:
#             - golden_action_plan:
#                 - natural_language_query: Show me the number of orders broken down by gender.
#                 - action_plan:
#                     - step: >-
#                         Run a SQL query on the table
#                         sqlgen-testing.thelook_ecommerce.orders to get a
#                         breakdown of order count by gender.
#                     - step: >-
#                         Create a vertical bar plot using the retrieved data,
#                         with one bar per gender.
#     - table:
#         - name: sqlgen-testing.thelook_ecommerce.users
#         - description: user of The Look fictitious e-commerce store.
#         - synonyms: customers
#         - tags: 'user, customer, buyer'
#         - fields:
#             - field:
#                 - name: id
#                 - description: unique identifier for each user
#             - field:
#                 - name: first_name
#                 - description: first name of the user
#                 - tag: person
#                 - sample_values: 'graham, adam, brian'
#             - field:
#                 - name: last_name
#                 - description: first name of the user
#                 - tag: person
#                 - sample_values: 'warmer, stilles, smith'
#             - field:
#                 - name: gender
#                 - description: gender of the user
#                 - sample_values:
#                     - male
#                     - female
#             - field:
#                 - name: email
#                 - description: email of the user
#                 - tag: contact
#                 - sample_values: 'brian@gmail.com, graham@gmail.com'
#         - golden_queries:
#             - golden_query:
#                 - natural_language_query: How many unique customers are there?
#                 - sql_query: >-
#                     SELECT COUNT(DISTINCT id) FROM
#                     sqlgen-testing.thelook_ecommerce.users
#             - golden_query:
#                 - natural_language_query: How many female users have gmail email id?
#                 - sql_query: >-
#                     SELECT COUNT(DISTINCT id) FROM
#                     sqlgen-testing.thelook_ecommerce.users WHERE users.gender =
#                     'F' AND users.email LIKE '%@gmail.com';
#     - relationships:
#         - relationship:
#             - name: earnings_to_user
#             - description: >-
#                 Sales table is related to the users table and can be joined for
#                 aggregated view.
#             - relationship_type: many-to-one
#             - join_type: left
#             - left_table: str sqlgen-testing.thelook_ecommerce.orders
#             - right_table: sqlgen-testing.thelook_ecommerce.users
#             - relationship_columns: '// Join columns - left_column:''user_id'' - right_column:''id'''
# - glossaries:
#     - glossary:
#         - term: male
#         - description: male gender
#     - glossary:
#         - term: female
#         - description: female gender
#     - glossary:
#         - term: complete
#         - description: complete status
#         - synonyms: 'finish, done, fulfilled'
#     - glossary:
#         - term: shipped
#         - description: shipped status
#     - glossary:
#         - term: returned
#         - description: returned status
#     - glossary:
#         - term: OMPF
#         - description: Order Management and Product Fulfillment
# - additional_descriptions:
#     - text: All the sales data is for Looker organization.
#     - text: 'Orders can be of three categories, food , clothes, electronics.'
# """

###############################################################################
######    Use the schema below to provide your own system instruction.   ######
###############################################################################

# """
# - system_instruction: str # Expected behavior of the agent. Eg. You are a sales analyst.
# - tables:
#     - table: # Table relevant for this agent.
#         - name: str # Name of the table.
#         - description: str # Description of the table.
#         - synonyms: list[str] # List of synonyms used to refer to the table.
#         - tags: list[str] # List of tags associated with the table.
#         - fields: # Fields in the table.
#             - field:
#             - name: str # Name of the column.
#             - description: str # Description of the column.
#             - synonyms: list[str] # List of synonyms used to refer to the column.
#             - tags: list[str] # List of tags associated with the column.
#             - sample_values: list[str] # List of sample values in the column.
#             - aggregations: list[str] # Any commonly used or default aggregations associated with the column.
#         - measures: # Measures for the table.
#             - measure:
#                 - name: str # Name of the measure.
#                 - description: str # Description of the measure.
#                 - exp: str # Expression to construct the measure.
#                 - synonyms: list[str] # List of synonyms used to refer to the measure.
#         - golden_queries: # Golden or popular queries for the table.
#             - golden_query:
#                 - natural_language_query: str # Natural language query.
#                 - sql_query: str # SQL query.
#         - golden_action_plans: # Golden action plans as the suggested steps to take (in order) to answer the query.
#           - golden_action_plan:
#             - natural_language_query: str # Natural language query.
#             - action_plan:
#               - step: str # Step to take.
#     - relationships: # Join relationships between tables.
#         - relationship:
#           - name: str # Name of the relationship.
#           - description: str # Description of the relationship.
#           - relationship_type: str # Relationship type: one-to-one, one-to-many, many-to-one, many-to-many.
#           - join_type: str # Join type: inner, outer, left, right, full
#           - left_table: str # Left table name.
#           - right_table: str # Right table name.
#           - relationship_columns: # Join columns.
#               - left_column: str # Join column from left table.
#               - right_column: str # Join column from right table.
# - glossaries: # Business glossary, jargon, etc.
#     - glossary:
#         - term: str # Name of the term. Term can be a word, phrase, abbreviation, etc.
#         - description: str # Description or definition of the term.
#         - synonyms: list[str] # List of synonyms for the term.
# - additional_descriptions:
#     - text: str # Any additional description that was not covered above.
# """

##################### OPTION 1 ###############################
#######                                                #######
#######        Fill the above YAML template for        #######
#######  system_instruction as shown in example above. #######
#######                                                #######
##############################################################

system_instruction = """"""

##################### OPTION 2 ###############################
#######                                                #######
#######         Leave system_instruction empty         #######
#######         or provide description in form         #######
#######                 of long string.                #######
#######                                                #######
##############################################################

# Define context for question.
# system_instruction is optional, can be used to steer the agent's answering
# strategy.
# fmt: off
system_instruction = "Think like an Analyst"  # @param {type:"string"}
# fmt: on

# @title Helper functions for processing the system_instruction yaml string

from collections import Counter
import re

import yaml
from yaml.scanner import ScannerError


def cleanup_path(input_string: str) -> str:
  """Removes bracketed numbers and preceding dots from an input string.

  Args:
    input_string: The input string.

  Returns:
    The string with bracketed numbers and preceding dots removed.
  """
  cleaned_string = re.sub(r"\[\d+]", "", input_string)
  if cleaned_string.startswith("."):
    cleaned_string = cleaned_string[1:]
  return cleaned_string


def is_valid_yaml(yaml_string: str) -> bool:
  try:
    yaml.safe_load(yaml_string)
    return True
  except ScannerError:
    return False
  except Exception as e:
    print(f"An unexpected error occurred: {e}")
    return False


def print_counter_desc(data: list[str]) -> None:
  # Count the occurrences of each item in the list.
  item_counts = Counter(data)

  # Sort the items in descending order based on their counts.
  sorted_items = sorted(
      item_counts.items(), key=lambda item: item[1], reverse=True
  )

  # Print the items and their counts.
  for item, count in sorted_items:
    print(f"{item}: {count}")


def extract_insights_from_system_instruction(yaml_string: str) -> None:
  """Extracts insights into used fields from a YAML string.

  Args:
    yaml_string: The YAML string to analyze.
  """
  if not is_valid_yaml(yaml_string):
    print("Received Invalid or not a YAML format. Cannot extract insights.")
    return

  try:
    data = yaml.safe_load(yaml_string)
    used_fields = []

    def traverse(node, path=""):
      if isinstance(node, dict):
        for key, value in node.items():
          new_path = f"{path}.{key}" if path else key
          traverse(value, new_path)
      elif isinstance(node, list):
        for i, item in enumerate(node):
          new_path = f"{path}[{i}]"
          traverse(item, new_path)
      else:
        used_fields.append(cleanup_path(path))

    traverse(data)
    print("Insights:")
    print("Number of unique fields used: ", len(set(used_fields)))
    print("Used fields in YAML:")
    print()
    print_counter_desc(used_fields)

  except ScannerError as e:
    print(f"Invalid YAML: {e}")
    return
  except Exception as e:
    print(f"An unexpected error occurred: {e}")
    return


# Printing insights assuming 'system_instruction' is defined above.
extract_insights_from_system_instruction(system_instruction)

# @title Create Data Agent

# fmt: off
data_agent_id = "data_agent_1"  # @param {type:"string"}
# fmt: on

data_agent = geminidataanalytics.DataAgent(
    data_analytics_agent=geminidataanalytics.DataAnalyticsAgent(
        published_context=published_context
    ),
)

request = geminidataanalytics.CreateDataAgentRequest(
    parent=f"projects/{billing_project}/locations/{location}",
    data_agent_id=data_agent_id,  # Optional
    data_agent=data_agent,
)

if bq_selected:
  if use_example_queries:
    request.data_agent.data_analytics_agent.published_context.example_queries = (
        example_queries
    )
  if use_glossary_terms:
    request.data_agent.data_analytics_agent.published_context.glossary_terms = (
        glossary_terms
    )
elif looker_selected and use_looker_golden_queries:
  request.data_agent.data_analytics_agent.published_context.looker_golden_queries = (
      looker_golden_queries
  )

try:
  operation = data_agent_client.create_data_agent(request=request)

  while not operation.done():
    print("Still in progress...")
    time.sleep(2)

  if operation.done():
    if operation.exception():
      print("Operation failed:", operation.exception())
    else:
      print("Operation succeeded:", operation.result())
except Exception as e:
  print(f"Error creating Data Agent: {e}")

# @title Create Conversation

# Initialize request argument(s)
# fmt: off
data_agent_id = "data_agent_1"  # @param {type:"string"}
conversation_id = "conversation_1"  # @param {type:"string"}
# fmt: on

conversation = geminidataanalytics.Conversation(
    agents=[
        data_chat_client.data_agent_path(
            billing_project, location, data_agent_id
        )
    ],
)

request = geminidataanalytics.CreateConversationRequest(
    parent=f"projects/{billing_project}/locations/{location}",
    conversation_id=conversation_id,
    conversation=conversation,
)

# Make the request
response = data_chat_client.create_conversation(request=request)

# Handle the response
print(response)

# Create a request containing a single user message -- your question.
# fmt: off
data_agent_id = "data_agent_1"  # @param {type:"string"}
conversation_id = "conversation_1"  # @param {type:"string"}
question = "Make a bar graph for the top 5 states by the total number of airports"  # @param {type:"string"}
# fmt: on

# Create the user message
messages = [
    geminidataanalytics.Message(
        user_message=geminidataanalytics.UserMessage(text=question)
    )
]

# Create a conversation_reference
conversation_reference = geminidataanalytics.ConversationReference(
    conversation=data_chat_client.conversation_path(
        billing_project, location, conversation_id
    ),
    data_agent_context=geminidataanalytics.DataAgentContext(
        data_agent=data_chat_client.data_agent_path(
            billing_project, location, data_agent_id
        ),
    ),
)

# Form the request
request = geminidataanalytics.ChatRequest(
    parent=f"projects/{billing_project}/locations/{location}",
    messages=messages,
    conversation_reference=conversation_reference,
)

if looker_selected:
  request.credentials = credentials

# Make the request
stream = data_chat_client.chat(request=request)

# Handle the response
for response in stream:
  show_message(response)

# Create a request containing a single user message -- your question.
# fmt: off
data_agent_id = "data_agent_1"  # @param {type:"string"}
question = "Make a bar graph for the top 5 states by the total number of airports"  # @param {type:"string"}
# fmt: on

# Create the message
messages = [
    geminidataanalytics.Message(
        user_message=geminidataanalytics.UserMessage(text=question)
    )
]

# Create the data agent context
data_agent_context = geminidataanalytics.DataAgentContext(
    data_agent=data_chat_client.data_agent_path(
        billing_project, location, data_agent_id
    ),
)

# Form the request
request = geminidataanalytics.ChatRequest(
    parent=f"projects/{billing_project}/locations/{location}",
    messages=messages,
    data_agent_context=data_agent_context,
)

if looker_selected:
  request.credentials = credentials

# Make the request
stream = data_chat_client.chat(request=request)

# Handle the response
for response in stream:
  show_message(response)

# Create a request containing a single user message -- your question.
# fmt: off
question = "How many airports are in CA with elevation greater than 1000?"  # @param {type:"string"}
# fmt: on

messages = [
    geminidataanalytics.Message(
        user_message=geminidataanalytics.UserMessage(text=question)
    )
]

request = geminidataanalytics.ChatRequest(
    inline_context=inline_context,
    parent=f"projects/{billing_project}/locations/{location}",
    messages=messages,
)
if bq_selected:
  if use_example_queries:
    request.inline_context.example_queries = example_queries
  if use_glossary_terms:
    request.inline_context.glossary_terms = glossary_terms
elif looker_selected:
  request.credentials = credentials
  if use_looker_golden_queries:
    request.inline_context.looker_golden_queries = looker_golden_queries

# Make the request
stream = data_chat_client.chat(request=request)

# Handle the response
for response in stream:
  show_message(response)

# @title Multi-turn Conversation

# Re-used across requests to track previous turns
conversation_messages = []

# fmt: off
context = "inline_context" # @param ["data_agent_context", "inline_context"]
# Only for data_agent_context
data_agent_id = "data_agent_1"  # @param {type:"string"}
# fmt: on

data_agent_context = geminidataanalytics.DataAgentContext(
    data_agent=data_chat_client.data_agent_path(
        billing_project, location, data_agent_id
    ),
)

# Helper function for calling the API
def multi_turn_conversation(msg):
  message = geminidataanalytics.Message(
      user_message=geminidataanalytics.UserMessage(text=msg)
  )

  # Send multi-turn request by including previous turns plus new message
  conversation_messages.append(message)

  request = geminidataanalytics.ChatRequest(
      parent=f"projects/{billing_project}/locations/{location}",
      messages=conversation_messages,
      **(
          {"data_agent_context": data_agent_context}
          if context == "data_agent_context"
          else {"inline_context": inline_context}
      ),
  )

  if looker_selected:
    request.credentials = credentials
    if context == "inline_context" and use_looker_golden_queries:
      inline_context.looker_golden_queries = looker_golden_queries
  elif bq_selected and context == "inline_context":
    if use_example_queries:
      inline_context.example_queries = example_queries
    if use_glossary_terms:
      inline_context.glossary_terms = glossary_terms

  # Make the request
  stream = data_chat_client.chat(request=request)

  # Handle the response
  for response in stream:
    show_message(response)
    if response.system_message or response.user_message:
      conversation_messages.append(response)

# Send first-turn request
multi_turn_conversation(
    "List of the top 5 states by the total number of airports"
)

# Send follow-up-turn request
multi_turn_conversation("Can you please make it a pie chart?")

# @title Get Data Agent

# Initialize request argument(s)

# fmt: off
data_agent_id = "data_agent_1"  # @param {type:"string"}
# fmt: on
request = geminidataanalytics.GetDataAgentRequest(
    name=data_agent_client.data_agent_path(
        billing_project, location, data_agent_id
    )
)

# Make the request
response = data_agent_client.get_data_agent(request=request)

# Handle the response
print(response)

# @title List Data Agents

# Initialize request argument(s)
request = geminidataanalytics.ListDataAgentsRequest(
    parent=f"projects/{billing_project}/locations/{location}",
    order_by="create_time desc", # Optional
)

# Make the request
data_agents = data_agent_client.list_data_agents(request=request)

# Handle the response
display_resource_list(data_agents)

# @title List Accessible Data Agents

# fmt: off
creator_filter = "NONE"  # @param ["NONE", "CREATOR_ONLY", "NOT_CREATOR_ONLY"]
# fmt: on

# Initialize request argument(s)
request = geminidataanalytics.ListAccessibleDataAgentsRequest(
    parent=f"projects/{billing_project}/locations/{location}",
    creator_filter=creator_filter,
)

# Make the request
data_agents = data_agent_client.list_accessible_data_agents(request=request)

# Handle the response
display_resource_list(data_agents)

# @title Update Data Agent

# Initialize request argument(s)

# fmt: off
data_agent_id = "data_agent_1"  # @param {type:"string"}
# fmt: on

data_agent = geminidataanalytics.DataAgent(
    data_analytics_agent=geminidataanalytics.DataAnalyticsAgent(
        published_context=published_context
    ),
    name=data_agent_client.data_agent_path(
        billing_project, location, data_agent_id
    ),
    description="This is my new updated description.",
)


update_mask = field_mask_pb2.FieldMask(
    paths=["description", "data_analytics_agent.published_context"]
)

request = geminidataanalytics.UpdateDataAgentRequest(
    data_agent=data_agent,
    update_mask=update_mask,
)

try:
  operation = data_agent_client.update_data_agent(request=request)

  while not operation.done():
    print("Still in progress...")
    time.sleep(2)

  if operation.done():
    if operation.exception():
      print("Operation failed:", operation.exception())
    else:
      print("Operation succeeded:", operation.result())
except Exception as e:
  print(f"Error updating Data Agent: {e}")

# @title [Agent Sharing] Set IAM Policy for Data Agent

"""PLEASE NOTE THIS API CALLS OVERRIDE EXISTING PERMISSION FOR THE RESOURCE.

For preserving existing policy in practice call Get IAM policy to fetch existing
policy and pass it along with additional changes in the call to Set IAM Policy
"""

# Inputs
# fmt: off
data_agent_id = "data_agent_1"  # @param {type:"string"}
role = "roles/geminidataanalytics.dataAgentEditor"  # @param {type:"string"}
users = "abc@google.com, wxyz@google.com"  # @param {type:"string"}
# fmt: on

resource = (
    f"projects/{billing_project}/locations/global/dataAgents/{data_agent_id}"
)

# Construct the IAM policy
binding = policy_pb2.Binding(
    role=role, members=[f"user:{i.strip()}" for i in users.split(",")]
)

policy = policy_pb2.Policy(bindings=[binding])

# Create the request
request = iam_policy_pb2.SetIamPolicyRequest(resource=resource, policy=policy)

# Send the request
try:
  response = data_agent_client.set_iam_policy(request=request)
  print("IAM Policy set successfully!")
  print(f"Response: {response}")
except Exception as e:
  print(f"Error setting IAM policy: {e}")

# @title [Agent Sharing] GET IAM Policy for Data Agent

# Initialize request argument(s)

# fmt: off
data_agent_id = "data_agent_1"  # @param {type:"string"}
# fmt: on

resource = (
    f"projects/{billing_project}/locations/global/dataAgents/{data_agent_id}"
)
request = iam_policy_pb2.GetIamPolicyRequest(
    resource=resource,
)
try:
  response = data_agent_client.get_iam_policy(request=request)
  print("IAM Policy fetched successfully!")
  print(f"Response: {response}")
except Exception as e:
  print(f"Error setting IAM policy: {e}")

# @title [Soft Delete] Delete Data Agent

# Initialize request argument(s)

# fmt: off
data_agent_id = "data_agent_1"  # @param {type:"string"}
# fmt: on

request = geminidataanalytics.DeleteDataAgentRequest(
    name=data_agent_client.data_agent_path(
        billing_project, location, data_agent_id
    )
)

try:
  # Make the request
  data_agent_client.delete_data_agent(request=request)
  print("Data Agent Deleted")
except Exception as e:
  print(f"Error deleting Data Agent: {e}")

# @title Get Conversation

# fmt: off
conversation_id = "conversation_1"  # @param {type:"string"}
# fmt: on

request = geminidataanalytics.GetConversationRequest(
    name=data_chat_client.conversation_path(
        billing_project, location, conversation_id
    ),
)

# Make the request
response = data_chat_client.get_conversation(request=request)

# Handle the response
print(response)

# @title List Conversation

request = geminidataanalytics.ListConversationsRequest(
    parent=f"projects/{billing_project}/locations/{location}",
)

# Make the request
conversations = data_chat_client.list_conversations(request=request)

# Handle the response
display_resource_list(conversations)

# @title List Messages

# fmt: off
conversation_id = "conversation_1"  # @param {type:"string"}
# fmt: on

request = geminidataanalytics.ListMessagesRequest(
    parent=data_chat_client.conversation_path(
        billing_project, location, conversation_id
    ),
)

# Make the request
response = data_chat_client.list_messages(request=request)

# Handle the response
print(response)

# @title Delete Conversation

# fmt: off
conversation_id = "conversation_1"  # @param {type:"string"}
# fmt: on

request = geminidataanalytics.DeleteConversationRequest(
    name=data_chat_client.conversation_path(
        billing_project, location, conversation_id
    ),
)

# Make the request
response = data_chat_client.delete_conversation(request=request)

# Handle the response
print(response)