# Copyright 2024 Google LLC All rights reserved.
#
# 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 google.api_core import grpc_helpers, grpc_helpers_async
import google.auth.credentials
import grpc

from google.cloud.aio._cross_sync import CrossSync
from google.cloud.spanner_admin_database_v1 import DatabaseDialect
from google.cloud.spanner_v1._async.database import Database
from google.cloud.spanner_v1.database import SPANNER_DATA_SCOPE
from google.cloud.spanner_v1.services.spanner.transports import (
    SpannerGrpcTransport,
    SpannerTransport,
)

if CrossSync.is_async:
    from google.cloud.spanner_v1._async.testing.interceptors import (
        MethodAbortAsyncInterceptor as MethodAbortInterceptor,
    )
    from google.cloud.spanner_v1._async.testing.interceptors import (
        MethodCountAsyncInterceptor as MethodCountInterceptor,
    )
    from google.cloud.spanner_v1._async.testing.interceptors import (
        XGoogRequestIDHeaderAsyncInterceptor as XGoogRequestIDHeaderInterceptor,
    )
    from google.cloud.spanner_v1.services.spanner.async_client import (
        SpannerAsyncClient as SpannerClient,
    )
else:
    from google.cloud.spanner_v1.services.spanner.client import SpannerClient
    from google.cloud.spanner_v1.testing.interceptors import (
        MethodAbortInterceptor,
        MethodCountInterceptor,
        XGoogRequestIDHeaderInterceptor,
    )

__CROSS_SYNC_OUTPUT__ = "google.cloud.spanner_v1.testing.database_test"


class TestDatabase(Database):
    """Representation of a Cloud Spanner Database. This class is only used for
    system testing as there is no support for interceptors in grpc client
    currently, and we don't want to make changes in the Database class for
    testing purpose as this is a hack to use interceptors in tests."""

    _interceptors = []

    def __init__(
        self,
        database_id,
        instance,
        ddl_statements=(),
        pool=None,
        logger=None,
        encryption_config=None,
        database_dialect=DatabaseDialect.DATABASE_DIALECT_UNSPECIFIED,
        database_role=None,
        enable_drop_protection=False,
    ):
        super().__init__(
            database_id,
            instance,
            ddl_statements,
            pool,
            logger,
            encryption_config,
            database_dialect,
            database_role,
            enable_drop_protection,
        )

        self._method_count_interceptor = MethodCountInterceptor()
        self._method_abort_interceptor = MethodAbortInterceptor()
        self._interceptors = [
            self._method_count_interceptor,
            self._method_abort_interceptor,
        ]

    @property
    def spanner_api(self):
        """Helper for session-related API calls."""
        if self._spanner_api is None:
            client = self._instance._client
            client_info = client._client_info
            client_options = client._client_options
            if self._instance.emulator_host is not None:
                if CrossSync.is_async:
                    self._x_goog_request_id_interceptor = (
                        XGoogRequestIDHeaderInterceptor()
                    )
                    self._interceptors.append(self._x_goog_request_id_interceptor)
                    channel = grpc.aio.insecure_channel(
                        self._instance.emulator_host, interceptors=self._interceptors
                    )
                else:
                    channel = grpc.insecure_channel(self._instance.emulator_host)
                    self._x_goog_request_id_interceptor = (
                        XGoogRequestIDHeaderInterceptor()
                    )
                    self._interceptors.append(self._x_goog_request_id_interceptor)
                    channel = grpc.intercept_channel(channel, *self._interceptors)

                transport = SpannerGrpcTransport(channel=channel)
                self._spanner_api = SpannerClient(
                    client_info=client_info,
                    transport=transport,
                )
                return self._spanner_api
            if self._instance.experimental_host is not None:
                self._x_goog_request_id_interceptor = XGoogRequestIDHeaderInterceptor()
                self._interceptors.append(self._x_goog_request_id_interceptor)

                from google.cloud.spanner_v1._async._helpers import (
                    _create_experimental_host_transport as _create_experimental_host_transport_async,
                )
                from google.cloud.spanner_v1._helpers import (
                    _create_experimental_host_transport as _create_experimental_host_transport_sync,
                )

                if CrossSync.is_async:
                    transport = _create_experimental_host_transport_async(
                        SpannerGrpcTransport,
                        self._instance.experimental_host,
                        client._use_plain_text,
                        client._ca_certificate,
                        client._client_certificate,
                        client._client_key,
                        self._interceptors,
                    )
                else:
                    transport = _create_experimental_host_transport_sync(
                        SpannerGrpcTransport,
                        self._instance.experimental_host,
                        client._use_plain_text,
                        client._ca_certificate,
                        client._client_certificate,
                        client._client_key,
                        self._interceptors,
                    )
                self._spanner_api = SpannerClient(
                    client_info=client_info,
                    transport=transport,
                    client_options=client_options,
                )
                return self._spanner_api
            credentials = client.credentials
            if isinstance(credentials, google.auth.credentials.Scoped):
                credentials = credentials.with_scopes((SPANNER_DATA_SCOPE,))
            self._spanner_api = self._create_spanner_client_for_tests(
                client_options,
                credentials,
            )
        return self._spanner_api

    def _create_spanner_client_for_tests(self, client_options, credentials):
        (
            api_endpoint,
            client_cert_source_func,
        ) = SpannerClient.get_mtls_endpoint_and_cert_source(client_options)

        if CrossSync.is_async:
            channel = grpc_helpers_async.create_channel(
                api_endpoint,
                credentials=credentials,
                credentials_file=client_options.credentials_file,
                quota_project_id=client_options.quota_project_id,
                default_scopes=SpannerTransport.AUTH_SCOPES,
                scopes=client_options.scopes,
                default_host=SpannerTransport.DEFAULT_HOST,
                interceptors=self._interceptors,
            )
        else:
            channel = grpc_helpers.create_channel(
                api_endpoint,
                credentials=credentials,
                credentials_file=client_options.credentials_file,
                quota_project_id=client_options.quota_project_id,
                default_scopes=SpannerTransport.AUTH_SCOPES,
                scopes=client_options.scopes,
                default_host=SpannerTransport.DEFAULT_HOST,
            )
            channel = grpc.intercept_channel(channel, *self._interceptors)

        transport = SpannerGrpcTransport(channel=channel)
        return SpannerClient(
            client_options=client_options,
            transport=transport,
        )

    def reset(self):
        if (
            hasattr(self, "_x_goog_request_id_interceptor")
            and self._x_goog_request_id_interceptor
        ):
            self._x_goog_request_id_interceptor.reset()
