# Copyright 2018, 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 functools
import logging
import threading
import time
import types as stdlib_types

import mock
import pytest

from google.api_core import bidi
from google.api_core import exceptions
from google.cloud.pubsub_v1 import types
from google.cloud.pubsub_v1.subscriber import client
from google.cloud.pubsub_v1.subscriber import message
from google.cloud.pubsub_v1.subscriber import scheduler
from google.cloud.pubsub_v1.subscriber._protocol import dispatcher
from google.cloud.pubsub_v1.subscriber._protocol import heartbeater
from google.cloud.pubsub_v1.subscriber._protocol import leaser
from google.cloud.pubsub_v1.subscriber._protocol import messages_on_hold
from google.cloud.pubsub_v1.subscriber._protocol import requests
from google.cloud.pubsub_v1.subscriber._protocol import streaming_pull_manager
from google.cloud.pubsub_v1.subscriber import exceptions as subscriber_exceptions
from google.cloud.pubsub_v1.subscriber import futures
from google.pubsub_v1 import types as gapic_types
import grpc
from google.rpc import status_pb2
from google.rpc import code_pb2
from google.rpc import error_details_pb2


@pytest.mark.parametrize(
    "exception,expected_cls",
    [
        (ValueError("meep"), ValueError),
        (
            mock.create_autospec(grpc.RpcError, instance=True),
            exceptions.GoogleAPICallError,
        ),
        ({"error": "RPC terminated"}, Exception),
        ("something broke", Exception),
    ],
)
def test__wrap_as_exception(exception, expected_cls):
    assert isinstance(
        streaming_pull_manager._wrap_as_exception(exception), expected_cls
    )


def test__wrap_callback_errors_no_error():
    msg = mock.create_autospec(message.Message, instance=True)
    callback = mock.Mock()
    on_callback_error = mock.Mock()

    streaming_pull_manager._wrap_callback_errors(callback, on_callback_error, msg)

    callback.assert_called_once_with(msg)
    msg.nack.assert_not_called()
    on_callback_error.assert_not_called()


def test__wrap_callback_errors_error():
    callback_error = ValueError("meep")

    msg = mock.create_autospec(message.Message, instance=True)
    callback = mock.Mock(side_effect=callback_error)
    on_callback_error = mock.Mock()

    streaming_pull_manager._wrap_callback_errors(callback, on_callback_error, msg)

    msg.nack.assert_called_once()
    on_callback_error.assert_called_once_with(callback_error)


def test_constructor_and_default_state():
    manager = streaming_pull_manager.StreamingPullManager(
        mock.sentinel.client, mock.sentinel.subscription
    )

    # Public state
    assert manager.is_active is False
    assert manager.flow_control == types.FlowControl()
    assert manager.dispatcher is None
    assert manager.leaser is None
    assert manager.ack_histogram is not None
    assert manager.ack_deadline == 10
    assert manager.load == 0

    # Private state
    assert manager._client == mock.sentinel.client
    assert manager._subscription == mock.sentinel.subscription
    assert manager._scheduler is not None
    assert manager._messages_on_hold is not None
    assert manager._client_id is not None


def test_constructor_with_options():
    manager = streaming_pull_manager.StreamingPullManager(
        mock.sentinel.client,
        mock.sentinel.subscription,
        flow_control=mock.sentinel.flow_control,
        scheduler=mock.sentinel.scheduler,
    )

    assert manager.flow_control == mock.sentinel.flow_control
    assert manager._scheduler == mock.sentinel.scheduler


def make_manager(**kwargs):
    client_ = mock.create_autospec(client.Client, instance=True)
    scheduler_ = mock.create_autospec(scheduler.Scheduler, instance=True)
    return streaming_pull_manager.StreamingPullManager(
        client_, "subscription-name", scheduler=scheduler_, **kwargs
    )


def complete_modify_ack_deadline_calls(dispatcher):
    def complete_futures(*args, **kwargs):
        modack_requests = args[0]
        for req in modack_requests:
            if req.future:
                req.future.set_result(subscriber_exceptions.AcknowledgeStatus.SUCCESS)

    dispatcher.modify_ack_deadline.side_effect = complete_futures


