// Copyright 2021 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.

#include "gtest/gtest.h"
#include "pw_rpc/internal/hash.h"
#include "pw_rpc/internal/test_utils.h"
#include "pw_rpc/nanopb/test_method_context.h"
#include "pw_rpc_nanopb_private/internal_test_utils.h"
#include "pw_rpc_test_protos/test.rpc.pb.h"

namespace pw::rpc {
namespace test {

class TestService final : public generated::TestService<TestService> {
 public:
  Status TestUnaryRpc(ServerContext&,
                      const pw_rpc_test_TestRequest& request,
                      pw_rpc_test_TestResponse& response) {
    response.value = request.integer + 1;
    return static_cast<Status::Code>(request.status_code);
  }

  void TestAnotherUnaryRpc(
      ServerContext& ctx,
      const pw_rpc_test_TestRequest& request,
      NanopbServerResponder<pw_rpc_test_TestResponse>& responder) {
    pw_rpc_test_TestResponse response{};
    responder.Finish(response, TestUnaryRpc(ctx, request, response));
  }

  static void TestServerStreamRpc(
      ServerContext&,
      const pw_rpc_test_TestRequest& request,
      ServerWriter<pw_rpc_test_TestStreamResponse>& writer) {
    for (int i = 0; i < request.integer; ++i) {
      writer.Write({.chunk = {}, .number = static_cast<uint32_t>(i)});
    }

    writer.Finish(static_cast<Status::Code>(request.status_code));
  }

  void TestClientStreamRpc(
      ServerContext&,
      ServerReader<pw_rpc_test_TestRequest, pw_rpc_test_TestStreamResponse>&
          new_reader) {
    reader = std::move(new_reader);
  }

  void TestBidirectionalStreamRpc(
      ServerContext&,
      ServerReaderWriter<pw_rpc_test_TestRequest,
                         pw_rpc_test_TestStreamResponse>& new_reader_writer) {
    reader_writer = std::move(new_reader_writer);
  }

