# 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
#
#     http://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.
#

# Code generated by the Google Gen AI SDK generator DO NOT EDIT.

import functools
import json
import logging
import mimetypes
import secrets
import time
from typing import Any, Iterator, Optional, Union
from urllib.parse import urlencode

from google import genai
from google.cloud import iam_credentials_v1  # type: ignore[attr-defined]
from google.genai import _api_module
from google.genai import _common
from google.genai import types as genai_types
from google.genai._common import get_value_by_path as getv
from google.genai._common import set_value_by_path as setv
from google.genai.pagers import Pager

from . import _agent_engines_utils
from . import types


logger = logging.getLogger("vertexai_genai.sandboxes")

logger.setLevel(logging.INFO)


def _CreateAgentEngineSandboxConfig_to_vertex(
    from_object: Union[dict[str, Any], object],
    parent_object: Optional[dict[str, Any]] = None,
) -> dict[str, Any]:
    to_object: dict[str, Any] = {}

    if getv(from_object, ["display_name"]) is not None:
        setv(parent_object, ["displayName"], getv(from_object, ["display_name"]))

    if getv(from_object, ["description"]) is not None:
        setv(parent_object, ["description"], getv(from_object, ["description"]))

    if getv(from_object, ["ttl"]) is not None:
        setv(parent_object, ["ttl"], getv(from_object, ["ttl"]))

    return to_object


def _CreateAgentEngineSandboxRequestParameters_to_vertex(
    from_object: Union[dict[str, Any], object],
    parent_object: Optional[dict[str, Any]] = None,
) -> dict[str, Any]:
    to_object: dict[str, Any] = {}
    if getv(from_object, ["name"]) is not None:
        setv(to_object, ["_url", "name"], getv(from_object, ["name"]))

    if getv(from_object, ["spec"]) is not None:
        setv(to_object, ["spec"], getv(from_object, ["spec"]))

    if getv(from_object, ["config"]) is not None:
        setv(
            to_object,
            ["config"],
            _CreateAgentEngineSandboxConfig_to_vertex(
                getv(from_object, ["config"]), to_object
            ),
        )

    return to_object


def _DeleteAgentEngineSandboxRequestParameters_to_vertex(
    from_object: Union[dict[str, Any], object],
    parent_object: Optional[dict[str, Any]] = None,
) -> dict[str, Any]:
    to_object: dict[str, Any] = {}
    if getv(from_object, ["name"]) is not None:
        setv(to_object, ["_url", "name"], getv(from_object, ["name"]))

    if getv(from_object, ["config"]) is not None:
        setv(to_object, ["config"], getv(from_object, ["config"]))

    return to_object


def _ExecuteCodeAgentEngineSandboxRequestParameters_to_vertex(
    from_object: Union[dict[str, Any], object],
    parent_object: Optional[dict[str, Any]] = None,
) -> dict[str, Any]:
    to_object: dict[str, Any] = {}
    if getv(from_object, ["name"]) is not None:
        setv(to_object, ["_url", "name"], getv(from_object, ["name"]))

    if getv(from_object, ["inputs"]) is not None:
        setv(to_object, ["inputs"], [item for item in getv(from_object, ["inputs"])])

    if getv(from_object, ["config"]) is not None:
        setv(to_object, ["config"], getv(from_object, ["config"]))

    return to_object


def _GetAgentEngineSandboxOperationParameters_to_vertex(
    from_object: Union[dict[str, Any], object],
    parent_object: Optional[dict[str, Any]] = None,
) -> dict[str, Any]:
    to_object: dict[str, Any] = {}
    if getv(from_object, ["operation_name"]) is not None:
        setv(
            to_object, ["_url", "operationName"], getv(from_object, ["operation_name"])
        )

    if getv(from_object, ["config"]) is not None:
        setv(to_object, ["config"], getv(from_object, ["config"]))

    return to_object


def _GetAgentEngineSandboxRequestParameters_to_vertex(
    from_object: Union[dict[str, Any], object],
    parent_object: Optional[dict[str, Any]] = None,
) -> dict[str, Any]:
    to_object: dict[str, Any] = {}
    if getv(from_object, ["name"]) is not None:
        setv(to_object, ["_url", "name"], getv(from_object, ["name"]))

    if getv(from_object, ["config"]) is not None:
        setv(to_object, ["config"], getv(from_object, ["config"]))

    return to_object


def _ListAgentEngineSandboxesConfig_to_vertex(
    from_object: Union[dict[str, Any], object],
    parent_object: Optional[dict[str, Any]] = None,
) -> dict[str, Any]:
    to_object: dict[str, Any] = {}

    if getv(from_object, ["page_size"]) is not None:
        setv(parent_object, ["_query", "pageSize"], getv(from_object, ["page_size"]))

    if getv(from_object, ["page_token"]) is not None:
        setv(parent_object, ["_query", "pageToken"], getv(from_object, ["page_token"]))

    if getv(from_object, ["filter"]) is not None:
        setv(parent_object, ["_query", "filter"], getv(from_object, ["filter"]))

    return to_object


