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

from __future__ import annotations

import functools
import hashlib
import inspect
import logging
import os
import random
import shutil
import string
import subprocess
import sys
import tempfile
import textwrap
from typing import List, NamedTuple, Optional, Sequence, TYPE_CHECKING

import requests

if TYPE_CHECKING:
    from bigframes.session import Session

import cloudpickle
import google.api_core.exceptions
import google.api_core.retry
from google.cloud import (
    bigquery,
    bigquery_connection_v1,
    functions_v2,
    resourcemanager_v3,
)
import google.iam.v1
from ibis.backends.bigquery.compiler import compiles
from ibis.backends.bigquery.datatypes import BigQueryType
from ibis.expr.datatypes.core import DataType as IbisDataType
from ibis.expr.datatypes.core import dtype as python_type_to_bigquery_type
import ibis.expr.operations as ops
import ibis.expr.rules as rlz

from bigframes import clients
import bigframes.constants as constants

logger = logging.getLogger(__name__)

# Protocol version 4 is available in python version 3.4 and above
# https://docs.python.org/3/library/pickle.html#data-stream-format
_pickle_protocol_version = 4

# Input and output types supported by BigQuery DataFrames remote functions.
# TODO(shobs): Extend the support to all types supported by BQ remote functions
# https://cloud.google.com/bigquery/docs/remote-functions#limitations
SUPPORTED_IO_PYTHON_TYPES = {bool, float, int, str}
SUPPORTED_IO_BIGQUERY_TYPEKINDS = {
    "BOOLEAN",
    "BOOL",
    "FLOAT",
    "FLOAT64",
    "INT64",
    "INTEGER",
    "STRING",
}


def get_remote_function_locations(bq_location):
    """Get BQ location and cloud functions region given a BQ client."""
    # TODO(shobs, b/274647164): Find the best way to determine default location.
    # For now let's assume that if no BQ location is set in the client then it
    # defaults to US multi region
    bq_location = bq_location.lower() if bq_location else "us"

    # Cloud function should be in the same region as the bigquery remote function
    cloud_function_region = bq_location

    # BigQuery has multi region but cloud functions does not.
    # Any region in the multi region that supports cloud functions should work
    # https://cloud.google.com/functions/docs/locations
    if bq_location == "us":
        cloud_function_region = "us-central1"
    elif bq_location == "eu":
        cloud_function_region = "europe-west1"

    return bq_location, cloud_function_region


def _get_hash(def_, package_requirements=None):
    "Get hash (32 digits alphanumeric) of a function."
    def_repr = cloudpickle.dumps(def_, protocol=_pickle_protocol_version)
    if package_requirements:
        for p in sorted(package_requirements):
            def_repr += p.encode()
    return hashlib.md5(def_repr).hexdigest()


def _run_system_command(command):
    program = subprocess.Popen(
        [command], stdout=subprocess.PIPE, stderr=subprocess.PIPE, shell=True
    )
    stdout, stderr = program.communicate()
    exit_code = program.wait()
    if exit_code:
        raise RuntimeError(
            f"Command: {command}\nOutput: {stdout.decode()}\nError: {stderr.decode()}"
            f"{constants.FEEDBACK_LINK}"
        )


def routine_ref_to_string_for_query(routine_ref: bigquery.RoutineReference) -> str:
    return f"`{routine_ref.project}.{routine_ref.dataset_id}`.{routine_ref.routine_id}"


class IbisSignature(NamedTuple):
    parameter_names: List[str]
    input_types: List[Optional[IbisDataType]]
    output_type: IbisDataType


def get_cloud_function_name(def_, uniq_suffix=None, package_requirements=None):
    "Get a name for the cloud function for the given user defined function."
    cf_name = _get_hash(def_, package_requirements)
    cf_name = f"bigframes-{cf_name}"  # for identification
    if uniq_suffix:
        cf_name = f"{cf_name}-{uniq_suffix}"
    return cf_name


