#!/usr/bin/env python3
# Copyright 2020 The Pigweed Authors
#
# 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.
"""Tests creating pw_rpc client."""

import unittest
import tempfile
from pathlib import Path
from typing import List, Tuple

from pw_protobuf_compiler import python_protos
from pw_rpc import client, ids, packets
from pw_status import Status

TEST_PROTO_1 = b"""\
syntax = "proto3";

package pw.call.test1;

message SomeMessage {
  uint32 magic_number = 1;
}

message AnotherMessage {
  enum Result {
    FAILED = 0;
    FAILED_MISERABLY = 1;
    I_DONT_WANT_TO_TALK_ABOUT_IT = 2;
  }

  Result result = 1;
  string payload = 2;
}

service PublicService {
  rpc SomeUnary(SomeMessage) returns (AnotherMessage) {}
  rpc SomeServerStreaming(SomeMessage) returns (stream AnotherMessage) {}
  rpc SomeClientStreaming(stream SomeMessage) returns (AnotherMessage) {}
  rpc SomeBidiStreaming(stream SomeMessage) returns (stream AnotherMessage) {}
}
"""

TEST_PROTO_2 = b"""\
syntax = "proto2";

package pw.call.test2;

message Request {
  optional float magic_number = 1;
}

message Response {
}

service Alpha {
  rpc Unary(Request) returns (Response) {}
}

service Bravo {
  rpc BidiStreaming(stream Request) returns (stream Response) {}
}
"""


