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

import datetime
import unittest

import mock

from google.api_core import grpc_helpers

class _Base(object):
    project = 'PROJECT'

    def _make_one(self, gapic_client=None, handwritten_client=None):
        from google.cloud.trace_v1.gapic import trace_service_client        
        channel = grpc_helpers.ChannelStub()
        if gapic_client is None:
            gapic_client = trace_service_client.TraceServiceClient(channel=channel)
        if handwritten_client is None:
            handwritten_client = mock.Mock()
        api = self._get_target_class()(gapic_client, handwritten_client)
        return channel, api


class Test__TraceAPI(_Base, unittest.TestCase):

    @staticmethod
    def _get_target_class():
        from google.cloud.trace.v1._gapic import _TraceAPI

        return _TraceAPI

    def test_constructor(self):
        from google.cloud.trace_v1.gapic import trace_service_client
        channel = grpc_helpers.ChannelStub()        
        gapic_client = trace_service_client.TraceServiceClient(channel=channel)
        _, api = self._make_one(gapic_client, mock.sentinel.client)
        self.assertIs(api._gapic_api, gapic_client)
        self.assertIs(api.client, mock.sentinel.client)

    def test_patch_traces(self):
        from google.cloud.trace_v1.gapic import trace_service_client
        from google.cloud.trace_v1.proto.trace_pb2 import (
            TraceSpan, Trace, Traces)
        from google.cloud.trace.v1._gapic import _traces_mapping_to_pb
        from google.cloud._helpers import _datetime_to_pb_timestamp

        trace_id = 'test_trace_id'
        span_id = 1234
        span_name = 'test_span_name'
        start_time = datetime.datetime.utcnow()
        end_time = datetime.datetime.utcnow()

        traces = {
            'traces': [
                {
                    'projectId': self.project,
                    'traceId': trace_id,
                    'spans': [
                        {
                            'spanId': span_id,
                            'name': span_name,
                            'startTime': start_time.isoformat() + 'Z',
                            'endTime': end_time.isoformat() + 'Z',
                        },
                    ],
                },
            ],
        }

        traces_pb = _traces_mapping_to_pb(traces)

        gapic_api = mock.Mock(spec=trace_service_client.TraceServiceClient)
        _, api = self._make_one(gapic_api, None)
        api.patch_traces(project_id=self.project, traces=traces)

        gapic_api.patch_traces.assert_called_with(self.project, traces_pb)

        call_args = gapic_api.patch_traces.call_args[0]
        self.assertEqual(len(call_args), 2)
        traces_called = call_args[1]
        self.assertEqual(len(traces_called.traces), 1)
        trace = traces_called.traces[0]

        self.assertEqual(len(trace.spans), 1)
        span = trace.spans[0]

        self.assertIsInstance(traces_called, Traces)
        self.assertEqual(trace.project_id, self.project)
        self.assertEqual(trace.trace_id, trace_id)
        self.assertIsInstance(trace, Trace)

        self.assertEqual(span.span_id, span_id)
        self.assertEqual(span.name, span_name)
        self.assertEqual(
            span.start_time,
            _datetime_to_pb_timestamp(start_time))
        self.assertEqual(
            span.end_time,
            _datetime_to_pb_timestamp(end_time))
        self.assertIsInstance(span, TraceSpan)

    def test_get_trace(self):
        from google.cloud.trace_v1.gapic import trace_service_client

        trace_id = 'test_trace_id'

        gapic_api = mock.Mock(spec=trace_service_client.TraceServiceClient)
        _, api = self._make_one(gapic_api, None)
        patch = mock.patch('google.cloud.trace.v1._gapic._parse_trace_pb',
                           return_value='fake_pb_result')

        with patch:
            api.get_trace(project_id=self.project, trace_id=trace_id)

        gapic_api.get_trace.assert_called_with(self.project, trace_id)

    def _make_trace_pb(
            self,
            project,
            trace_id,
            span_id,
            span_name,
            start_time,
            end_time,
            parent_span_id,
            labels):
        from google.cloud.trace.v1._gapic import _traces_mapping_to_pb

        span_kind = 2

        traces = {
            'traces': [
                {
                    'projectId': project,
                    'traceId': trace_id,
                    'spans': [
                        {
                            'spanId': span_id,
                            'name': span_name,
                            'startTime': start_time,
                            'endTime': end_time,
                            'kind': span_kind,
                            'parentSpanId': parent_span_id,
                            'labels': labels,
                        },
                    ],
                },
            ],
        }

        traces_pb = _traces_mapping_to_pb(traces)
        trace_pb = traces_pb.traces
        return trace_pb

    def test_list_traces(self):
        from google.api_core.page_iterator import GRPCIterator
        from google.cloud._helpers import _rfc3339_to_datetime
        from google.cloud._helpers import UTC
        from google.cloud.trace_v1.gapic import trace_service_client
        from google.cloud.trace_v1.gapic.enums import ListTracesRequest as Enum
        from google.cloud.trace_v1.proto import trace_pb2

        trace_id = 'test_trace_id'
        span_id = 1234
        span_name = 'test_span_name'
        span_kind = 'RPC_CLIENT'
        parent_span_id = 123
        start_ts = datetime.datetime.utcnow()
        end_ts = datetime.datetime.utcnow()
        labels = {
            '/http/status_code': '200',
            '/component': 'HTTP load balancer',
        }
        size = 10
        view_type = Enum.ViewType.COMPLETE
        token = 'TOKEN'

        trace_pb = self._make_trace_pb(
            self.project,
            trace_id,
            span_id,
            span_name,
            start_ts.isoformat() + 'Z',
            end_ts.isoformat() + 'Z',
            parent_span_id,
            labels)

        gapic_api = mock.Mock(spec=trace_service_client.TraceServiceClient)
        gapic_api.list_traces = mock.create_autospec(gapic_api.list_traces)        
        channel, api = self._make_one()

        channel.ListTraces.response = trace_pb2.ListTracesResponse(traces=[trace_pb[0]])
        iterator = api.list_traces(
            project_id=self.project,
            view=view_type,
            page_size=size,
            page_token=token)

        traces = list(iterator)

        self.assertEqual(len(traces), 1)
        trace = traces[0]

        self.assertEqual(len(trace['spans']), 1)
        span = trace['spans'][0]

        self.assertEqual(trace['projectId'], self.project)
        self.assertEqual(trace['traceId'], trace_id)

        self.assertEqual(span['spanId'], str(span_id))
        self.assertEqual(span['name'], span_name)

        self.assertEqual(
            _rfc3339_to_datetime(span['startTime']),
            start_ts.replace(tzinfo=UTC))
        self.assertEqual(
            _rfc3339_to_datetime(span['endTime']),
            end_ts.replace(tzinfo=UTC))
        self.assertEqual(span['kind'], span_kind)
        self.assertEqual(span['parentSpanId'], str(parent_span_id))
        self.assertEqual(span['labels'], labels)

        self.assertEqual(len(channel.ListTraces.requests), 1)
        request = channel.ListTraces.requests[0]

        self.assertEqual(request.project_id, self.project)
        self.assertEqual(request.view, view_type)
        self.assertEqual(request.page_size, size)
        self.assertEqual(request.start_time.ToDatetime(), datetime.datetime(1970, 1, 1, 0, 0))
        self.assertEqual(request.end_time.ToDatetime(), datetime.datetime(1970, 1, 1, 0, 0))
        self.assertEqual(request.filter, '')
        self.assertEqual(request.order_by, '')