def get_remote_function_name(def_, uniq_suffix=None, package_requirements=None):
    "Get a name for the BQ remote function for the given user defined function."
    bq_rf_name = _get_hash(def_, package_requirements)
    bq_rf_name = f"bigframes_{bq_rf_name}"  # for identification
    if uniq_suffix:
        bq_rf_name = f"{bq_rf_name}_{uniq_suffix}"
    return bq_rf_name


class RemoteFunctionClient:
    # Wait time (in seconds) for an IAM binding to take effect after creation
    _iam_wait_seconds = 120

    def __init__(
        self,
        gcp_project_id,
        cloud_function_region,
        cloud_functions_client,
        bq_location,
        bq_dataset,
        bq_client,
        bq_connection_client,
        bq_connection_id,
        cloud_resource_manager_client,
    ):
        self._gcp_project_id = gcp_project_id
        self._cloud_function_region = cloud_function_region
        self._cloud_functions_client = cloud_functions_client
        self._bq_location = bq_location
        self._bq_dataset = bq_dataset
        self._bq_client = bq_client
        self._bq_connection_id = bq_connection_id
        self._bq_connection_manager = clients.BqConnectionManager(
            bq_connection_client, cloud_resource_manager_client
        )

    def create_bq_remote_function(
        self, input_args, input_types, output_type, endpoint, bq_function_name
    ):
        """Create a BigQuery remote function given the artifacts of a user defined
        function and the http endpoint of a corresponding cloud function."""
        self._bq_connection_manager.create_bq_connection(
            self._gcp_project_id,
            self._bq_location,
            self._bq_connection_id,
            "run.invoker",
        )

        # Create BQ function
        # https://cloud.google.com/bigquery/docs/reference/standard-sql/remote-functions#create_a_remote_function_2
        bq_function_args = []
        bq_function_return_type = BigQueryType.from_ibis(output_type)

        # We are expecting the input type annotations to be 1:1 with the input args
        for idx, name in enumerate(input_args):
            bq_function_args.append(
                f"{name} {BigQueryType.from_ibis(input_types[idx])}"
            )
        create_function_ddl = f"""
            CREATE OR REPLACE FUNCTION `{self._gcp_project_id}.{self._bq_dataset}`.{bq_function_name}({','.join(bq_function_args)})
            RETURNS {bq_function_return_type}
            REMOTE WITH CONNECTION `{self._gcp_project_id}.{self._bq_location}.{self._bq_connection_id}`
            OPTIONS (
              endpoint = "{endpoint}",
              max_batching_rows = 1000
            )"""

        logger.info(f"Creating BQ remote function: {create_function_ddl}")

        # Make sure the dataset exists. I.e. if it doesn't exist, go ahead and
        # create it
        dataset = bigquery.Dataset(
            bigquery.DatasetReference.from_string(
                self._bq_dataset, default_project=self._gcp_project_id
            )
        )
        dataset.location = self._bq_location
        try:
            # This check does not require bigquery.datasets.create IAM
            # permission. So, if the data set already exists, then user can work
            # without having that permission.
            self._bq_client.get_dataset(dataset)
        except google.api_core.exceptions.NotFound:
            # This requires bigquery.datasets.create IAM permission
            self._bq_client.create_dataset(dataset, exists_ok=True)

        # TODO: Use session._start_query() so we get progress bar
        query_job = self._bq_client.query(create_function_ddl)  # Make an API request.
        query_job.result()  # Wait for the job to complete.

        logger.info(f"Created remote function {query_job.ddl_target_routine}")

    def get_cloud_function_fully_qualified_parent(self):
        "Get the fully qualilfied parent for a cloud function."
        return self._cloud_functions_client.common_location_path(
            self._gcp_project_id, self._cloud_function_region
        )

    def get_cloud_function_fully_qualified_name(self, name):
        "Get the fully qualilfied name for a cloud function."
        return self._cloud_functions_client.function_path(
            self._gcp_project_id, self._cloud_function_region, name
        )

    def get_cloud_function_endpoint(self, name):
        """Get the http endpoint of a cloud function if it exists."""
        fully_qualified_name = self.get_cloud_function_fully_qualified_name(name)
        try:
            response = self._cloud_functions_client.get_function(
                name=fully_qualified_name
            )
            return response.service_config.uri
        except google.api_core.exceptions.NotFound:
            pass
        return None

    def generate_udf_code(self, def_, dir):
        """Generate serialized bytecode using cloudpickle given a udf."""
        udf_code_file_name = "udf.py"
        udf_bytecode_file_name = "udf.cloudpickle"

        # original code, only for debugging purpose
        udf_code = textwrap.dedent(inspect.getsource(def_))
        udf_code_file_path = os.path.join(dir, udf_code_file_name)
        with open(udf_code_file_path, "w") as f:
            f.write(udf_code)

        # serialized bytecode
        udf_bytecode_file_path = os.path.join(dir, udf_bytecode_file_name)
        with open(udf_bytecode_file_path, "wb") as f:
            cloudpickle.dump(def_, f, protocol=_pickle_protocol_version)

        return udf_code_file_name, udf_bytecode_file_name

    def generate_cloud_function_main_code(self, def_, dir):
        """Get main.py code for the cloud function for the given user defined function."""

        # Pickle the udf with all its dependencies
        udf_code_file, udf_bytecode_file = self.generate_udf_code(def_, dir)
        handler_func_name = "udf_http"

        # We want to build a cloud function that works for BQ remote functions,
        # where we receive `calls` in json which is a batch of rows from BQ SQL.
        # The number and the order of values in each row is expected to exactly
        # match to the number and order of arguments in the udf , e.g. if the udf is
        #   def foo(x: int, y: str):
        #     ...
        # then the http request body could look like
        # {
        #   ...
        #   "calls" : [
        #     [123, "hello"],
        #     [456, "world"]
        #   ]
        #   ...
        # }
        # https://cloud.google.com/bigquery/docs/reference/standard-sql/remote-functions#input_format
        code_template = textwrap.dedent(
            """\
        import cloudpickle
        import json

        # original udf code is in {udf_code_file}
        # serialized udf code is in {udf_bytecode_file}
        with open("{udf_bytecode_file}", "rb") as f:
          udf = cloudpickle.load(f)

        def {handler_func_name}(request):
          request_json = request.get_json(silent=True)
          calls = request_json["calls"]
          replies = []
          for call in calls:
            reply = udf(*call)
            replies.append(reply)
          return_json = json.dumps({{"replies" : replies}})
          return return_json
        """
        )

        code = code_template.format(
            udf_code_file=udf_code_file,
            udf_bytecode_file=udf_bytecode_file,
            handler_func_name=handler_func_name,
        )

        main_py = os.path.join(dir, "main.py")
        with open(main_py, "w") as f:
            f.write(code)
        logger.debug(f"Wrote {os.path.abspath(main_py)}:\n{open(main_py).read()}")

        return handler_func_name

    def generate_cloud_function_code(self, def_, dir, package_requirements=None):
        """Generate the cloud function code for a given user defined function."""

        # requirements.txt
        requirements = ["cloudpickle >= 2.1.0"]
        if package_requirements:
            requirements.extend(package_requirements)
        requirements = sorted(requirements)
        requirements_txt = os.path.join(dir, "requirements.txt")
        with open(requirements_txt, "w") as f:
            f.write("\n".join(requirements))

        # main.py
        entry_point = self.generate_cloud_function_main_code(def_, dir)
        return entry_point

    def create_cloud_function(self, def_, cf_name, package_requirements=None):
        """Create a cloud function from the given user defined function."""

        # Build and deploy folder structure containing cloud function
        with tempfile.TemporaryDirectory() as dir:
            entry_point = self.generate_cloud_function_code(
                def_, dir, package_requirements
            )
            archive_path = shutil.make_archive(dir, "zip", dir)

            # We are creating cloud function source code from the currently running
            # python version. Use the same version to deploy. This is necessary
            # because cloudpickle serialization done in one python version and
            # deserialization done in another python version doesn't work.
            # TODO(shobs): Figure out how to achieve version compatibility, specially
            # when pickle (internally used by cloudpickle) guarantees that:
            # https://docs.python.org/3/library/pickle.html#:~:text=The%20pickle%20serialization%20format%20is,unique%20breaking%20change%20language%20boundary.
            python_version = "python{}{}".format(
                sys.version_info.major, sys.version_info.minor
            )

            # Determine an upload URL for user code
            upload_url_request = functions_v2.GenerateUploadUrlRequest()
            upload_url_request.parent = self.get_cloud_function_fully_qualified_parent()
            upload_url_response = self._cloud_functions_client.generate_upload_url(
                request=upload_url_request
            )

            # Upload the code to GCS
            with open(archive_path, "rb") as f:
                response = requests.put(
                    upload_url_response.upload_url,
                    data=f,
                    headers={"content-type": "application/zip"},
                )
                if response.status_code != 200:
                    raise RuntimeError(
                        "Failed to upload user code. code={}, reason={}, text={}".format(
                            response.status_code, response.reason, response.text
                        )
                    )

            # Deploy Cloud Function
            create_function_request = functions_v2.CreateFunctionRequest()
            create_function_request.parent = (
                self.get_cloud_function_fully_qualified_parent()
            )
            create_function_request.function_id = cf_name
            function = functions_v2.Function()
            function.name = self.get_cloud_function_fully_qualified_name(cf_name)
            function.build_config = functions_v2.BuildConfig()
            function.build_config.runtime = python_version
            function.build_config.entry_point = entry_point
            function.build_config.source = functions_v2.Source()
            function.build_config.source.storage_source = functions_v2.StorageSource()
            function.build_config.source.storage_source.bucket = (
                upload_url_response.storage_source.bucket
            )
            function.build_config.source.storage_source.object_ = (
                upload_url_response.storage_source.object_
            )
            function.service_config = functions_v2.ServiceConfig()
            function.service_config.available_memory = "1024M"
            function.service_config.timeout_seconds = 600
            create_function_request.function = function

            # Create the cloud function and wait for it to be ready to use
            try:
                operation = self._cloud_functions_client.create_function(
                    request=create_function_request
                )
                operation.result()

                # Cleanup
                os.remove(archive_path)
            except google.api_core.exceptions.AlreadyExists:
                # If a cloud function with the same name already exists, let's
                # update it
                update_function_request = functions_v2.UpdateFunctionRequest()
                update_function_request.function = function
                operation = self._cloud_functions_client.update_function(
                    request=update_function_request
                )
                operation.result()

        # Fetch the endpoint of the just created function
        endpoint = self.get_cloud_function_endpoint(cf_name)
        if not endpoint:
            raise ValueError(
                f"Couldn't fetch the http endpoint. {constants.FEEDBACK_LINK}"
            )

        logger.info(
            f"Successfully created cloud function {cf_name} with uri ({endpoint})"
        )
        return endpoint

    def provision_bq_remote_function(
        self,
        def_,
        input_types,
        output_type,
        reuse,
        name,
        package_requirements,
    ):
        """Provision a BigQuery remote function."""
        # If reuse of any existing function with the same name (indicated by the
        # same hash of its source code) is not intended, then attach a unique
        # suffix to the intended function name to make it unique.
        uniq_suffix = None
        if not reuse:
            uniq_suffix = "".join(
                random.choices(string.ascii_lowercase + string.digits, k=8)
            )

        # Derive the name of the cloud function underlying the intended BQ
        # remote function
        cloud_function_name = get_cloud_function_name(
            def_, uniq_suffix, package_requirements
        )
        cf_endpoint = self.get_cloud_function_endpoint(cloud_function_name)

        # Create the cloud function if it does not exist
        if not cf_endpoint:
            cf_endpoint = self.create_cloud_function(
                def_, cloud_function_name, package_requirements
            )
        else:
            logger.info(f"Cloud function {cloud_function_name} already exists.")

        # Derive the name of the remote function
        remote_function_name = name
        if not remote_function_name:
            remote_function_name = get_remote_function_name(
                def_, uniq_suffix, package_requirements
            )
        rf_endpoint, rf_conn = self.get_remote_function_specs(remote_function_name)

        # Create the BQ remote function in following circumstances:
        # 1. It does not exist
        # 2. It exists but the existing remote function has different
        #    configuration than intended
        if not rf_endpoint or (
            rf_endpoint != cf_endpoint or rf_conn != self._bq_connection_id
        ):
            input_args = inspect.getargs(def_.__code__).args
            if len(input_args) != len(input_types):
                raise ValueError(
                    "Exactly one type should be provided for every input arg."
                )
            self.create_bq_remote_function(
                input_args, input_types, output_type, cf_endpoint, remote_function_name
            )
        else:
            logger.info(f"Remote function {remote_function_name} already exists.")

        return remote_function_name, cloud_function_name

    def get_remote_function_specs(self, remote_function_name):
        """Check whether a remote function already exists for the udf."""
        http_endpoint = None
        bq_connection = None
        routines = self._bq_client.list_routines(
            f"{self._gcp_project_id}.{self._bq_dataset}"
        )
        try:
            for routine in routines:
                if routine.reference.routine_id == remote_function_name:
                    # TODO(shobs): Use first class properties when they are available
                    # https://github.com/googleapis/python-bigquery/issues/1552
                    rf_options = routine._properties.get("remoteFunctionOptions")
                    if rf_options:
                        http_endpoint = rf_options.get("endpoint")
                        bq_connection = rf_options.get("connection")
                        if bq_connection:
                            bq_connection = os.path.basename(bq_connection)
                    break
        except google.api_core.exceptions.NotFound:
            # The dataset might not exist, in which case the http_endpoint doesn't, either.
            # Note: list_routines doesn't make an API request until we iterate on the response object.
            pass
        return (http_endpoint, bq_connection)


