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

import json
import textwrap

import altair as alt
from google.colab import auth
from IPython.display import HTML, display
import pandas as pd
from pygments import formatters, highlight, lexers
import requests

auth.authenticate_user()

access_token = !gcloud auth application-default print-access-token
headers = {
    "Authorization": f"Bearer {access_token[0]}",
    "Content-Type": "application/json",
}

# @title Billing Project and Prompt (System Instruction)

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

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

bigquery_data_sources = {}
looker_data_sources = {}
looker_studio_data_sources = {}

# @title BigQuery Datasource

bigquery_data_sources = {
    "bq": {
        "tableReferences": [
            {
                "projectId": "bigquery-public-data",
                "datasetId": "faa",
                "tableId": "us_airports",
            },
            {
                "projectId": "bigquery-public-data",
                "datasetId": "san_francisco",
                "tableId": "street_trees",
            },
            # Add more table references here
        ]
    }
}

# Optional example queries (only leveraged for BigQuery datasources currently)
# Example queries can be parameterized in which case parameters must be defined.
# Each parameter requires a name (exactly as it appears in the question and SQL),
# description, and a valid dataType
# (https://docs.cloud.google.com/bigquery/docs/reference/standard-sql/data-types#data_type_list).
example_queries = [
    {
        "naturalLanguageQuestion": (
            "What is the highest observed positive longitude?"
        ),
        "sqlQuery": "SELECT MAX(longitude) FROM airports",
    },
    {
        "naturalLanguageQuestion": (
            "How many airports are in {state} with elevation"
            " greater than {elevation}?"
        ),
        "sqlQuery": (
            "SELECT COUNT(*) FROM `bigquery-public-data.faa.us_airports` WHERE"
            " LOWER(state_abbreviation) = @state AND elevation > @elevation"
        ),
        "parameters": [
            {
                "name": "state",
                "dataType": "STRING",
                "description": "State abbreviation in lower case",
            },
            {
                "name": "elevation",
                "dataType": "FLOAT64",
                "description": "Elevation in feet",
            },
        ],
    },
]

