# Copyright 2025 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 sys

# Additional authentication is required for Google Colab
if "google.colab" in sys.modules:
    # Authenticate user to Google Cloud
    from google.colab import auth

    auth.authenticate_user()

# Use the environment variable if the user doesn't provide Project ID.
import os

# fmt: off
PROJECT_ID = "your-project-id"  # @param {type: "string", placeholder: "[your-project-id]", isTemplate: true}
# fmt: on
if not PROJECT_ID or PROJECT_ID == "[your-project-id]":
    PROJECT_ID = str(os.environ.get("GOOGLE_CLOUD_PROJECT"))
LOCATION = os.environ.get("GOOGLE_CLOUD_REGION", "us-central1")

# Initialize the aiplatform package
from google.cloud import aiplatform

aiplatform.init(project=PROJECT_ID, location=LOCATION)

%%bigquery --project $PROJECT_ID
CREATE SCHEMA import_example_dataset; -- OPTIONS (location=$LOCATION);
CREATE TABLE import_example_dataset.test_table (
  id INTEGER,
  embedding ARRAY <FLOAT64>,
  allow_column STRING,
  deny_column STRING,
  int_column INTEGER,
  float_column FLOAT64,
  metadata_column STRING
);
INSERT INTO import_example_dataset.test_table(id, embedding, allow_column, deny_column, int_column, float_column, metadata_column) VALUES
(1, [0.1, 0.1, 0.1], "allow1", "deny1", 1, 0.1, "metadata1"),
(2, [0.2, 0.2, 0.2], "allow2", "deny2", 2, 0.2, "metadata2");
SELECT * FROM import_example_dataset.test_table;

my_index = aiplatform.MatchingEngineIndex.create_tree_ah_index(
    display_name="import_test_index_name",
    dimensions=3,
    approximate_neighbors_count=10,
    index_update_method="BATCH_UPDATE",
)

import requests

# Get token to use for REST request
gcloud_token = !gcloud auth print-access-token

url = f"https://{LOCATION}-aiplatform.googleapis.com/v1beta1/{my_index.resource_name}:import"
headers = {
    "Content-Type": "application/json",
    "Authorization": f"Bearer {gcloud_token[0]}",
}
request = {
    "is_complete_overwrite": True,
    "config": {
        "big_query_source_config": {
            "table_path": f"bq://{PROJECT_ID}.import_example_dataset.test_table",
            "datapoint_field_mapping": {
                "id_column": "id",
                "embedding_column": "embedding",
                "restricts": [
                    {
                        "namespace": "restrict",
                        "allow_column": ["allow_column"],
                        "deny_column": ["deny_column"],
                    },
                ],
                "numeric_restricts": [
                    {
                        "namespace": "int_restrict",
                        "value_column": "int_column",
                        "value_type": "INT",
                    },
                    {
                        "namespace": "float_restrict",
                        "value_column": "float_column",
                        "value_type": "FLOAT",
                    },
                ],
                # The import may include metadata if your project has been allow-listed
                # for the VS metadata preview.
                # "embedding_metadata": "metadata_column"
            },
        }
    },
}

try:
    # Make the POST request
    response = requests.post(url, headers=headers, json=request)

    # Check the response status code
    if response.status_code == 200:
        print("Import request successful!")
        print(response.json())
    else:
        print(f"Import request failed with status code: {response.status_code}")
        print(response.text)

except requests.exceptions.RequestException as e:
    print(f"An error occurred: {e}")

operation = response.json()["name"]
response = requests.get(
    f"https://{LOCATION}-aiplatform.googleapis.com/v1beta1/{operation}", headers=headers
)

if "done" in response.json():
    if "error" in response.json():
        print("Import failed")
        print(response.json()["error"])
    else:
        print("Import succeeded!")
        print(response.json())
else:
    print("Import still in progress")
print(response.json())