def remote_function_node(
    routine_ref: bigquery.RoutineReference, ibis_signature: IbisSignature
):
    """Creates an Ibis node representing a remote function call."""

    fields = {
        name: rlz.ValueOf(None if type_ == "ANY TYPE" else type_)
        for name, type_ in zip(
            ibis_signature.parameter_names, ibis_signature.input_types
        )
    }

    fields["dtype"] = ibis_signature.output_type  # type: ignore
    fields["shape"] = rlz.shape_like("args")

    node = type(routine_ref_to_string_for_query(routine_ref), (ops.ValueOp,), fields)  # type: ignore

    @compiles(node)
    def compile_node(t, op):
        return "{}({})".format(node.__name__, ", ".join(map(t.translate, op.args)))

    def f(*args, **kwargs):
        return node(*args, **kwargs).to_expr()

    f.bigframes_remote_function = str(routine_ref)  # type: ignore

    return f


class UnsupportedTypeError(ValueError):
    def __init__(self, type_, supported_types):
        self.type = type_
        self.supported_types = supported_types


def ibis_type_from_python_type(t: type) -> IbisDataType:
    if t not in SUPPORTED_IO_PYTHON_TYPES:
        raise UnsupportedTypeError(t, SUPPORTED_IO_PYTHON_TYPES)
    return python_type_to_bigquery_type(t)