def _ListAgentEngineSandboxesRequestParameters_to_vertex(
    from_object: Union[dict[str, Any], object],
    parent_object: Optional[dict[str, Any]] = None,
) -> dict[str, Any]:
    to_object: dict[str, Any] = {}
    if getv(from_object, ["name"]) is not None:
        setv(to_object, ["_url", "name"], getv(from_object, ["name"]))

    if getv(from_object, ["config"]) is not None:
        setv(
            to_object,
            ["config"],
            _ListAgentEngineSandboxesConfig_to_vertex(
                getv(from_object, ["config"]), to_object
            ),
        )

    return to_object


class Sandboxes(_api_module.BaseModule):

    def _create(
        self,
        *,
        name: str,
        spec: Optional[types.SandboxEnvironmentSpecOrDict] = None,
        config: Optional[types.CreateAgentEngineSandboxConfigOrDict] = None,
    ) -> types.AgentEngineSandboxOperation:
        """
        Creates a new sandbox in the Agent Engine.
        """

        parameter_model = types._CreateAgentEngineSandboxRequestParameters(
            name=name,
            spec=spec,
            config=config,
        )

        request_url_dict: Optional[dict[str, str]]
        if not self._api_client.vertexai:
            raise ValueError("This method is only supported in the Vertex AI client.")
        else:
            request_dict = _CreateAgentEngineSandboxRequestParameters_to_vertex(
                parameter_model
            )
            request_url_dict = request_dict.get("_url")
            if request_url_dict:
                path = "{name}/sandboxEnvironments".format_map(request_url_dict)
            else:
                path = "{name}/sandboxEnvironments"

        query_params = request_dict.get("_query")
        if query_params:
            path = f"{path}?{urlencode(query_params)}"
        # TODO: remove the hack that pops config.
        request_dict.pop("config", None)

        http_options: Optional[types.HttpOptions] = None
        if (
            parameter_model.config is not None
            and parameter_model.config.http_options is not None
        ):
            http_options = parameter_model.config.http_options

        request_dict = _common.convert_to_dict(request_dict)
        request_dict = _common.encode_unserializable_types(request_dict)

        response = self._api_client.request("post", path, request_dict, http_options)

        response_dict = {} if not response.body else json.loads(response.body)

        return_value = types.AgentEngineSandboxOperation._from_response(
            response=response_dict, kwargs=parameter_model.model_dump()
        )

        self._api_client._verify_response(return_value)
        return return_value

    def _delete(
        self,
        *,
        name: str,
        config: Optional[types.DeleteAgentEngineSandboxConfigOrDict] = None,
    ) -> types.DeleteAgentEngineSandboxOperation:
        """
        Delete an Agent Engine sandbox.

        Args:
            name (str):
                Required. The name of the Agent Engine sandbox to be deleted. Format:
                `projects/{project}/locations/{location}/reasoningEngines/{resource_id}/sandboxEnvironments/{sandbox}`.

        """

        parameter_model = types._DeleteAgentEngineSandboxRequestParameters(
            name=name,
            config=config,
        )

        request_url_dict: Optional[dict[str, str]]
        if not self._api_client.vertexai:
            raise ValueError("This method is only supported in the Vertex AI client.")
        else:
            request_dict = _DeleteAgentEngineSandboxRequestParameters_to_vertex(
                parameter_model
            )
            request_url_dict = request_dict.get("_url")
            if request_url_dict:
                path = "{name}".format_map(request_url_dict)
            else:
                path = "{name}"

        query_params = request_dict.get("_query")
        if query_params:
            path = f"{path}?{urlencode(query_params)}"
        # TODO: remove the hack that pops config.
        request_dict.pop("config", None)

        http_options: Optional[types.HttpOptions] = None
        if (
            parameter_model.config is not None
            and parameter_model.config.http_options is not None
        ):
            http_options = parameter_model.config.http_options

        request_dict = _common.convert_to_dict(request_dict)
        request_dict = _common.encode_unserializable_types(request_dict)

        response = self._api_client.request("delete", path, request_dict, http_options)

        response_dict = {} if not response.body else json.loads(response.body)

        return_value = types.DeleteAgentEngineSandboxOperation._from_response(
            response=response_dict, kwargs=parameter_model.model_dump()
        )

        self._api_client._verify_response(return_value)
        return return_value

    def _execute_code(
        self,
        *,
        name: str,
        inputs: Optional[list[types.ChunkOrDict]] = None,
        config: Optional[types.ExecuteCodeAgentEngineSandboxConfigOrDict] = None,
    ) -> types.ExecuteSandboxEnvironmentResponse:
        """
        Execute code in an Agent Engine sandbox.
        """

        parameter_model = types._ExecuteCodeAgentEngineSandboxRequestParameters(
            name=name,
            inputs=inputs,
            config=config,
        )

        request_url_dict: Optional[dict[str, str]]
        if not self._api_client.vertexai:
            raise ValueError("This method is only supported in the Vertex AI client.")
        else:
            request_dict = _ExecuteCodeAgentEngineSandboxRequestParameters_to_vertex(
                parameter_model
            )
            request_url_dict = request_dict.get("_url")
            if request_url_dict:
                path = "{name}/:execute".format_map(request_url_dict)
            else:
                path = "{name}/:execute"

        query_params = request_dict.get("_query")
        if query_params:
            path = f"{path}?{urlencode(query_params)}"
        # TODO: remove the hack that pops config.
        request_dict.pop("config", None)

        http_options: Optional[types.HttpOptions] = None
        if (
            parameter_model.config is not None
            and parameter_model.config.http_options is not None
        ):
            http_options = parameter_model.config.http_options

        request_dict = _common.convert_to_dict(request_dict)
        request_dict = _common.encode_unserializable_types(request_dict)

        response = self._api_client.request("post", path, request_dict, http_options)

        response_dict = {} if not response.body else json.loads(response.body)

        return_value = types.ExecuteSandboxEnvironmentResponse._from_response(
            response=response_dict, kwargs=parameter_model.model_dump()
        )

        self._api_client._verify_response(return_value)
        return return_value

    def _get(
        self,
        *,
        name: str,
        config: Optional[types.GetAgentEngineSandboxConfigOrDict] = None,
    ) -> types.SandboxEnvironment:
        """
        Gets an agent engine sandbox.

        Args:
            name (str): Required. A fully-qualified resource name or ID such as
              "projects/123/locations/us-central1/reasoningEngines/456/sandboxEnvironments/789"
              or a shortened name such as "reasoningEngines/456/sandboxEnvironments/789".

        """

        parameter_model = types._GetAgentEngineSandboxRequestParameters(
            name=name,
            config=config,
        )

        request_url_dict: Optional[dict[str, str]]
        if not self._api_client.vertexai:
            raise ValueError("This method is only supported in the Vertex AI client.")
        else:
            request_dict = _GetAgentEngineSandboxRequestParameters_to_vertex(
                parameter_model
            )
            request_url_dict = request_dict.get("_url")
            if request_url_dict:
                path = "{name}".format_map(request_url_dict)
            else:
                path = "{name}"

        query_params = request_dict.get("_query")
        if query_params:
            path = f"{path}?{urlencode(query_params)}"
        # TODO: remove the hack that pops config.
        request_dict.pop("config", None)

        http_options: Optional[types.HttpOptions] = None
        if (
            parameter_model.config is not None
            and parameter_model.config.http_options is not None
        ):
            http_options = parameter_model.config.http_options

        request_dict = _common.convert_to_dict(request_dict)
        request_dict = _common.encode_unserializable_types(request_dict)

        response = self._api_client.request("get", path, request_dict, http_options)

        response_dict = {} if not response.body else json.loads(response.body)

        return_value = types.SandboxEnvironment._from_response(
            response=response_dict, kwargs=parameter_model.model_dump()
        )

        self._api_client._verify_response(return_value)
        return return_value

    def _list(
        self,
        *,
        name: str,
        config: Optional[types.ListAgentEngineSandboxesConfigOrDict] = None,
    ) -> types.ListAgentEngineSandboxesResponse:
        """
        Lists Agent Engine sandboxes.

        Args:
            name (str): Required. The name of the Agent Engine to list sessions for. Format:
                `projects/{project}/locations/{location}/reasoningEngines/{resource_id}`.
            config (ListAgentEngineSandboxesConfig):
                Optional. Additional configurations for listing the Agent Engine sandboxes.

        Returns:
            ListReasoningEnginesSandboxesResponse: The requested Agent Engine sandboxes.

        """

        parameter_model = types._ListAgentEngineSandboxesRequestParameters(
            name=name,
            config=config,
        )

        request_url_dict: Optional[dict[str, str]]
        if not self._api_client.vertexai:
            raise ValueError("This method is only supported in the Vertex AI client.")
        else:
            request_dict = _ListAgentEngineSandboxesRequestParameters_to_vertex(
                parameter_model
            )
            request_url_dict = request_dict.get("_url")
            if request_url_dict:
                path = "{name}/sandboxEnvironments".format_map(request_url_dict)
            else:
                path = "{name}/sandboxEnvironments"

        query_params = request_dict.get("_query")
        if query_params:
            path = f"{path}?{urlencode(query_params)}"
        # TODO: remove the hack that pops config.
        request_dict.pop("config", None)

        http_options: Optional[types.HttpOptions] = None
        if (
            parameter_model.config is not None
            and parameter_model.config.http_options is not None
        ):
            http_options = parameter_model.config.http_options

        request_dict = _common.convert_to_dict(request_dict)
        request_dict = _common.encode_unserializable_types(request_dict)

        response = self._api_client.request("get", path, request_dict, http_options)

        response_dict = {} if not response.body else json.loads(response.body)

        return_value = types.ListAgentEngineSandboxesResponse._from_response(
            response=response_dict, kwargs=parameter_model.model_dump()
        )

        self._api_client._verify_response(return_value)
        return return_value

    def _get_sandbox_operation(
        self,
        *,
        operation_name: str,
        config: Optional[types.GetAgentEngineOperationConfigOrDict] = None,
    ) -> types.AgentEngineSandboxOperation:
        parameter_model = types._GetAgentEngineSandboxOperationParameters(
            operation_name=operation_name,
            config=config,
        )

        request_url_dict: Optional[dict[str, str]]
        if not self._api_client.vertexai:
            raise ValueError("This method is only supported in the Vertex AI client.")
        else:
            request_dict = _GetAgentEngineSandboxOperationParameters_to_vertex(
                parameter_model
            )
            request_url_dict = request_dict.get("_url")
            if request_url_dict:
                path = "{operationName}".format_map(request_url_dict)
            else:
                path = "{operationName}"

        query_params = request_dict.get("_query")
        if query_params:
            path = f"{path}?{urlencode(query_params)}"
        # TODO: remove the hack that pops config.
        request_dict.pop("config", None)

        http_options: Optional[types.HttpOptions] = None
        if (
            parameter_model.config is not None
            and parameter_model.config.http_options is not None
        ):
            http_options = parameter_model.config.http_options

        request_dict = _common.convert_to_dict(request_dict)
        request_dict = _common.encode_unserializable_types(request_dict)

        response = self._api_client.request("get", path, request_dict, http_options)

        response_dict = {} if not response.body else json.loads(response.body)

        return_value = types.AgentEngineSandboxOperation._from_response(
            response=response_dict, kwargs=parameter_model.model_dump()
        )

        self._api_client._verify_response(return_value)
        return return_value

    def create(
        self,
        *,
        name: str,
        poll_interval_seconds: float = 0.1,
        spec: Optional[types.SandboxEnvironmentSpecOrDict] = None,
        config: Optional[types.CreateAgentEngineSandboxConfigOrDict] = None,
    ) -> types.AgentEngineSandboxOperation:
        """Creates a new sandbox in the Agent Engine.

        Args:
            name (str):
                Required. The name of the agent engine to create sandbox for.
                projects/{project}/locations/{location}/reasoningEngines/{resource_id}
            poll_interval_seconds (float):
                Optional. The interval in seconds to poll for sandbox creation
                completion.
            spec (SandboxEnvironmentSpec):
                Optional. The specification for the sandbox to create.
            config (CreateAgentEngineSandboxConfigOrDict):
                Optional. The configuration for the sandbox.

        Returns:
            AgentEngineSandboxOperation: The operation for creating the sandbox.
        """
        if spec:
            computer_use = False
            if isinstance(spec, dict):
                computer_use = spec.get("computer_use_environment") is not None
            elif hasattr(spec, "computer_use_environment"):
                computer_use = True

            if computer_use:
                logging.warning(
                    "The computer_use_environment feature in the sandboxes module is experimental and may change in future versions."
                )
        operation = self._create(
            name=name,
            spec=spec,
            config=config,
        )
        if config is None:
            config = types.CreateAgentEngineSandboxConfig()
        elif isinstance(config, dict):
            config = types.CreateAgentEngineSandboxConfig.model_validate(config)
        if config.wait_for_completion:
            if not operation.done:
                operation = _agent_engines_utils._await_operation(
                    operation_name=operation.name,
                    get_operation_fn=self._get_sandbox_operation,
                    poll_interval_seconds=poll_interval_seconds,
                )
            # We need to make a call to get the sandbox because the operation
            # response might not contain the relevant fields.
            if not operation.response:
                raise ValueError("Error retrieving sandbox.")
            operation.response = self.get(name=operation.response.name)
        return operation

    def list(
        self,
        *,
        name: str,
        config: Optional[types.ListAgentEngineSandboxesConfigOrDict] = None,
    ) -> Iterator[types.SandboxEnvironment]:
        """Lists Agent Engine sandboxes.

        Args:
            name (str):
                Required. The name of the agent engine to list sandboxes for.
                projects/{project}/locations/{location}/reasoningEngines/{resource_id}/SandboxEnvironments/{sandbox_id}
            config (ListAgentEngineSandboxConfig):
                Optional. The configuration for the sandboxes to list.

        Returns:
            Iterable[SandboxEnvironment]: An iterable of agent engine sandboxes.
        """
        return Pager(
            "sandbox_environments",
            functools.partial(self._list, name=name),
            self._list(name=name, config=config),
            config,
        )

    def execute_code(
        self,
        *,
        name: str,
        input_data: dict[str, Any],
        config: Optional[types.ExecuteCodeAgentEngineSandboxConfigOrDict] = None,
    ) -> types.ExecuteSandboxEnvironmentResponse:
        """Executes code in the Agent Engine sandbox.

        Args:
            name (str):
                Required. The name of the agent engine sandbox to run code in.
                projects/{project}/locations/{location}/reasoningEngines/{resource_id}/SandboxEnvironments/{sandbox_id}
            input_data (dict[str, Any]):
                Required. The input to the code to execute.
            config (ExecuteCodeAgentEngineSandboxConfigOrDict):
                Optional. The configuration for the sandboxes to run code in.

        Returns:
            ExecuteSandboxEnvironmentResponse: The response from executing the code.
        """
        input_chunks = []

        if input_data.get("code") is not None:
            code = input_data.get("code", "")
            json_code = json.dumps({"code": code}).encode("utf-8")
            input_chunks.append(
                types.Chunk(
                    mime_type="application/json",
                    data=json_code,
                )
            )

        for file in input_data.get("files", []):
            file_name = file.get("name", "")
            input_chunks.append(
                types.Chunk(
                    mime_type=file.get("mimeType", ""),
                    data=file.get("content", b""),
                    metadata={"attributes": {"file_name": file_name.encode("utf-8")}},
                )
            )

        response = self._execute_code(
            name=name,
            inputs=input_chunks,
            config=config,
        )

        output_chunks = []
        if response.outputs is not None:
            for output in response.outputs:
                if output.mime_type is None:
                    # if mime_type is not available, try to guess the mime_type from the file_name.
                    if (
                        output.metadata is not None
                        and output.metadata.attributes is not None
                    ):
                        file_name = output.metadata.attributes.get(
                            "file_name", b""
                        ).decode("utf-8")
                        mime_type, _ = mimetypes.guess_type(file_name)
                        output.mime_type = mime_type
                output_chunks.append(output)

        response = types.ExecuteSandboxEnvironmentResponse(outputs=output_chunks)

        return response

    def get(
        self,
        *,
        name: str,
        config: Optional[types.GetAgentEngineSandboxConfigOrDict] = None,
    ) -> types.SandboxEnvironment:
        """Gets an agent engine sandbox.
        Args:
          name (str):
              Required. A fully-qualified resource name or ID such as
              projects/{project}/locations/{location}/reasoningEngines/{resource_id}/SandboxEnvironments/{sandbox_id}
              or a shortened name such as "reasoningEngines/{resource_id}/sandboxEnvironments/{sandbox_id}".
          config (GetAgentEngineSandboxConfigOrDict):
              Optional. The configuration for the sandbox to get.

        """
        return self._get(name=name, config=config)

    def delete(
        self,
        *,
        name: str,
        config: Optional[types.DeleteAgentEngineSandboxConfigOrDict] = None,
    ) -> types.DeleteAgentEngineSandboxOperation:
        """Deletes an agent engine sandbox.
        Args:
          name (str):
              Required. A fully-qualified resource name or ID such as
              projects/{project}/locations/{location}/reasoningEngines/{resource_id}/SandboxEnvironments/{sandbox_id}
              or a shortened name such as "reasoningEngines/{resource_id}/sandboxEnvironments/{sandbox_id}".
          config (DeleteAgentEngineSandboxConfigOrDict):
              Optional. The configuration for the sandbox to delete.
        """
        return self._delete(name=name, config=config)

    def generate_access_token(
        self,
        service_account_email: str,
        sandbox_id: str,
        port: str = "8080",
        timeout: int = 3600,
    ) -> str:
        """Signs a JWT with a Google Cloud service account.

        Args:
            service_account_email (str):
                Required. The email of the service account to use for signing.
            sandbox_id (str):
                Required. The resource name of the sandbox to generate a token for.
            port (str):
                Optional. The port to use for the token. Defaults to "8080".
            timeout (int):
                Optional. The timeout in seconds for the token. Defaults to 3600.

        Returns:
            str: The signed JWT.
        """
        client = iam_credentials_v1.IAMCredentialsClient()
        name = f"projects/-/serviceAccounts/{service_account_email}"
        custom_claims = {"port": port, "sandbox_id": sandbox_id}
        payload = {
            "iat": int(time.time()),
            "exp": int(time.time()) + timeout,
            "iss": service_account_email,
            "nonce": secrets.randbelow(1000000000) + 1,
            "aud": "vmaas-proxy-api",  # default audience for sandbox proxy
            **custom_claims,
        }
        request = iam_credentials_v1.SignJwtRequest(
            name=name,
            payload=json.dumps(payload),
        )
        response = client.sign_jwt(request=request)
        return response.signed_jwt  # type: ignore[no-any-return]

    def send_command(
        self,
        *,
        http_method: str,
        access_token: str,
        sandbox_environment: types.SandboxEnvironment,
        path: Optional[str] = None,
        query_params: Optional[dict[str, object]] = None,
        headers: Optional[dict[str, str]] = None,
        request_dict: Optional[dict[str, object]] = None,
    ) -> genai_types.HttpResponse:
        """Sends a command to the sandbox.

        Args:
            http_method (str):
                Required. The HTTP method to use for the command.
            access_token (str):
                Required. The access token to use for authorization.
            sandbox_environment (types.SandboxEnvironment):
                Required. The sandbox environment to send the command to.
            path (str):
                Optional. The path to send the command to.
            query_params (dict[str, object]):
                Optional. The query parameters to include in the command.
            headers (dict[str, str]):
                Optional. The headers to include in the command.
            request_dict (dict[str, object]):
                Optional. The request body to include in the command.

        Returns:
            genai_types.HttpResponse: The response from the sandbox.
        """
        headers = headers or {}
        request_dict = request_dict or {}
        connection_info = sandbox_environment.connection_info
        if not connection_info:
            raise ValueError("Connection info is not available.")
        if connection_info.load_balancer_hostname:
            endpoint = "https://" + connection_info.load_balancer_hostname
        elif connection_info.load_balancer_ip:
            endpoint = "http://" + connection_info.load_balancer_ip
        else:
            raise ValueError("Load balancer hostname or ip is not available.")

        path = path or ""
        if query_params:
            path = f"{path}?{urlencode(query_params)}"
        headers["Authorization"] = f"Bearer {access_token}"
        endpoint = endpoint + path if path.startswith("/") else endpoint + "/" + path
        http_options = genai_types.HttpOptions(headers=headers, base_url=endpoint)
        http_client = genai.Client(vertexai=True, http_options=http_options)
        # Full path is constructed in this function. The passed in path into request
        # function will not be used.
        response = http_client._api_client.request(http_method, path, request_dict)
        return genai_types.HttpResponse(
            headers=response.headers,
            body=response.body,
        )

    def generate_browser_ws_headers(
        self,
        sandbox_environment: types.SandboxEnvironment,
        service_account_email: str,
        timeout: int = 3600,
    ) -> tuple[str, dict[str, str]]:
        """Generates the websocket upgrade headers for the browser.

        Args:
            sandbox_environment (types.SandboxEnvironment):
                Required. The sandbox environment to generate websocket headers for.
            service_account_email (str):
                Required. The email of the service account to use for signing.
            timeout (int):
                Optional. The timeout in seconds for the token. Defaults to 3600.

        Returns:
            tuple[str, dict[str, str]]: A tuple containing the websocket URL and
            the headers for websocket upgrade.
        """
        sandbox_id = sandbox_environment.name
        # port 8080 is the default port for http endpoint.
        http_access_token = self.generate_access_token(
            service_account_email, sandbox_id, "8080", timeout
        )
        response = self.send_command(
            http_method="GET",
            access_token=http_access_token,
            sandbox_environment=sandbox_environment,
            path="/cdp_ws_endpoint",
        )
        if not response:
            raise ValueError("Failed to get the websocket endpoint.")
        body_dict = json.loads(response.body)
        ws_path = body_dict["endpoint"]

        ws_url = "wss://test-us-central1.autopush-sandbox.vertexai.goog"
        if sandbox_environment and sandbox_environment.connection_info:
            connection_info = sandbox_environment.connection_info
            if connection_info.load_balancer_hostname:
                ws_url = "wss://" + connection_info.load_balancer_hostname
            elif connection_info.load_balancer_ip:
                ws_url = "ws://" + connection_info.load_balancer_ip
            else:
                raise ValueError("Load balancer hostname or ip is not available.")
        ws_url = ws_url + "/" + ws_path

        # port 9222 is the default port for the browser websocket endpoint.
        ws_access_token = self.generate_access_token(
            service_account_email, sandbox_id, "9222", timeout
        )

        headers = {}
        headers["Sec-WebSocket-Protocol"] = f"binary, {ws_access_token}"
        return ws_url, headers


