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

%pip install -q google-cloud-aiplatform[adk,agent_engines]
%pip install --upgrade --force-reinstall -q google-cloud-aiplatform[evaluation]

import sys

if "google.colab" in sys.modules:
    from google.colab import auth

    auth.authenticate_user()

import os
from urllib.parse import urlparse

import vertexai
from vertexai import Client
from google.cloud import storage
from google.genai import types as genai_types

def get_config_value(initial_value: str, placeholder: str, env_var_name: str) -> str:
    """Gets a configuration value or environment variable if unspecified."""
    if not initial_value or initial_value == placeholder:
        return os.environ.get(env_var_name)
    return initial_value


def get_or_create_gcs_bucket(project_id: str, gcs_dest: str) -> str:
    """Retrieves GCS bucket or creates a default."""
    gcs_dest = gcs_dest or f"{project_id}/agent-evaluation"
    storage_client = storage.Client(project=project_id)
    bucket = gcs_dest.replace("gs://", "").split("/")[0]
    if not storage_client.lookup_bucket(bucket):
        print(f"Creating bucket: {bucket}")
        storage_client.create_bucket(bucket)
    if not gcs_dest.startswith("gs://"):
        return f"gs://{gcs_dest}"
    return gcs_dest


# Configuration
# fmt: off
PROJECT_ID = ""  # @param {type: "string", placeholder: "[your-project-id]", isTemplate: true}
PROJECT_ID = get_config_value(PROJECT_ID, "[your-project-id]", "GOOGLE_CLOUD_PROJECT")
assert PROJECT_ID, "Please specify a valid project"

LOCATION = ""  # @param {type: "string", placeholder: "[us-central1]", isTemplate: true}
LOCATION = get_config_value(LOCATION, None, "GOOGLE_CLOUD_REGION")
assert LOCATION, "Please specify a valid location"

GCS_DEST = ""  # @param {type: "string", placeholder: "[your-gcs-bucket]", isTemplate: true}
# fmt: on
GCS_DEST = get_config_value(
    GCS_DEST, "[your-gcs-bucket]", "GOOGLE_CLOUD_STORAGE_BUCKET"
)
GCS_DEST = get_or_create_gcs_bucket(PROJECT_ID, GCS_DEST)
assert GCS_DEST, "Please specify a valid GCS destination"
STAGING_BUCKET = f"gs://{urlparse(GCS_DEST).netloc}"

# Initialize SDK
vertexai.init(project=PROJECT_ID, location=LOCATION)
client = Client(
    project=PROJECT_ID,
    location=LOCATION,
    http_options=genai_types.HttpOptions(api_version="v1beta1"),
)

import time

import pandas as pd
from google.adk import Agent
from vertexai import types

# Define Agent Tools
def search_products(query: str):
    """Searches for products based on a query.

    Args:
        query: The search query.

    Returns:
        A list of products that match the query.
    """
    # Mock response for demonstration
    if "headphones" in query.lower():
        return {"products": [{"name": "Wireless Headphones", "id": "B08H8H8H8H"}]}
    return {"products": []}


def get_product_details(product_id: str):
    """Gets the details for a given product ID.

    Args:
        product_id: The ID of the product.

    Returns:
        The details of the product.
    """
    if product_id == "B08H8H8H8H":
        return {"details": "Noise-cancelling, 20-hour battery life."}
    return {"error": "Product not found."}


def add_to_cart(product_id: str, quantity: int):
    """Adds a specified quantity of a product to the cart.

    Args:
        product_id: The ID of the product.
        quantity: The quantity to add to the cart.

    Returns:
        A status message indicating the addition to the cart.
    """
    return {"status": f"Added {quantity} of {product_id} to cart."}


# Define Agent
ecommerce_agent = Agent(
    model="gemini-2.5-flash",
    name="ecommerce_agent",
    instruction="You are an ecommerce expert",
    tools=[search_products, get_product_details, add_to_cart],
)

# Deploy Agent
app = vertexai.agent_engines.AdkApp(
    agent=ecommerce_agent,
)
agent_engine = client.agent_engines.create(
    agent=app,
    config={
        "staging_bucket": STAGING_BUCKET,
        "requirements": ["google-cloud-aiplatform[adk,agent_engines]"],
        "env_vars": {"GOOGLE_CLOUD_AGENT_ENGINE_ENABLE_TELEMETRY": "true"},
    },
)
agent_engine_resource_name = agent_engine.api_resource.name

session_inputs = types.evals.SessionInput(
    user_id="user_123",
    state={},
)
ecommerce_prompts = [
    "Search for 'noise-cancelling headphones'.",
    "Show me the details for product 'B08H8H8H8H'.",
    "Add one pair of 'B08H8H8H8H' to my shopping cart.",
    "Find 'wireless ear buds' and then add the first result to my cart.",
    "I need a new laptop for work, can you find one with at least 16GB of RAM?",
]
ecommerce_dataset = pd.DataFrame(
    {
        "prompt": ecommerce_prompts,
        "session_inputs": [session_inputs] * len(ecommerce_prompts),
    }
)

# Run inference
agent_dataset_with_inference = client.evals.run_inference(
    agent=agent_engine_resource_name,
    src=ecommerce_dataset,
)
# Display inference results
agent_dataset_with_inference.show()

# Create agent_info from Agent definition and deployed resource name
ecommerce_agent_info = types.evals.AgentInfo.load_from_agent(
    ecommerce_agent, agent_engine_resource_name
)

# Evaluate Agent Dataset
evaluation_run = client.evals.create_evaluation_run(
    dataset=agent_dataset_with_inference,
    agent_info=ecommerce_agent_info,
    metrics=[
        types.RubricMetric.FINAL_RESPONSE_QUALITY,
        types.RubricMetric.TOOL_USE_QUALITY,
        types.RubricMetric.HALLUCINATION,
        types.RubricMetric.SAFETY,
    ],
    dest=GCS_DEST,
)

# Display status and results
evaluation_run.show()

completed_states = set(
    [
        "SUCCEEDED",
        "FAILED",
        "CANCELLED",
    ]
)

while evaluation_run.state not in completed_states:
    evaluation_run.show()
    evaluation_run = client.evals.get_evaluation_run(name=evaluation_run.name)
    time.sleep(5)
evaluation_run = client.evals.get_evaluation_run(
    name=evaluation_run.name, include_evaluation_items=True
)

evaluation_run.show()

# Create agent_info from Agent definition and deployed resource name
ecommerce_agent_info = types.evals.AgentInfo.load_from_agent(
    ecommerce_agent, agent_engine_resource_name
)

# Evaluate Agent Dataset
eval_result = client.evals.evaluate(
    dataset=agent_dataset_with_inference,
    agent_info=ecommerce_agent_info,
    metrics=[
        types.RubricMetric.FINAL_RESPONSE_QUALITY,
        types.RubricMetric.TOOL_USE_QUALITY,
        types.RubricMetric.HALLUCINATION,
        types.RubricMetric.SAFETY,
    ],
)

# Display results
eval_result.show()