def ibis_type_from_type_kind(tk: bigquery.StandardSqlTypeNames) -> IbisDataType:
    if tk not in SUPPORTED_IO_BIGQUERY_TYPEKINDS:
        raise UnsupportedTypeError(tk, SUPPORTED_IO_BIGQUERY_TYPEKINDS)
    return BigQueryType.to_ibis(tk)


def ibis_signature_from_python_signature(
    signature: inspect.Signature,
    input_types: Sequence[type],
    output_type: type,
) -> IbisSignature:
    return IbisSignature(
        parameter_names=list(signature.parameters.keys()),
        input_types=[ibis_type_from_python_type(t) for t in input_types],
        output_type=ibis_type_from_python_type(output_type),
    )


class ReturnTypeMissingError(ValueError):
    pass


def ibis_signature_from_routine(routine: bigquery.Routine) -> IbisSignature:
    if not routine.return_type:
        raise ReturnTypeMissingError

    return IbisSignature(
        parameter_names=[arg.name for arg in routine.arguments],
        input_types=[
            ibis_type_from_type_kind(arg.data_type.type_kind) if arg.data_type else None
            for arg in routine.arguments
        ],
        output_type=ibis_type_from_type_kind(routine.return_type.type_kind),
    )


class DatasetMissingError(ValueError):
    pass