class AsyncSandboxes(_api_module.BaseModule):

    async def _create(
        self,
        *,
        name: str,
        spec: Optional[types.SandboxEnvironmentSpecOrDict] = None,
        config: Optional[types.CreateAgentEngineSandboxConfigOrDict] = None,
    ) -> types.AgentEngineSandboxOperation:
        """
        Creates a new sandbox in the Agent Engine.
        """

        parameter_model = types._CreateAgentEngineSandboxRequestParameters(
            name=name,
            spec=spec,
            config=config,
        )

        request_url_dict: Optional[dict[str, str]]
        if not self._api_client.vertexai:
            raise ValueError("This method is only supported in the Vertex AI client.")
        else:
            request_dict = _CreateAgentEngineSandboxRequestParameters_to_vertex(
                parameter_model
            )
            request_url_dict = request_dict.get("_url")
            if request_url_dict:
                path = "{name}/sandboxEnvironments".format_map(request_url_dict)
            else:
                path = "{name}/sandboxEnvironments"

        query_params = request_dict.get("_query")
        if query_params:
            path = f"{path}?{urlencode(query_params)}"
        # TODO: remove the hack that pops config.
        request_dict.pop("config", None)

        http_options: Optional[types.HttpOptions] = None
        if (
            parameter_model.config is not None
            and parameter_model.config.http_options is not None
        ):
            http_options = parameter_model.config.http_options

        request_dict = _common.convert_to_dict(request_dict)
        request_dict = _common.encode_unserializable_types(request_dict)

        response = await self._api_client.async_request(
            "post", path, request_dict, http_options
        )

        response_dict = {} if not response.body else json.loads(response.body)

        return_value = types.AgentEngineSandboxOperation._from_response(
            response=response_dict, kwargs=parameter_model.model_dump()
        )

        self._api_client._verify_response(return_value)
        return return_value

    async def _delete(
        self,
        *,
        name: str,
        config: Optional[types.DeleteAgentEngineSandboxConfigOrDict] = None,
    ) -> types.DeleteAgentEngineSandboxOperation:
        """
        Delete an Agent Engine sandbox.

        Args:
            name (str):
                Required. The name of the Agent Engine sandbox to be deleted. Format:
                `projects/{project}/locations/{location}/reasoningEngines/{resource_id}/sandboxEnvironments/{sandbox}`.

        """

        parameter_model = types._DeleteAgentEngineSandboxRequestParameters(
            name=name,
            config=config,
        )

        request_url_dict: Optional[dict[str, str]]
        if not self._api_client.vertexai:
            raise ValueError("This method is only supported in the Vertex AI client.")
        else:
            request_dict = _DeleteAgentEngineSandboxRequestParameters_to_vertex(
                parameter_model
            )
            request_url_dict = request_dict.get("_url")
            if request_url_dict:
                path = "{name}".format_map(request_url_dict)
            else:
                path = "{name}"

        query_params = request_dict.get("_query")
        if query_params:
            path = f"{path}?{urlencode(query_params)}"
        # TODO: remove the hack that pops config.
        request_dict.pop("config", None)

        http_options: Optional[types.HttpOptions] = None
        if (
            parameter_model.config is not None
            and parameter_model.config.http_options is not None
        ):
            http_options = parameter_model.config.http_options

        request_dict = _common.convert_to_dict(request_dict)
        request_dict = _common.encode_unserializable_types(request_dict)

        response = await self._api_client.async_request(
            "delete", path, request_dict, http_options
        )

        response_dict = {} if not response.body else json.loads(response.body)

        return_value = types.DeleteAgentEngineSandboxOperation._from_response(
            response=response_dict, kwargs=parameter_model.model_dump()
        )

        self._api_client._verify_response(return_value)
        return return_value

    async def _execute_code(
        self,
        *,
        name: str,
        inputs: Optional[list[types.ChunkOrDict]] = None,
        config: Optional[types.ExecuteCodeAgentEngineSandboxConfigOrDict] = None,
    ) -> types.ExecuteSandboxEnvironmentResponse:
        """
        Execute code in an Agent Engine sandbox.
        """

        parameter_model = types._ExecuteCodeAgentEngineSandboxRequestParameters(
            name=name,
            inputs=inputs,
            config=config,
        )

        request_url_dict: Optional[dict[str, str]]
        if not self._api_client.vertexai:
            raise ValueError("This method is only supported in the Vertex AI client.")
        else:
            request_dict = _ExecuteCodeAgentEngineSandboxRequestParameters_to_vertex(
                parameter_model
            )
            request_url_dict = request_dict.get("_url")
            if request_url_dict:
                path = "{name}/:execute".format_map(request_url_dict)
            else:
                path = "{name}/:execute"

        query_params = request_dict.get("_query")
        if query_params:
            path = f"{path}?{urlencode(query_params)}"
        # TODO: remove the hack that pops config.
        request_dict.pop("config", None)

        http_options: Optional[types.HttpOptions] = None
        if (
            parameter_model.config is not None
            and parameter_model.config.http_options is not None
        ):
            http_options = parameter_model.config.http_options

        request_dict = _common.convert_to_dict(request_dict)
        request_dict = _common.encode_unserializable_types(request_dict)

        response = await self._api_client.async_request(
            "post", path, request_dict, http_options
        )

        response_dict = {} if not response.body else json.loads(response.body)

        return_value = types.ExecuteSandboxEnvironmentResponse._from_response(
            response=response_dict, kwargs=parameter_model.model_dump()
        )

        self._api_client._verify_response(return_value)
        return return_value

    async def _get(
        self,
        *,
        name: str,
        config: Optional[types.GetAgentEngineSandboxConfigOrDict] = None,
    ) -> types.SandboxEnvironment:
        """
        Gets an agent engine sandbox.

        Args:
            name (str): Required. A fully-qualified resource name or ID such as
              "projects/123/locations/us-central1/reasoningEngines/456/sandboxEnvironments/789"
              or a shortened name such as "reasoningEngines/456/sandboxEnvironments/789".

        """

        parameter_model = types._GetAgentEngineSandboxRequestParameters(
            name=name,
            config=config,
        )

        request_url_dict: Optional[dict[str, str]]
        if not self._api_client.vertexai:
            raise ValueError("This method is only supported in the Vertex AI client.")
        else:
            request_dict = _GetAgentEngineSandboxRequestParameters_to_vertex(
                parameter_model
            )
            request_url_dict = request_dict.get("_url")
            if request_url_dict:
                path = "{name}".format_map(request_url_dict)
            else:
                path = "{name}"

        query_params = request_dict.get("_query")
        if query_params:
            path = f"{path}?{urlencode(query_params)}"
        # TODO: remove the hack that pops config.
        request_dict.pop("config", None)

        http_options: Optional[types.HttpOptions] = None
        if (
            parameter_model.config is not None
            and parameter_model.config.http_options is not None
        ):
            http_options = parameter_model.config.http_options

        request_dict = _common.convert_to_dict(request_dict)
        request_dict = _common.encode_unserializable_types(request_dict)

        response = await self._api_client.async_request(
            "get", path, request_dict, http_options
        )

        response_dict = {} if not response.body else json.loads(response.body)

        return_value = types.SandboxEnvironment._from_response(
            response=response_dict, kwargs=parameter_model.model_dump()
        )

        self._api_client._verify_response(return_value)
        return return_value

    async def _list(
        self,
        *,
        name: str,
        config: Optional[types.ListAgentEngineSandboxesConfigOrDict] = None,
    ) -> types.ListAgentEngineSandboxesResponse:
        """
        Lists Agent Engine sandboxes.

        Args:
            name (str): Required. The name of the Agent Engine to list sessions for. Format:
                `projects/{project}/locations/{location}/reasoningEngines/{resource_id}`.
            config (ListAgentEngineSandboxesConfig):
                Optional. Additional configurations for listing the Agent Engine sandboxes.

        Returns:
            ListReasoningEnginesSandboxesResponse: The requested Agent Engine sandboxes.

        """

        parameter_model = types._ListAgentEngineSandboxesRequestParameters(
            name=name,
            config=config,
        )

        request_url_dict: Optional[dict[str, str]]
        if not self._api_client.vertexai:
            raise ValueError("This method is only supported in the Vertex AI client.")
        else:
            request_dict = _ListAgentEngineSandboxesRequestParameters_to_vertex(
                parameter_model
            )
            request_url_dict = request_dict.get("_url")
            if request_url_dict:
                path = "{name}/sandboxEnvironments".format_map(request_url_dict)
            else:
                path = "{name}/sandboxEnvironments"

        query_params = request_dict.get("_query")
        if query_params:
            path = f"{path}?{urlencode(query_params)}"
        # TODO: remove the hack that pops config.
        request_dict.pop("config", None)

        http_options: Optional[types.HttpOptions] = None
        if (
            parameter_model.config is not None
            and parameter_model.config.http_options is not None
        ):
            http_options = parameter_model.config.http_options

        request_dict = _common.convert_to_dict(request_dict)
        request_dict = _common.encode_unserializable_types(request_dict)

        response = await self._api_client.async_request(
            "get", path, request_dict, http_options
        )

        response_dict = {} if not response.body else json.loads(response.body)

        return_value = types.ListAgentEngineSandboxesResponse._from_response(
            response=response_dict, kwargs=parameter_model.model_dump()
        )

        self._api_client._verify_response(return_value)
        return return_value

    async def _get_sandbox_operation(
        self,
        *,
        operation_name: str,
        config: Optional[types.GetAgentEngineOperationConfigOrDict] = None,
    ) -> types.AgentEngineSandboxOperation:
        parameter_model = types._GetAgentEngineSandboxOperationParameters(
            operation_name=operation_name,
            config=config,
        )

        request_url_dict: Optional[dict[str, str]]
        if not self._api_client.vertexai:
            raise ValueError("This method is only supported in the Vertex AI client.")
        else:
            request_dict = _GetAgentEngineSandboxOperationParameters_to_vertex(
                parameter_model
            )
            request_url_dict = request_dict.get("_url")
            if request_url_dict:
                path = "{operationName}".format_map(request_url_dict)
            else:
                path = "{operationName}"

        query_params = request_dict.get("_query")
        if query_params:
            path = f"{path}?{urlencode(query_params)}"
        # TODO: remove the hack that pops config.
        request_dict.pop("config", None)

        http_options: Optional[types.HttpOptions] = None
        if (
            parameter_model.config is not None
            and parameter_model.config.http_options is not None
        ):
            http_options = parameter_model.config.http_options

        request_dict = _common.convert_to_dict(request_dict)
        request_dict = _common.encode_unserializable_types(request_dict)

        response = await self._api_client.async_request(
            "get", path, request_dict, http_options
        )

        response_dict = {} if not response.body else json.loads(response.body)

        return_value = types.AgentEngineSandboxOperation._from_response(
            response=response_dict, kwargs=parameter_model.model_dump()
        )

        self._api_client._verify_response(return_value)
        return return_value