class ClientTest(unittest.TestCase):
    """Tests the pw_rpc Python client."""
    def setUp(self):
        self._proto_dir = tempfile.TemporaryDirectory(prefix='proto_test')
        protos = []

        for i, contents in enumerate([TEST_PROTO_1, TEST_PROTO_2]):
            protos.append(Path(self._proto_dir.name, f'test_{i}.proto'))
            protos[-1].write_bytes(contents)

        self._protos = python_protos.Library(
            python_protos.compile_and_import(protos))

        self._impl = client.SimpleSynchronousClient()
        self._client = client.Client.from_modules(
            self._impl, [client.Channel(1, self._handle_request)],
            self._protos.modules())

        self._last_request: packets.RpcPacket = None
        self._next_packets: List[Tuple[bytes, bool]] = []

    def tearDown(self):
        self._proto_dir.cleanup()

    def _enqueue_response(self,
                          channel_id: int,
                          service_id: int,
                          method_id: int,
                          status: Status = Status.OK,
                          response=b'',
                          valid=True):
        if isinstance(response, bytes):
            payload = response
        else:
            payload = response.SerializeToString()

        self._next_packets.append(
            (packets.RpcPacket(channel_id=channel_id,
                               service_id=service_id,
                               method_id=method_id,
                               status=status.value,
                               payload=payload).SerializeToString(), valid))

    def _handle_request(self, data: bytes):
        self._last_request = packets.decode(data)

        self.assertTrue(self._next_packets)
        for packet, valid in self._next_packets:
            self.assertEqual(valid, self._client.process_packet(packet))

    def _last_payload(self, message_type):
        self.assertIsNotNone(self._last_request)
        message = message_type()
        message.ParseFromString(self._last_request.payload)
        return message

    def test_access_service_client_as_attribute_or_index(self):
        self.assertIs(
            self._client.channel(1).call.PublicService,
            self._client.channel(1).call['PublicService'])
        self.assertIs(
            self._client.channel(1).call.PublicService,
            self._client.channel(1).call[ids.calculate('PublicService')])

    def test_access_method_client_as_attribute_or_index(self):
        self.assertIs(
            self._client.channel(1).call.Alpha.Unary,
            self._client.channel(1).call['Alpha']['Unary'])
        self.assertIs(
            self._client.channel(1).call.Alpha.Unary,
            self._client.channel(1).call['Alpha'][ids.calculate('Unary')])

    def test_check_for_presence_of_services(self):
        self.assertIn('PublicService', self._client.channel(1).call)
        self.assertIn(ids.calculate('PublicService'),
                      self._client.channel(1).call)
        self.assertNotIn('NotAService', self._client.channel(1).call)
        self.assertNotIn(-1213, self._client.channel(1).call)

    def test_check_for_presence_of_methods(self):
        self.assertIn('SomeUnary', self._client.channel(1).call.PublicService)
        self.assertIn(ids.calculate('SomeUnary'),
                      self._client.channel(1).call.PublicService)

        self.assertNotIn('Unary', self._client.channel(1).call.PublicService)
        self.assertNotIn(12345, self._client.channel(1).call.PublicService)

    def test_invoke_unary_rpc(self):
        rpcs = self._client.channel(1).call
        method = rpcs.PublicService.SomeUnary.method

        for _ in range(3):
            self._enqueue_response(1, method.service.id, method.id,
                                   Status.ABORTED,
                                   method.response_type(payload='0_o'))

            status, response = rpcs.PublicService.SomeUnary(magic_number=6)

            self.assertEqual(
                6,
                self._last_payload(method.request_type).magic_number)

            self.assertIs(Status.ABORTED, status)
            self.assertEqual('0_o', response.payload)

    def test_ignore_bad_packets_with_pending_rpc(self):
        rpcs = self._client.channel(1).call
        method = rpcs.PublicService.SomeUnary.method
        service_id = method.service.id

        # Unknown channel
        self._enqueue_response(999, service_id, method.id, valid=False)
        # Bad service
        self._enqueue_response(1, 999, method.id, valid=False)
        # Bad method
        self._enqueue_response(1, service_id, 999, valid=False)
        # For RPC not pending (valid=True because the packet is processed)
        self._enqueue_response(1,
                               service_id,
                               rpcs.PublicService.SomeBidiStreaming.method.id,
                               valid=True)

        self._enqueue_response(1, service_id, method.id, valid=True)

        status, response = rpcs.PublicService.SomeUnary(magic_number=6)
        self.assertIs(Status.OK, status)
        self.assertEqual('', response.payload)

    def test_pass_none_if_payload_fails_to_decode(self):
        rpcs = self._client.channel(1).call
        method = rpcs.PublicService.SomeUnary.method

        self._enqueue_response(1,
                               method.service.id,
                               method.id,
                               Status.OK,
                               b'INVALID DATA!!!',
                               valid=True)

        status, response = rpcs.PublicService.SomeUnary(magic_number=6)
        self.assertIs(status, Status.OK)
        self.assertIsNone(response)

    def test_call_method_with_both_message_and_kwargs(self):
        req = self._client.services['Alpha'].methods['Unary'].request_type()

        with self.assertRaisesRegex(TypeError, r'either'):
            self._client.channel(1).call.Alpha.Unary(req, magic_number=1.0)

    def test_call_method_with_wrong_type(self):
        with self.assertRaisesRegex(TypeError, r'pw\.call\.test2\.Request'):
            self._client.channel(1).call.Alpha.Unary('This is str!')

    def test_call_method_with_incorrect_message_type(self):
        msg = self._protos.packages.pw.call.test1.AnotherMessage()
        with self.assertRaisesRegex(TypeError,
                                    r'pw\.call\.test1\.SomeMessage'):
            self._client.channel(1).call.PublicService.SomeUnary(msg)

    def test_process_packet_invalid_proto_data(self):
        self.assertFalse(self._client.process_packet(b'NOT a packet!'))

    def test_process_packet_unrecognized_channel(self):
        self.assertFalse(
            self._client.process_packet(
                packets.encode((123, 456, 789),
                               self._protos.packages.pw.call.test2.Request())))

    def test_process_packet_unrecognized_service(self):
        self.assertFalse(
            self._client.process_packet(
                packets.encode((1, 456, 789),
                               self._protos.packages.pw.call.test2.Request())))

    def test_process_packet_unrecognized_method(self):
        self.assertFalse(
            self._client.process_packet(
                packets.encode((1, next(iter(self._client.services)).id, 789),
                               self._protos.packages.pw.call.test2.Request())))


if __name__ == '__main__':
    unittest.main()