def get_routine_reference(
    routine_ref_str: str, bigquery_client: bigquery.Client, session: Optional[Session]
) -> bigquery.RoutineReference:
    try:
        # Handle cases "<project_id>.<dataset_name>.<routine_name>" and
        # "<dataset_name>.<routine_name>".
        return bigquery.RoutineReference.from_string(
            routine_ref_str,
            default_project=bigquery_client.project,
        )
    except ValueError:
        # Handle case of "<routine_name>".
        if not session:
            raise DatasetMissingError

        dataset_ref = bigquery.DatasetReference(
            bigquery_client.project, session._anonymous_dataset.dataset_id
        )
        return dataset_ref.routine(routine_ref_str)


# Inspired by @udf decorator implemented in ibis-bigquery package
# https://github.com/ibis-project/ibis-bigquery/blob/main/ibis_bigquery/udf/__init__.py
# which has moved as @js to the ibis package
# https://github.com/ibis-project/ibis/blob/master/ibis/backends/bigquery/udf/__init__.py
def remote_function(
    input_types: Sequence[type],
    output_type: type,
    session: Optional[Session] = None,
    bigquery_client: Optional[bigquery.Client] = None,
    bigquery_connection_client: Optional[
        bigquery_connection_v1.ConnectionServiceClient
    ] = None,
    cloud_functions_client: Optional[functions_v2.FunctionServiceClient] = None,
    resource_manager_client: Optional[resourcemanager_v3.ProjectsClient] = None,
    dataset: Optional[str] = None,
    bigquery_connection: Optional[str] = None,
    reuse: bool = True,
    name: Optional[str] = None,
    packages: Optional[Sequence[str]] = None,
):
    """Decorator to turn a user defined function into a BigQuery remote function.

    .. deprecated:: 0.0.1
       This is an internal method. Please use :func:`bigframes.pandas.remote_function` instead.

    .. note::
        Please make sure following is setup before using this API:

        1. Have the below APIs enabled for your project:

            * BigQuery Connection API
            * Cloud Functions API
            * Cloud Run API
            * Cloud Build API
            * Artifact Registry API
            * Cloud Resource Manager API

           This can be done from the cloud console (change `PROJECT_ID` to yours):
           https://console.cloud.google.com/apis/enableflow?apiid=bigqueryconnection.googleapis.com,cloudfunctions.googleapis.com,run.googleapis.com,cloudbuild.googleapis.com,artifactregistry.googleapis.com,cloudresourcemanager.googleapis.com&project=PROJECT_ID

           Or from the gcloud CLI:

           `$ gcloud services enable bigqueryconnection.googleapis.com cloudfunctions.googleapis.com run.googleapis.com cloudbuild.googleapis.com artifactregistry.googleapis.com cloudresourcemanager.googleapis.com`

        2. Have following IAM roles enabled for you:

            * BigQuery Data Editor (roles/bigquery.dataEditor)
            * BigQuery Connection Admin (roles/bigquery.connectionAdmin)
            * Cloud Functions Developer (roles/cloudfunctions.developer)
            * Service Account User (roles/iam.serviceAccountUser) on the service account `PROJECT_NUMBER-compute@developer.gserviceaccount.com`
            * Storage Object Viewer (roles/storage.objectViewer)
            * Project IAM Admin (roles/resourcemanager.projectIamAdmin) (Only required if the bigquery connection being used is not pre-created and is created dynamically with user credentials.)

        3. Either the user has setIamPolicy privilege on the project, or a BigQuery connection is pre-created with necessary IAM role set:

            1. To create a connection, follow https://cloud.google.com/bigquery/docs/reference/standard-sql/remote-functions#create_a_connection
            2. To set up IAM, follow https://cloud.google.com/bigquery/docs/reference/standard-sql/remote-functions#grant_permission_on_function

               Alternatively, the IAM could also be setup via the gcloud CLI:

               `$ gcloud projects add-iam-policy-binding PROJECT_ID --member="serviceAccount:CONNECTION_SERVICE_ACCOUNT_ID" --role="roles/run.invoker"`.

    Args:
        input_types list(type):
            List of input data types in the user defined function.
        output_type type:
            Data type of the output in the user defined function.
        session (bigframes.Session, Optional):
            BigQuery DataFrames session to use for getting default project,
            dataset and BigQuery connection.
        bigquery_client (google.cloud.bigquery.Client, Optional):
            Client to use for BigQuery operations. If this param is not provided
            then bigquery client from the session would be used.
        bigquery_connection_client (google.cloud.bigquery_connection_v1.ConnectionServiceClient, Optional):
            Client to use for cloud functions operations. If this param is not
            provided then functions client from the session would be used.
        cloud_functions_client (google.cloud.functions_v2.FunctionServiceClient, Optional):
            Client to use for BigQuery connection operations. If this param is
            not provided then bigquery connection client from the session would
            be used.
        resource_manager_client (google.cloud.resourcemanager_v3.ProjectsClient, Optional):
            Client to use for cloud resource management operations, e.g. for
            getting and setting IAM roles on cloud resources. If this param is
            not provided then resource manager client from the session would be
            used.
        dataset (str, Optional.):
            Dataset in which to create a BigQuery remote function. It should be in
            `<project_id>.<dataset_name>` or `<dataset_name>` format. If this
            parameter is not provided then session dataset id is used.
        bigquery_connection (str, Optional):
            Name of the BigQuery connection in the form of `CONNECTION_ID` or
            `LOCATION.CONNECTION_ID` or `PROJECT_ID.LOCATION.CONNECTION_ID`.
            If this param is not provided then the bigquery connection from the session
            would be used. If it is pre created in the same location as the
            `bigquery_client.location` then it would be used, otherwise it is created
            dynamically using the `bigquery_connection_client` assuming the user has necessary
            priviliges. The PROJECT_ID should be the same as the BigQuery connection project.
        reuse (bool, Optional):
            Reuse the remote function if is already exists.
            `True` by default, which results in reusing an existing remote
            function and corresponding cloud function (if any) that was
            previously created for the same udf.
            Setting it to `False` forces the creation of a unique remote function.
            If the required remote function does not exist then it would be
            created irrespective of this param.
        name (str, Optional):
            Explicit name of the persisted BigQuery remote function. Use it with
            caution, because two users working in the same project and dataset
            could overwrite each other's remote functions if they use the same
            persistent name.
        packages (str[], Optional):
            Explicit name of the external package dependencies. Each dependency
            is added to the `requirements.txt` as is, and can be of the form
            supported in https://pip.pypa.io/en/stable/reference/requirements-file-format/.

    """
    import bigframes.pandas as bpd

    session = session or bpd.get_global_session()

    # A BigQuery client is required to perform BQ operations
    if not bigquery_client:
        bigquery_client = session.bqclient
    if not bigquery_client:
        raise ValueError(
            "A bigquery client must be provided, either directly or via session. "
            f"{constants.FEEDBACK_LINK}"
        )

    # A BigQuery connection client is required to perform BQ connection operations
    if not bigquery_connection_client:
        bigquery_connection_client = session.bqconnectionclient
    if not bigquery_connection_client:
        raise ValueError(
            "A bigquery connection client must be provided, either directly or via session. "
            f"{constants.FEEDBACK_LINK}"
        )

    # A cloud functions client is required to perform cloud functions operations
    if not cloud_functions_client:
        cloud_functions_client = session.cloudfunctionsclient
    if not cloud_functions_client:
        raise ValueError(
            "A cloud functions client must be provided, either directly or via session. "
            f"{constants.FEEDBACK_LINK}"
        )

    # A resource manager client is required to get/set IAM operations
    if not resource_manager_client:
        resource_manager_client = session.resourcemanagerclient
    if not resource_manager_client:
        raise ValueError(
            "A resource manager client must be provided, either directly or via session. "
            f"{constants.FEEDBACK_LINK}"
        )

    # BQ remote function must be persisted, for which we need a dataset
    # https://cloud.google.com/bigquery/docs/reference/standard-sql/remote-functions#:~:text=You%20cannot%20create%20temporary%20remote%20functions.
    if dataset:
        dataset_ref = bigquery.DatasetReference.from_string(
            dataset, default_project=bigquery_client.project
        )
    else:
        dataset_ref = session._anonymous_dataset

    bq_location, cloud_function_region = get_remote_function_locations(
        bigquery_client.location
    )

    # A connection is required for BQ remote function
    # https://cloud.google.com/bigquery/docs/reference/standard-sql/remote-functions#create_a_remote_function
    if not bigquery_connection:
        bigquery_connection = session._bq_connection  # type: ignore

    bigquery_connection = clients.BqConnectionManager.resolve_full_connection_name(
        bigquery_connection,
        default_project=dataset_ref.project,
        default_location=bq_location,
    )
    # Guaranteed to be the form of <project>.<location>.<connection_id>
    (
        gcp_project_id,
        bq_connection_location,
        bq_connection_id,
    ) = bigquery_connection.split(".")
    if gcp_project_id.casefold() != dataset_ref.project.casefold():
        raise ValueError(
            "The project_id does not match BigQuery connection gcp_project_id: "
            f"{dataset_ref.project}."
        )
    if bq_connection_location.casefold() != bq_location.casefold():
        raise ValueError(
            "The location does not match BigQuery connection location: "
            f"{bq_location}."
        )

    def wrapper(f):
        if not callable(f):
            raise TypeError("f must be callable, got {}".format(f))

        signature = inspect.signature(f)
        ibis_signature = ibis_signature_from_python_signature(
            signature, input_types, output_type
        )

        remote_function_client = RemoteFunctionClient(
            dataset_ref.project,
            cloud_function_region,
            cloud_functions_client,
            bq_location,
            dataset_ref.dataset_id,
            bigquery_client,
            bigquery_connection_client,
            bq_connection_id,
            resource_manager_client,
        )

        rf_name, cf_name = remote_function_client.provision_bq_remote_function(
            f,
            ibis_signature.input_types,
            ibis_signature.output_type,
            reuse,
            name,
            packages,
        )

        node = remote_function_node(dataset_ref.routine(rf_name), ibis_signature)

        node = functools.wraps(f)(node)
        node.__signature__ = signature
        node.bigframes_cloud_function = (
            remote_function_client.get_cloud_function_fully_qualified_name(cf_name)
        )

        return node

    return wrapper