glossary_terms = [
    {
        "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"}
use_example_queries = True  # @param {type:"boolean"}

# @title Looker Datasource

"""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)
"""

client_id = "<add client_id here>"  # @param {type:"string"}
client_secret = "<add client_secret here>"  # @param {type:"string"}
looker_credentials = {
    "oauth": {
        "secret": {
            "client_id": client_id,
            "client_secret": client_secret,
        }
    }
}

# Uncomment this and comment out the above one if you are using access_token for authentication
# access_token = "<add access_token here>"  # @param {type:"string"}
# looker_credentials = {
#     "oauth": {
#         "token": {
#           "access_token": access_token,
#         }
#     }
# }

# 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"}


""" Looker datasources 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_data_sources = {
    "looker": {
        "explore_references": [
            {
                "looker_instance_uri": looker_instance_uri_1,
                "lookml_model": lookml_model_1,
                "explore": explore_1,
            },
            # {
            #     "looker_instance_uri": looker_instance_uri_2,
            #     "lookml_model": lookml_model_2,
            #     "explore": explore_2,
            # },
            # Add up to 5 total explore references here
        ],
    }
}


""" Looker datasource 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_data_sources = {
#     "looker": {
#       "explore_references": [
#           {
#             "private_looker_instance_info": {
#                 "looker_instance_id": looker_instance_id_1,
#                 "service_directory_name": looker_service_directory_name_1,
#             },
#             "lookml_model": lookml_model_1,
#             "explore": explore_1,
#           },
#           # {
#           #   "private_looker_instance_info": {
#           #       "looker_instance_id": looker_instance_id_2,
#           #       "service_directory_name": looker_service_directory_name_2,
#           #   },
#           #   "lookml_model": lookml_model_2,
#           #   "explore": explore_2,
#           # },
# .          # Add up to 5 total explore references here
#       ],
# }


looker_golden_queries = [
    {
        "natural_language_questions": (
            "What is the highest observed positive longitude?"
        ),
        "looker_query": {
            "model": "airports",
            "explore": "airports",
            "fields": "airports.longitude",
            "filters": {
                "field": "airports.longitude",
                "value": ">0",
            },
            "sorts": "airports.longitude desc",
            "limit": "1",
        },
    },
    {
        "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": {
            "model": "airports",
            "explore": "airports",
            "fields": "airports.city",
            "fields": "airports.code",
            "filters": {
                "field": "airports.major",
                "value": "Y",
            },
            "filters": {
                "field": "airports.state",
                "value": "CA",
            },
        },
    },
]

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

# @title Looker Studio Datasource

looker_studio_data_sources = {
    "studio": {
        "studio_references": [{"datasource_id": "your_studio_datasource_id"}]
    }
}

# @title Select datasource

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

datasource_map = {
    "bigquery_data_sources": bigquery_data_sources,
    "looker_data_sources": looker_data_sources,
    "looker_studio_data_sources": looker_studio_data_sources,
}

datasource_references = datasource_map[selected_datasource]

bq_selected = False
looker_selected = False
looker_studio_selected = False

if selected_datasource == "looker_data_sources":
  looker_selected = True
elif selected_datasource == "bigquery_data_sources":
  bq_selected = True
elif selected_datasource == "looker_studio_data_sources":
  looker_studio_selected = True

# fmt: off
environment = "prod-global"  # @param ["prod-global", "prod-us", "prod-eu", "prod-us-east4", "staging-global", "autopush-global"]
location = ""
api_version = "v1"
base_url = ""
# fmt: on

if environment == "prod-us":
  base_url = "https://geminidataanalytics.us.rep.googleapis.com"
  location = "us"
elif environment == "prod-eu":
  base_url = "https://geminidataanalytics.eu.rep.googleapis.com"
  location = "eu"
elif environment == "prod-us-east4":
  base_url = "https://geminidataanalytics-us-east4.googleapis.com"
  location = "us-east4"
elif environment == "autopush-global":
  base_url = "https://autopush-geminidataanalytics.sandbox.googleapis.com"
  location = "global"
elif environment == "staging-global":
  base_url = "https://staging-geminidataanalytics.sandbox.googleapis.com"
  location = "global"
else:
  base_url = "https://geminidataanalytics.googleapis.com"
  location = "global"

# @title Parser for List methods

def display_resource_list(resource_list):
    for i, resource in enumerate(resource_list):
      resource_name = resource.get("name").split("/")[-1]
      resource_display_name = resource.get("displayName")
      formatted_json = json.dumps(resource, 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 Streaming Chat Messages


def is_json(str):
  try:
    json_object = json.loads(str)
  except ValueError:
    return False
  return True


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 get_property(data, field_name, default=""):
  return data[field_name] if field_name in data else default


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


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


def format_bq_table_ref(table_ref):
  return "{}.{}.{}".format(
      table_ref["projectId"], table_ref["datasetId"], table_ref["tableId"]
  )


def format_looker_table_ref(table_ref):
  return "lookmlModel: {}, explore: {}".format(
      table_ref["lookmlModel"], table_ref["explore"]
  )


def display_datasource(datasource):
  source_name = ""

  if "studioDatasourceId" in datasource:
    source_name = datasource["studioDatasourceId"]
  elif "lookerExploreReference" in datasource:
    source_name = format_looker_table_ref(datasource["lookerExploreReference"])
  else:
    source_name = format_bq_table_ref(datasource["bigqueryTableReference"])

  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("Query name: {}".format(query["name"]))
    if "question" in query:
      print("Question: {}".format(query["question"]))
    if "datasources" in query:
      print("Data sources:")
      for datasource in query["datasources"]:
        display_datasource(datasource)
  elif "generatedSql" in resp:
    display_section_title("SQL generated")
    print(resp["generatedSql"])
  elif "result" in resp:
    display_section_title("Data retrieved")

    fields = map(
        lambda field: get_property(field, "name"),
        resp["result"]["schema"]["fields"],
    )
    dict = {}

    for field in fields:
      dict[field] = [get_property(el, field) for el in resp["result"]["data"]]

    display(pd.DataFrame(dict))


def handle_chart_response(resp):
  if "query" in resp and "instructions" in resp["query"]:
    print(resp["query"]["instructions"])
  elif "result" in resp:
    vegaConfig = resp["result"]["vegaConfig"]
    alt.Chart.from_json(json.dumps(vegaConfig)).display()


def handle_error(resp):
  display_section_title("Error")
  print("Code: {}".format(resp["code"]))
  print("Message: {}".format(resp["message"]))


def get_stream(url, payload):
  s = requests.Session()

  acc = ""

  with s.post(url, json=payload, headers=headers, stream=True) as resp:
    for line in resp.iter_lines():
      if not line:
        continue

      decoded_line = str(line, encoding="utf-8")

      if decoded_line == "[{":
        acc = "{"
      elif decoded_line == "}]":
        acc += "}"
      elif decoded_line == ",":
        continue
      else:
        acc += decoded_line

      if not is_json(acc):
        continue

      data_json = json.loads(acc)

      if "systemMessage" not in data_json:
        if "error" in data_json:
          handle_error(data_json["error"])
        continue

      if "text" in data_json["systemMessage"]:
        handle_text_response(data_json["systemMessage"]["text"])
      elif "schema" in data_json["systemMessage"]:
        handle_schema_response(data_json["systemMessage"]["schema"])
      elif "data" in data_json["systemMessage"]:
        handle_data_response(data_json["systemMessage"]["data"])
      elif "chart" in data_json["systemMessage"]:
        handle_chart_response(data_json["systemMessage"]["chart"])
      else:
        colored_json = highlight(
            acc, lexers.JsonLexer(), formatters.TerminalFormatter()
        )
        print(colored_json)
      print("\n")
      acc = ""

# @title get_stream function for multi-turn conversation

# NOTE: This methods is same as get_stream() method present in Streaming Chat Messages section.
# The only difference is that - here we are storing the response in an array "conversation_messages" to save the conversation.


def get_stream_multi_turn(url, payload, conversation_messages):
  s = requests.Session()

  acc = ""

  with s.post(url, json=payload, headers=headers, stream=True) as resp:
    for line in resp.iter_lines():
      if not line:
        continue

      decoded_line = str(line, encoding="utf-8")

      if decoded_line == "[{":
        acc = "{"
      elif decoded_line == "}]":
        acc += "}"
      elif decoded_line == ",":
        continue
      else:
        acc += decoded_line

      if not is_json(acc):
        continue

      data_json = json.loads(acc)
      # Store the response to be used in next iteration.
      conversation_messages.append(data_json)

      if "systemMessage" not in data_json:
        if "error" in data_json:
          handle_error(data_json["error"])
        continue

      if "text" in data_json["systemMessage"]:
        handle_text_response(data_json["systemMessage"]["text"])
      elif "schema" in data_json["systemMessage"]:
        handle_schema_response(data_json["systemMessage"]["schema"])
      elif "data" in data_json["systemMessage"]:
        handle_data_response(data_json["systemMessage"]["data"])
      elif "chart" in data_json["systemMessage"]:
        handle_chart_response(data_json["systemMessage"]["chart"])
      else:
        colored_json = highlight(
            acc, lexers.JsonLexer(), formatters.TerminalFormatter()
        )
        print(colored_json)
      print("\n")
      acc = ""

# @title Create Data Agent

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

data_agent_url = f"{base_url}/{api_version}/projects/{billing_project}/locations/{location}/dataAgents"

data_agent_payload = {
    "name": (
        f"projects/{billing_project}/locations/{location}/dataAgents/{data_agent_id}"
    ),  # Optional
    "description": "This is the description of data_agent.",  # Optional
    "data_analytics_agent": {
        "published_context": {
            "datasource_references": datasource_references,
            "system_instruction": system_instruction,
            "options": {
                "analysis": {
                    # Optional - if wanting to use advanced analysis with python.
                    # Default is False.
                    "python": {"enabled": False}
                }
            },
        }
    },
}

if bq_selected:
  if use_example_queries:
    data_agent_payload["data_analytics_agent"]["published_context"][
        "example_queries"
    ] = example_queries
  if use_glossary_terms:
    data_agent_payload["data_analytics_agent"]["published_context"][
        "glossary_terms"
    ] = glossary_terms
elif looker_selected and use_looker_golden_queries:
  data_agent_payload["data_analytics_agent"]["published_context"][
      "looker_golden_queries"
  ] = looker_golden_queries

params = {"data_agent_id": data_agent_id}  # Optional

data_agent_response = requests.post(
    data_agent_url, params=params, json=data_agent_payload, headers=headers
)

if data_agent_response.status_code == 200:
  print("Data Agent created successfully!")
  print(json.dumps(data_agent_response.json(), indent=2))
else:
  print(f"Error creating Data Agent: {data_agent_response.status_code}")
  print(data_agent_response.text)

# @title Get Data Agent

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

data_agent_url = f"{base_url}/{api_version}/projects/{billing_project}/locations/{location}/dataAgents/{data_agent_id}"

data_agent_response = requests.get(data_agent_url, headers=headers)

if data_agent_response.status_code == 200:
  print("Fetched Data Agent successfully!")
  print(json.dumps(data_agent_response.json(), indent=2))
else:
  print(f"Error: {data_agent_response.status_code}")
  print(data_agent_response.text)

# @title List Data Agents
data_agent_url = f"{base_url}/{api_version}/projects/{billing_project}/locations/{location}/dataAgents"

params = {
    "orderBy": "create_time desc", # Optional
}
data_agent_response = requests.get(data_agent_url, headers=headers, params=params)

if data_agent_response.status_code == 200:
  print("Data Agents Listed successfully!")
  agents = data_agent_response.json().get("dataAgents", [])
  display_resource_list(agents)
else:
  print(f"Error Listing Data Agents: {data_agent_response.status_code}")
  print(data_agent_response.text)

# @title List Accessible Data Agents

data_agent_url = f"{base_url}/{api_version}/projects/{billing_project}/locations/{location}/dataAgents:listAccessible"

# Choose the desired creator filter
# fmt: off
creator_filter = "NONE"  # @param ["NONE", "CREATOR_ONLY", "NOT_CREATOR_ONLY"]
# fmt: on

# Add the creator_filter as a query parameter
params = {"creator_filter": creator_filter}

data_agent_response = requests.get(
    data_agent_url, headers=headers, params=params
)

if data_agent_response.status_code == 200:
  print("Data Agents Listed successfully!")
  agents = data_agent_response.json().get("dataAgents", [])
  display_resource_list(agents)
else:
  print(f"Error Listing Data Agents: {data_agent_response.status_code}")
  print(data_agent_response.text)

# @title Update Data Agent

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

data_agent_url = f"{base_url}/{api_version}/projects/{billing_project}/locations/{location}/dataAgents/{data_agent_id}"

payload = {
    "description": "This is the latest description of data_agent.",
    "data_analytics_agent": {
        "published_context": {
            "datasource_references": datasource_references,
            "system_instruction": system_instruction,
        }
    },
}

fields = ["description", "data_analytics_agent"]
params = {"updateMask": ",".join(fields)}

data_agent_response = requests.patch(
    data_agent_url, headers=headers, params=params, json=payload
)

if data_agent_response.status_code == 200:
  print("Data Agent updated successfully!")
  print(json.dumps(data_agent_response.json(), indent=2))
else:
  print(f"Error Updating Data Agent: {data_agent_response.status_code}")
  print(data_agent_response.text)

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


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

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

# fmt: off
data_agent_id = "data_agent_1"  # @param {type:"string"}
users = "abc@google.com, wxyz@google.com"  # @param {type:"string"}
# fmt: on

data_agent_url = f"{base_url}/{api_version}/projects/{billing_project}/locations/{location}/dataAgents/{data_agent_id}:setIamPolicy"


# Request body
payload = {
    "policy": {
        "bindings": [{
            "role": "roles/geminidataanalytics.dataAgentEditor",
            "members": [f"user:{i.strip()}" for i in users.split(",")],
        }]
    }
}

data_agent_response = requests.post(
    data_agent_url, headers=headers, json=payload
)

if data_agent_response.status_code == 200:
  print("IAM Policy set successfully!")
  print(json.dumps(data_agent_response.json(), indent=2))
else:
  print(f"Error setting IAM policy: {data_agent_response.status_code}")
  print(data_agent_response.text)

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

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

data_agent_url = f"{base_url}/{api_version}/projects/{billing_project}/locations/{location}/dataAgents/{data_agent_id}:getIamPolicy"

# Request body
payload = {
    "resource": (
        f"projects/{billing_project}/locations/{location}/dataAgents/{data_agent_id}"
    )
}

data_agent_response = requests.post(
    data_agent_url, headers=headers, json=payload
)

if data_agent_response.status_code == 200:
  print("IAM Policy fetched successfully!")
  print(json.dumps(data_agent_response.json(), indent=2))
else:
  print(f"Error fetching IAM policy: {data_agent_response.status_code}")
  print(data_agent_response.text)

# @title [Soft Delete] Delete Data Agent

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

data_agent_url = f"{base_url}/{api_version}/projects/{billing_project}/locations/{location}/dataAgents/{data_agent_id}"

data_agent_response = requests.delete(data_agent_url, headers=headers)

if data_agent_response.status_code == 200:
  print("Data Agent deleted successfully!")
  print(json.dumps(data_agent_response.json(), indent=2))
else:
  print(f"Error Deleting Data Agent: {data_agent_response.status_code}")
  print(data_agent_response.text)

# @title Create Conversation

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

conversation_url = f"{base_url}/{api_version}/projects/{billing_project}/locations/{location}/conversations"

conversation_payload = {
    "agents": [
        f"projects/{billing_project}/locations/{location}/dataAgents/{data_agent_id}"
    ],
    "name": (
        f"projects/{billing_project}/locations/{location}/conversations/{conversation_id}"
    ),
}
params = {"conversation_id": conversation_id}

conversation_response = requests.post(
    conversation_url, headers=headers, params=params, json=conversation_payload
)

if conversation_response.status_code == 200:
  print("Conversation created successfully!")
  print(json.dumps(conversation_response.json(), indent=2))
else:
  print(f"Error creating Conversation: {conversation_response.status_code}")
  print(conversation_response.text)

# @title Get Conversation

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

conversation_url = f"{base_url}/{api_version}/projects/{billing_project}/locations/{location}/conversations/{conversation_id}"

conversation_response = requests.get(conversation_url, headers=headers)

# Handle the response
if conversation_response.status_code == 200:
  print("Conversation fetched successfully!")
  print(json.dumps(conversation_response.json(), indent=2))
else:
  print(f"Error while fetching conversation: {conversation_response.status_code}")
  print(conversation_response.text)

# @title List Conversation

conversation_url = f"{base_url}/{api_version}/projects/{billing_project}/locations/{location}/conversations"


conversation_response = requests.get(conversation_url, headers=headers)

# Handle the response
if conversation_response.status_code == 200:
  print("Conversation fetched successfully!")
  conversations = conversation_response.json().get("conversations", [])
  display_resource_list(conversations)
else:
  print(f"Error while fetching conversations: {conversation_response.status_code}")
  print(conversation_response.text)

# @title [Stateful] Chat using Conversation

data_agent_id = "data_agent_1"  # @param {type:"string"}
conversation_id = "conversation_1"  # @param {type:"string"}

# fmt: off
question = "Make a bar graph for the top 5 states by the total number of airports"  # @param {type:"string"}
# fmt: on

chat_url = f"{base_url}/{api_version}/projects/{billing_project}/locations/{location}:chat"

# Construct the payload
chat_payload = {
    "parent": f"projects/{billing_project}/locations/global",
    "messages": [{"userMessage": {"text": question}}],
    "conversation_reference": {
        "conversation": (
            f"projects/{billing_project}/locations/{location}/conversations/{conversation_id}"
        ),
        "data_agent_context": {
            "data_agent": (
                f"projects/{billing_project}/locations/{location}/dataAgents/{data_agent_id}"
            ),
        },
    },
}

if looker_selected:
  chat_payload["credentials"] = looker_credentials

# Call get_stream method to stream the response
get_stream(chat_url, chat_payload)

# @title [Stateful] Chat using Data Agent

data_agent_id = "data_agent_1"  # @param {type:"string"}
# fmt: off
question = "Make a bar graph for the top 5 states by the total number of airports"  # @param {type:"string"}
# fmt: on

chat_url = f"{base_url}/{api_version}/projects/{billing_project}/locations/{location}:chat"

# Construct the payload
chat_payload = {
    "parent": f"projects/{billing_project}/locations/global",
    "messages": [{"userMessage": {"text": question}}],
    "data_agent_context": {
        "data_agent": (
            f"projects/{billing_project}/locations/{location}/dataAgents/{data_agent_id}"
        ),
    },
}

if looker_selected:
  chat_payload["credentials"] = looker_credentials

# Call get_stream method to stream the response
get_stream(chat_url, chat_payload)

# @title [Stateless] Chat using Inline Context

chat_url = (
    f"{base_url}/{api_version}/projects/{billing_project}/locations/global:chat"
)

# fmt: off
question = "How many airports are in CA with elevation greater than 1000?"  # @param {type:"string"}
# fmt: on

# Construct the payload
chat_payload = {
    "parent": f"projects/{billing_project}/locations/global",
    "messages": [{"userMessage": {"text": question}}],
    "inline_context": {
        "datasource_references": datasource_references,
        "options": {
            "analysis": {
                # Optional - if wanting to use advanced analysis with python.
                # Default is False.
                "python": {"enabled": False}
            }
        },
    },
}

if bq_selected:
  if use_example_queries:
    chat_payload["inline_context"]["example_queries"] = example_queries
  if use_glossary_terms:
    chat_payload["inline_context"]["glossary_terms"] = glossary_terms
elif looker_selected:
  chat_payload["credentials"] = looker_credentials
  if use_looker_golden_queries:
    chat_payload["inline_context"]["looker_golden_queries"] = looker_golden_queries

# Call get_stream method to stream the response
get_stream(chat_url, chat_payload)

# @title List Messages

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

conversation_url = f"{base_url}/{api_version}/projects/{billing_project}/locations/{location}/conversations/{conversation_id}/messages"

conversation_response = requests.get(conversation_url, headers=headers)

# Handle the response
if conversation_response.status_code == 200:
  print("Conversation fetched successfully!")
  print(json.dumps(conversation_response.json(), indent=2))
else:
  print(f"Error while fetching conversation: {conversation_response.status_code}")
  print(conversation_response.text)

# @title Delete Conversation

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

conversation_url = f"{base_url}/{api_version}/projects/{billing_project}/locations/{location}/conversations/{conversation_id}"

conversation_response = requests.delete(conversation_url, headers=headers)

# Handle the response
if conversation_response.status_code == 200:
  print("Conversation deleted successfully!")
  print(json.dumps(conversation_response.json(), indent=2))
else:
  print(f"Error while deleting conversation: {conversation_response.status_code}")
  print(conversation_response.text)

# @title Multi-turn Conversation helper function
chat_url = (
    f"{base_url}/{api_version}/projects/{billing_project}/locations/global:chat"
)

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

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


# Helper function for calling the API
def multi_turn_conversation(msg):
  userMessage = {"userMessage": {"text": msg}}

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

  # Construct the payload
  chat_payload = (
      {
          "parent": f"projects/{billing_project}/locations/global",
          "messages": conversation_messages,
          "data_agent_context": {
              "data_agent": (
                  f"projects/{billing_project}/locations/{location}/dataAgents/{data_agent_id}"
              ),
          },
      }
      if context == "data_agent_context"
      else {
          "parent": f"projects/{billing_project}/locations/global",
          "messages": conversation_messages,
          "inline_context": {
              "datasource_references": datasource_references,
          },
      }
  )

  if looker_selected:
    chat_payload["credentials"] = looker_credentials
    if context == "inline_context" and use_looker_golden_queries:
      chat_payload["inline_context"][
          "looker_golden_queries"
      ] = looker_golden_queries
  elif bq_selected and context == "inline_context":
    if use_example_queries:
      chat_payload["inline_context"]["example_queries"] = example_queries
    if use_glossary_terms:
      chat_payload["inline_context"]["glossary_terms"] = glossary_terms

  # Call get_stream_multi_turn method to stream the response
  get_stream_multi_turn(chat_url, chat_payload, conversation_messages)

# 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?")