def fake_leaser_add(leaser, init_msg_count=0, assumed_msg_size=10):
    """Add a simplified fake add() method to a leaser instance.

    The fake add() method actually increases the leaser's internal message count
    by one for each message, and the total bytes by ``assumed_msg_size`` for
    each message (regardless of the actual message size).
    """

    def fake_add(self, items):
        self.message_count += len(items)
        self.bytes += len(items) * assumed_msg_size

    leaser.message_count = init_msg_count
    leaser.bytes = init_msg_count * assumed_msg_size
    leaser.add = stdlib_types.MethodType(fake_add, leaser)


def test__obtain_ack_deadline_no_custom_flow_control_setting():
    from google.cloud.pubsub_v1.subscriber._protocol import histogram

    manager = make_manager()

    # Make sure that min_duration_per_lease_extension and
    # max_duration_per_lease_extension is disabled.
    manager._flow_control = types.FlowControl(
        min_duration_per_lease_extension=0, max_duration_per_lease_extension=0
    )

    deadline = manager._obtain_ack_deadline(maybe_update=True)
    assert deadline == histogram.MIN_ACK_DEADLINE

    # When we get some historical data, the deadline is adjusted.
    manager.ack_histogram.add(histogram.MIN_ACK_DEADLINE * 2)
    deadline = manager._obtain_ack_deadline(maybe_update=True)
    assert deadline == histogram.MIN_ACK_DEADLINE * 2

    # Adding just a single additional data point does not yet change the deadline.
    manager.ack_histogram.add(histogram.MIN_ACK_DEADLINE)
    deadline = manager._obtain_ack_deadline(maybe_update=True)
    assert deadline == histogram.MIN_ACK_DEADLINE * 2


def test__obtain_ack_deadline_with_max_duration_per_lease_extension():
    from google.cloud.pubsub_v1.subscriber._protocol import histogram

    manager = make_manager()
    manager._flow_control = types.FlowControl(
        max_duration_per_lease_extension=histogram.MIN_ACK_DEADLINE + 1
    )
    manager.ack_histogram.add(histogram.MIN_ACK_DEADLINE * 3)  # make p99 value large

    # The deadline configured in flow control should prevail.
    deadline = manager._obtain_ack_deadline(maybe_update=True)
    assert deadline == histogram.MIN_ACK_DEADLINE + 1


def test__obtain_ack_deadline_with_min_duration_per_lease_extension():
    from google.cloud.pubsub_v1.subscriber._protocol import histogram

    manager = make_manager()
    manager._flow_control = types.FlowControl(
        min_duration_per_lease_extension=histogram.MAX_ACK_DEADLINE
    )
    manager.ack_histogram.add(histogram.MIN_ACK_DEADLINE)  # make p99 value small

    # The deadline configured in flow control should prevail.
    deadline = manager._obtain_ack_deadline(maybe_update=True)
    assert deadline == histogram.MAX_ACK_DEADLINE


def test__obtain_ack_deadline_with_max_duration_per_lease_extension_too_low():
    from google.cloud.pubsub_v1.subscriber._protocol import histogram

    manager = make_manager()
    manager._flow_control = types.FlowControl(
        max_duration_per_lease_extension=histogram.MIN_ACK_DEADLINE - 1
    )

    # The deadline configured in flow control should be adjusted to the minimum allowed.
    deadline = manager._obtain_ack_deadline(maybe_update=True)
    assert deadline == histogram.MIN_ACK_DEADLINE


def test__obtain_ack_deadline_with_min_duration_per_lease_extension_too_high():
    from google.cloud.pubsub_v1.subscriber._protocol import histogram

    manager = make_manager()
    manager._flow_control = types.FlowControl(
        min_duration_per_lease_extension=histogram.MAX_ACK_DEADLINE + 1
    )

    # The deadline configured in flow control should be adjusted to the maximum allowed.
    deadline = manager._obtain_ack_deadline(maybe_update=True)
    assert deadline == histogram.MAX_ACK_DEADLINE