def read_gbq_function(
    function_name: str,
    session: Optional[Session] = None,
    bigquery_client: Optional[bigquery.Client] = None,
):
    """
    Read an existing BigQuery function and prepare it for use in future queries.
    """

    # A BigQuery client is required to perform BQ operations
    if not bigquery_client and session:
        bigquery_client = session.bqclient
    if not bigquery_client:
        raise ValueError(
            "A bigquery client must be provided, either directly or via session. "
            f"{constants.FEEDBACK_LINK}"
        )

    try:
        routine_ref = get_routine_reference(function_name, bigquery_client, session)
    except DatasetMissingError:
        raise ValueError(
            "Project and dataset must be provided, either directly or via session. "
            f"{constants.FEEDBACK_LINK}"
        )

    # Find the routine and get its arguments.
    try:
        routine = bigquery_client.get_routine(routine_ref)
    except google.api_core.exceptions.NotFound:
        raise ValueError(f"Unknown function '{routine_ref}'. {constants.FEEDBACK_LINK}")

    try:
        ibis_signature = ibis_signature_from_routine(routine)
    except ReturnTypeMissingError:
        raise ValueError(
            "Function return type must be specified. {constants.FEEDBACK_LINK}"
        )
    except UnsupportedTypeError as e:
        raise ValueError(
            f"Type {e.type} not supported, supported types are {e.supported_types}. "
            f"{constants.FEEDBACK_LINK}"
        )

    return remote_function_node(routine_ref, ibis_signature)