  ServerReader<pw_rpc_test_TestRequest, pw_rpc_test_TestStreamResponse> reader;
  ServerReaderWriter<pw_rpc_test_TestRequest, pw_rpc_test_TestStreamResponse>
      reader_writer;
};

}  // namespace test

namespace {

using internal::ClientContextForTest;

TEST(NanopbCodegen, CompilesProperly) {
  test::TestService service;
  EXPECT_EQ(service.id(), internal::Hash("pw.rpc.test.TestService"));
  EXPECT_STREQ(service.name(), "TestService");
}

TEST(NanopbCodegen, Server_InvokeUnaryRpc) {
  PW_NANOPB_TEST_METHOD_CONTEXT(test::TestService, TestUnaryRpc) context;

  EXPECT_EQ(OkStatus(),
            context.call({.integer = 123, .status_code = OkStatus().code()}));

  EXPECT_EQ(124, context.response().value);

  EXPECT_EQ(Status::InvalidArgument(),
            context.call({.integer = 999,
                          .status_code = Status::InvalidArgument().code()}));
  EXPECT_EQ(1000, context.response().value);
}

TEST(NanopbCodegen, Server_InvokeAsyncUnaryRpc) {
  PW_NANOPB_TEST_METHOD_CONTEXT(test::TestService, TestAnotherUnaryRpc) context;

  context.call({.integer = 123, .status_code = OkStatus().code()});

  EXPECT_EQ(OkStatus(), context.status());
  EXPECT_EQ(124, context.response().value);

  context.call(
      {.integer = 999, .status_code = Status::InvalidArgument().code()});
  EXPECT_EQ(Status::InvalidArgument(), context.status());
  EXPECT_EQ(1000, context.response().value);
}

TEST(NanopbCodegen, Server_InvokeServerStreamingRpc) {
  PW_NANOPB_TEST_METHOD_CONTEXT(test::TestService, TestServerStreamRpc) context;

  context.call({.integer = 0, .status_code = Status::Aborted().code()});

  EXPECT_EQ(Status::Aborted(), context.status());
  EXPECT_TRUE(context.done());
  EXPECT_TRUE(context.responses().empty());
  EXPECT_EQ(0u, context.total_stream_packets());

  context.call({.integer = 4, .status_code = OkStatus().code()});

  ASSERT_EQ(4u, context.responses().size());
  ASSERT_EQ(4u, context.total_stream_packets());

  for (size_t i = 0; i < context.responses().size(); ++i) {
    EXPECT_EQ(context.responses()[i].number, i);
  }

  EXPECT_EQ(OkStatus().code(), context.status());
}

TEST(NanopbCodegen,
     Server_InvokeServerStreamingRpc_ContextKeepsFixedNumberOfResponses) {
  PW_NANOPB_TEST_METHOD_CONTEXT(test::TestService, TestServerStreamRpc, 3)
  context;

  ASSERT_EQ(3u, context.responses().max_size());

  context.call({.integer = 5, .status_code = Status::NotFound().code()});

  ASSERT_EQ(3u, context.responses().size());
  ASSERT_EQ(5u, context.total_stream_packets());

  EXPECT_EQ(context.responses()[0].number, 0u);
  EXPECT_EQ(context.responses()[1].number, 1u);
  EXPECT_EQ(context.responses()[2].number, 4u);
}

TEST(NanopbCodegen, Server_InvokeServerStreamingRpc_ManualWriting) {
  PW_NANOPB_TEST_METHOD_CONTEXT(test::TestService, TestServerStreamRpc, 3)
  context;

  ASSERT_EQ(3u, context.responses().max_size());

  auto writer = context.writer();

  writer.Write({.chunk = {}, .number = 3});
  writer.Write({.chunk = {}, .number = 6});
  writer.Write({.chunk = {}, .number = 9});

  EXPECT_FALSE(context.done());

  writer.Finish(Status::Cancelled());
  ASSERT_TRUE(context.done());
  EXPECT_EQ(Status::Cancelled(), context.status());

  ASSERT_EQ(3u, context.responses().size());
  ASSERT_EQ(3u, context.total_stream_packets());

  EXPECT_EQ(context.responses()[0].number, 3u);
  EXPECT_EQ(context.responses()[1].number, 6u);
  EXPECT_EQ(context.responses()[2].number, 9u);
}

TEST(NanopbCodegen, Server_InvokeClientStreamingRpc) {
  PW_NANOPB_TEST_METHOD_CONTEXT(test::TestService, TestClientStreamRpc) context;

  context.call();

  pw_rpc_test_TestRequest request = {};
  context.service().reader.set_on_next(
      [&request](const pw_rpc_test_TestRequest& req) { request = req; });

  context.SendClientStream({.integer = -99, .status_code = 10});
  EXPECT_EQ(request.integer, -99);
  EXPECT_EQ(request.status_code, 10u);

  ASSERT_EQ(OkStatus(),
            context.service().reader.Finish({.chunk = {}, .number = 3},
                                            Status::Unimplemented()));
  EXPECT_EQ(Status::Unimplemented(), context.status());
  EXPECT_EQ(context.response().number, 3u);
}

TEST(NanopbCodegen, Server_InvokeBidirectionalStreamingRpc) {
  PW_NANOPB_TEST_METHOD_CONTEXT(test::TestService, TestBidirectionalStreamRpc)
  context;

  context.call();

  pw_rpc_test_TestRequest request = {};
  context.service().reader_writer.set_on_next(
      [&request](const pw_rpc_test_TestRequest& req) { request = req; });

  context.SendClientStream({.integer = -99, .status_code = 10});
  EXPECT_EQ(request.integer, -99);
  EXPECT_EQ(request.status_code, 10u);

  ASSERT_EQ(OkStatus(),
            context.service().reader_writer.Write({.chunk = {}, .number = 2}));
  EXPECT_EQ(context.responses()[0].number, 2u);

  ASSERT_EQ(OkStatus(),
            context.service().reader_writer.Finish(Status::NotFound()));
  EXPECT_EQ(Status::NotFound(), context.status());
}

using TestServiceClient = test::nanopb::TestServiceClient;

TEST(NanopbCodegen, Client_GeneratesCallAliases) {
  static_assert(
      std::is_same_v<TestServiceClient::TestUnaryRpcCall,
                     NanopbClientCall<
                         internal::UnaryCallbacks<pw_rpc_test_TestResponse>>>);
  static_assert(
      std::is_same_v<TestServiceClient::TestServerStreamRpcCall,
                     NanopbClientCall<internal::ServerStreamingCallbacks<
                         pw_rpc_test_TestStreamResponse>>>);
}

TEST(NanopbCodegen, ClientCall_DefaultConstructor) {
  TestServiceClient::TestUnaryRpcCall unary_call;
  TestServiceClient::TestServerStreamRpcCall server_streaming_call;
}

TEST(NanopbCodegen, Client_InvokesUnaryRpcWithCallback) {
  constexpr uint32_t kServiceId = internal::Hash("pw.rpc.test.TestService");
  constexpr uint32_t kMethodId = internal::Hash("TestUnaryRpc");

  ClientContextForTest<128, 128, 99, kServiceId, kMethodId> context;

  TestServiceClient test_client(context.client(), context.channel().id());

  struct {
    Status last_status = Status::Unknown();
    int response_value = -1;
  } result;

  auto call = test_client.TestUnaryRpc(
      {.integer = 123, .status_code = 0},
      [&result](const pw_rpc_test_TestResponse& response, Status status) {
        result.last_status = status;
        result.response_value = response.value;
      });

  EXPECT_TRUE(call.active());

  EXPECT_EQ(context.output().packet_count(), 1u);
  auto packet = context.output().sent_packet();
  EXPECT_EQ(packet.channel_id(), context.channel().id());
  EXPECT_EQ(packet.service_id(), kServiceId);
  EXPECT_EQ(packet.method_id(), kMethodId);
  PW_DECODE_PB(pw_rpc_test_TestRequest, sent_proto, packet.payload());
  EXPECT_EQ(sent_proto.integer, 123);

  PW_ENCODE_PB(pw_rpc_test_TestResponse, response, .value = 42);
  context.SendResponse(OkStatus(), response);
  EXPECT_EQ(result.last_status, OkStatus());
  EXPECT_EQ(result.response_value, 42);
}

TEST(NanopbCodegen, Client_InvokesServerStreamingRpcWithCallback) {
  constexpr uint32_t kServiceId = internal::Hash("pw.rpc.test.TestService");
  constexpr uint32_t kMethodId = internal::Hash("TestServerStreamRpc");

  ClientContextForTest<128, 128, 99, kServiceId, kMethodId> context;

  TestServiceClient test_client(context.client(), context.channel().id());

  struct {
    bool active = true;
    Status stream_status = Status::Unknown();
    int response_value = -1;
  } result;

  auto call = test_client.TestServerStreamRpc(
      {.integer = 123, .status_code = 0},
      [&result](const pw_rpc_test_TestStreamResponse& response) {
        result.active = true;
        result.response_value = response.number;
      },
      [&result](Status status) {
        result.active = false;
        result.stream_status = status;
      });

  EXPECT_TRUE(call.active());

  EXPECT_EQ(context.output().packet_count(), 1u);
  auto packet = context.output().sent_packet();
  EXPECT_EQ(packet.channel_id(), context.channel().id());
  EXPECT_EQ(packet.service_id(), kServiceId);
  EXPECT_EQ(packet.method_id(), kMethodId);
  PW_DECODE_PB(pw_rpc_test_TestRequest, sent_proto, packet.payload());
  EXPECT_EQ(sent_proto.integer, 123);

  PW_ENCODE_PB(
      pw_rpc_test_TestStreamResponse, response, .chunk = {}, .number = 11u);
  context.SendServerStream(response);
  EXPECT_TRUE(result.active);
  EXPECT_EQ(result.response_value, 11);

  context.SendResponse(Status::NotFound());
  EXPECT_FALSE(result.active);
  EXPECT_EQ(result.stream_status, Status::NotFound());
}

TEST(NanopbCodegen, Client_StaticMethod_InvokesUnaryRpcWithCallback) {
  constexpr uint32_t kServiceId = internal::Hash("pw.rpc.test.TestService");
  constexpr uint32_t kMethodId = internal::Hash("TestUnaryRpc");

  ClientContextForTest<128, 128, 99, kServiceId, kMethodId> context;

  struct {
    Status last_status = Status::Unknown();
    int response_value = -1;
  } result;

  auto call = test::pw_rpc::nanopb::TestService::TestUnaryRpc(
      context.client(),
      context.channel().id(),
      {.integer = 123, .status_code = 0},
      [&result](const pw_rpc_test_TestResponse& response, Status status) {
        result.last_status = status;
        result.response_value = response.value;
      });

  EXPECT_TRUE(call.active());

  EXPECT_EQ(context.output().packet_count(), 1u);
  auto packet = context.output().sent_packet();
  EXPECT_EQ(packet.channel_id(), context.channel().id());
  EXPECT_EQ(packet.service_id(), kServiceId);
  EXPECT_EQ(packet.method_id(), kMethodId);
  PW_DECODE_PB(pw_rpc_test_TestRequest, sent_proto, packet.payload());
  EXPECT_EQ(sent_proto.integer, 123);

  PW_ENCODE_PB(pw_rpc_test_TestResponse, response, .value = 42);
  context.SendResponse(OkStatus(), response);
  EXPECT_EQ(result.last_status, OkStatus());
  EXPECT_EQ(result.response_value, 42);
}

TEST(NanopbCodegen, Client_StaticMethod_InvokesServerStreamingRpcWithCallback) {
  constexpr uint32_t kServiceId = internal::Hash("pw.rpc.test.TestService");
  constexpr uint32_t kMethodId = internal::Hash("TestServerStreamRpc");

  ClientContextForTest<128, 128, 99, kServiceId, kMethodId> context;

  struct {
    bool active = true;
    Status stream_status = Status::Unknown();
    int response_value = -1;
  } result;

  auto call = test::pw_rpc::nanopb::TestService::TestServerStreamRpc(
      context.client(),
      context.channel().id(),
      {.integer = 123, .status_code = 0},
      [&result](const pw_rpc_test_TestStreamResponse& response) {
        result.active = true;
        result.response_value = response.number;
      },
      [&result](Status status) {
        result.active = false;
        result.stream_status = status;
      });

  EXPECT_TRUE(call.active());

  EXPECT_EQ(context.output().packet_count(), 1u);
  auto packet = context.output().sent_packet();
  EXPECT_EQ(packet.channel_id(), context.channel().id());
  EXPECT_EQ(packet.service_id(), kServiceId);
  EXPECT_EQ(packet.method_id(), kMethodId);
  PW_DECODE_PB(pw_rpc_test_TestRequest, sent_proto, packet.payload());
  EXPECT_EQ(sent_proto.integer, 123);

  PW_ENCODE_PB(
      pw_rpc_test_TestStreamResponse, response, .chunk = {}, .number = 11u);
  context.SendServerStream(response);
  EXPECT_TRUE(result.active);
  EXPECT_EQ(result.response_value, 11);

  context.SendResponse(Status::NotFound());
  EXPECT_FALSE(result.active);
  EXPECT_EQ(result.stream_status, Status::NotFound());
}

TEST(NanopbCodegen, Client_InvalidChannel_DoesntMakeCall) {
  constexpr uint32_t kServiceId = internal::Hash("pw.rpc.test.TestService");
  constexpr uint32_t kMethodId = internal::Hash("TestUnaryRpc");

  constexpr uint32_t kBadChannelId = 97;

  ClientContextForTest<128, 128, 99, kServiceId, kMethodId> context;
  TestServiceClient test_client(context.client(), kBadChannelId);

  struct {
    Status last_status = Status::Unknown();
    int response_value = -1;
  } result;

  Status error = Status::Unknown();

  auto call = test_client.TestUnaryRpc(
      {.integer = 123, .status_code = 0},
      [&result](const pw_rpc_test_TestResponse& response, Status status) {
        result.last_status = status;
        result.response_value = response.value;
      },
      [&error](Status rpc_error) { error = rpc_error; });

  EXPECT_FALSE(call.active());

  EXPECT_EQ(context.output().packet_count(), 0u);
  EXPECT_EQ(result.last_status, Status::Unknown());
  EXPECT_EQ(result.response_value, -1);
  EXPECT_EQ(error, Status::Unavailable());
}

}  // namespace
}  // namespace pw::rpc