def test__obtain_ack_deadline_with_exactly_once_enabled():
    manager = make_manager()
    manager._flow_control = types.FlowControl(
        min_duration_per_lease_extension=0  # leave as default value
    )
    manager._exactly_once_enabled = True
    manager.ack_histogram.add(
        10
    )  # reduce p99 value below 60s min for exactly_once subscriptions

    deadline = manager._obtain_ack_deadline(maybe_update=True)
    # Since the 60-second min ack_deadline value for exactly_once subscriptions
    # seconds is higher than the histogram value, the deadline should be 60 sec.
    assert deadline == 60


def test__obtain_ack_deadline_with_min_duration_per_lease_extension_with_exactly_once_enabled():
    from google.cloud.pubsub_v1.subscriber._protocol import histogram

    manager = make_manager()
    manager._flow_control = types.FlowControl(
        min_duration_per_lease_extension=histogram.MAX_ACK_DEADLINE
    )
    manager._exactly_once_enabled = True
    manager.ack_histogram.add(histogram.MIN_ACK_DEADLINE)  # make p99 value small

    # The deadline configured in flow control should prevail.
    deadline = manager._obtain_ack_deadline(maybe_update=True)
    # User-defined custom min ack_deadline value takes precedence over
    # exactly_once default of 60 seconds.
    assert deadline == histogram.MAX_ACK_DEADLINE


def test__obtain_ack_deadline_no_value_update():
    manager = make_manager()

    # Make sure that max_duration_per_lease_extension is disabled.
    manager._flow_control = types.FlowControl(max_duration_per_lease_extension=0)

    manager.ack_histogram.add(21)
    deadline = manager._obtain_ack_deadline(maybe_update=True)
    assert deadline == 21

    for _ in range(5):
        manager.ack_histogram.add(35)  # Gather some new ACK data.

    deadline = manager._obtain_ack_deadline(maybe_update=False)
    assert deadline == 21  # still the same

    # Accessing the value through the ack_deadline property has no side effects either.
    assert manager.ack_deadline == 21

    # Updating the ack deadline is reflected on ack_deadline wrapper, too.
    deadline = manager._obtain_ack_deadline(maybe_update=True)
    assert manager.ack_deadline == deadline == 35


def test_client_id():
    manager1 = make_manager()
    request1 = manager1._get_initial_request(stream_ack_deadline_seconds=10)
    client_id_1 = request1.client_id
    assert client_id_1

    manager2 = make_manager()
    request2 = manager2._get_initial_request(stream_ack_deadline_seconds=10)
    client_id_2 = request2.client_id
    assert client_id_2

    assert client_id_1 != client_id_2


def test_streaming_flow_control():
    manager = make_manager(
        flow_control=types.FlowControl(max_messages=10, max_bytes=1000)
    )
    request = manager._get_initial_request(stream_ack_deadline_seconds=10)
    assert request.max_outstanding_messages == 10
    assert request.max_outstanding_bytes == 1000


def test_streaming_flow_control_use_legacy_flow_control():
    manager = make_manager(
        flow_control=types.FlowControl(max_messages=10, max_bytes=1000),
        use_legacy_flow_control=True,
    )
    request = manager._get_initial_request(stream_ack_deadline_seconds=10)
    assert request.max_outstanding_messages == 0
    assert request.max_outstanding_bytes == 0


def test_maybe_pause_consumer_wo_consumer_set():
    manager = make_manager(
        flow_control=types.FlowControl(max_messages=10, max_bytes=1000)
    )
    manager.maybe_pause_consumer()  # no raise
    # Ensure load > 1
    _leaser = manager._leaser = mock.create_autospec(leaser.Leaser)
    _leaser.message_count = 100
    _leaser.bytes = 10000
    manager.maybe_pause_consumer()  # no raise


def test_lease_load_and_pause():
    manager = make_manager(
        flow_control=types.FlowControl(max_messages=10, max_bytes=1000)
    )
    manager._leaser = leaser.Leaser(manager)