# Copyright 2021 Google LLC
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     https://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

import datetime
import decimal
import uuid

import google.api_core.exceptions
import test_utils.retry

from arrivy.google.cloud._helpers import UTC


_naive = datetime.datetime(2016, 12, 5, 12, 41, 9)
_naive_microseconds = datetime.datetime(2016, 12, 5, 12, 41, 9, 250000)
_stamp = "%s %s" % (_naive.date().isoformat(), _naive.time().isoformat())
_stamp_microseconds = _stamp + ".250000"
_zoned = _naive.replace(tzinfo=UTC)
_zoned_microseconds = _naive_microseconds.replace(tzinfo=UTC)
_numeric = decimal.Decimal("123456789.123456789")


# Examples of most data types to test with query() and DB-API.
STANDARD_SQL_EXAMPLES = [
    ("SELECT 1", 1),
    ("SELECT 1.3", 1.3),
    ("SELECT TRUE", True),
    ('SELECT "ABC"', "ABC"),
    ('SELECT CAST("foo" AS BYTES)', b"foo"),
    ('SELECT TIMESTAMP "%s"' % (_stamp,), _zoned),
    (
        'SELECT TIMESTAMP "%s"' % (_stamp_microseconds,),
        _zoned_microseconds,
    ),
    ('SELECT DATETIME(TIMESTAMP "%s")' % (_stamp,), _naive),
    (
        'SELECT DATETIME(TIMESTAMP "%s")' % (_stamp_microseconds,),
        _naive_microseconds,
    ),
    ('SELECT DATE(TIMESTAMP "%s")' % (_stamp,), _naive.date()),
    ('SELECT TIME(TIMESTAMP "%s")' % (_stamp,), _naive.time()),
    ('SELECT NUMERIC "%s"' % (_numeric,), _numeric),
    ("SELECT (1, 2)", {"_field_1": 1, "_field_2": 2}),
    (
        "SELECT ((1, 2), (3, 4), 5)",
        {
            "_field_1": {"_field_1": 1, "_field_2": 2},
            "_field_2": {"_field_1": 3, "_field_2": 4},
            "_field_3": 5,
        },
    ),
    ("SELECT [1, 2, 3]", [1, 2, 3]),
    (
        "SELECT ([1, 2], 3, [4, 5])",
        {"_field_1": [1, 2], "_field_2": 3, "_field_3": [4, 5]},
    ),
    (
        "SELECT [(1, 2, 3), (4, 5, 6)]",
        [
            {"_field_1": 1, "_field_2": 2, "_field_3": 3},
            {"_field_1": 4, "_field_2": 5, "_field_3": 6},
        ],
    ),
    (
        "SELECT [([1, 2, 3], 4), ([5, 6], 7)]",
        [{"_field_1": [1, 2, 3], "_field_2": 4}, {"_field_1": [5, 6], "_field_2": 7}],
    ),
    ("SELECT ARRAY(SELECT STRUCT([1, 2]))", [{"_field_1": [1, 2]}]),
    ("SELECT ST_GeogPoint(1, 2)", "POINT(1 2)"),
]


def temp_suffix():
    now = datetime.datetime.now()
    return f"{now.strftime('%Y%m%d%H%M%S')}_{uuid.uuid4().hex[:8]}"


def _rate_limit_exceeded(forbidden):
    """Predicate: pass only exceptions with 'rateLimitExceeded' as reason."""
    return any(error["reason"] == "rateLimitExceeded" for error in forbidden._errors)


# We need to wait to stay within the rate limits.
# The alternative outcome is a 403 Forbidden response from upstream, which
# they return instead of the more appropriate 429.
# See https://cloud.google.com/bigquery/quota-policy
retry_403 = test_utils.retry.RetryErrors(
    google.api_core.exceptions.Forbidden,
    error_predicate=_rate_limit_exceeded,
)