class Test__parse_trace_pb(unittest.TestCase):

    @staticmethod
    def _call_fut(*args, **kwargs):
        from google.cloud.trace.v1._gapic import _parse_trace_pb

        return _parse_trace_pb(*args, **kwargs)

    def test_registered_type(self):
        from google.cloud.trace_v1.proto.trace_pb2 import (
            TraceSpan, Trace)
        from google.protobuf.timestamp_pb2 import Timestamp

        project = u'PROJECT'
        trace_id = u'test_trace_id'
        span_id = 1234
        span_name = u'test_span_name'
        start_time = '2017-06-24T00:12:50.369990Z'
        end_time = '2017-06-24T00:13:39.633255Z'
        start_seconds = 1498263170
        start_nanos = 369990000
        end_seconds = 1498263219
        end_nanos = 633255000

        start_time_pb = Timestamp(seconds=start_seconds, nanos=start_nanos)
        end_time_pb = Timestamp(seconds=end_seconds, nanos=end_nanos)

        span_pb = TraceSpan(
            span_id=span_id,
            name=span_name,
            start_time=start_time_pb,
            end_time=end_time_pb)

        trace_pb = Trace(
            project_id=project,
            trace_id=trace_id,
            spans=[span_pb])

        parse_result = self._call_fut(trace_pb)

        expected_result = {
            'projectId': project,
            'traceId': trace_id,
            'spans': [
                {
                    'spanId': str(span_id),
                    'name': span_name,
                    'startTime': start_time,
                    'endTime': end_time,
                },
            ],
        }

        self.assertEqual(parse_result, expected_result)

    @mock.patch('google.cloud.trace.v1._gapic.MessageToDict',
                side_effect=TypeError)
    def test_unregistered_type(self, msg_to_dict_mock):
        trace_pb = mock.Mock(spec=['HasField'])
        trace_pb.HasField.return_value = False
        with self.assertRaises(TypeError):
            self._call_fut(trace_pb)


class Test_make_trace_api(unittest.TestCase):

    def _call_fut(self, client):
        from google.cloud.trace.v1._gapic import make_trace_api

        return make_trace_api(client)

    def test_it(self):
        from google.cloud.trace.v1._gapic import _TraceAPI
        from google.cloud._http import DEFAULT_USER_AGENT

        credentials = object()
        client = mock.Mock(_credentials=credentials, spec=['_credentials'])
        generated_api_kwargs = []
        generated = object()

        def generated_api(**kwargs):
            generated_api_kwargs.append(kwargs)
            return generated

        host = 'foo.apis.invalid'
        generated_api.SERVICE_ADDRESS = host

        patch_api = mock.patch(
            'google.cloud.trace.v1._gapic.trace_service_client.TraceServiceClient',
            new=generated_api)

        with patch_api:
            trace_api = self._call_fut(client)

        self.assertEqual(len(generated_api_kwargs), 1)

        self.assertIsInstance(trace_api, _TraceAPI)
        self.assertIs(trace_api._gapic_api, generated)
        self.assertIs(trace_api.client, client)
