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

#include "internal/platform/wifi_direct.h"

#include <cstddef>
#include <optional>
#include <string>

#include "gtest/gtest.h"
#include "absl/strings/match.h"
#include "absl/strings/string_view.h"
#include "connections/implementation/flags/nearby_connections_feature_flags.h"
#include "internal/flags/nearby_flags.h"
#include "internal/platform/byte_array.h"
#include "internal/platform/cancellation_flag.h"
#include "internal/platform/exception.h"
#include "internal/platform/feature_flags.h"
#include "internal/platform/input_stream.h"
#include "internal/platform/medium_environment.h"
#include "internal/platform/output_stream.h"
#include "internal/platform/single_thread_executor.h"
#include "internal/platform/wifi_credential.h"

namespace nearby {
namespace {

using FeatureFlags = FeatureFlags::Flags;

constexpr FeatureFlags kTestCases[] = {
    FeatureFlags{
        .enable_cancellation_flag = true,
    },
    FeatureFlags{
        .enable_cancellation_flag = false,
    },
};

constexpr absl::string_view kServiceName = "NC-WifiDirectTest";
constexpr absl::string_view kPin = "b592f7d3";
constexpr absl::string_view kIp = "123.234.23.1";
constexpr const size_t kPort = 20;
constexpr absl::string_view kData = "ABCD";
constexpr const size_t kChunkSize = 10;

TEST(WifiDirectCredentialsTest, SetGetServiceName) {
  std::string service_name(kServiceName);
  WifiDirectCredentials wifi_direct_credentials;
  wifi_direct_credentials.SetServiceName(service_name);

  EXPECT_EQ(wifi_direct_credentials.GetServiceName(), kServiceName);
}

TEST(WifiDirectCredentialsTest, SetGetPin) {
  std::string pin(kPin);
  WifiDirectCredentials wifi_direct_credentials;
  wifi_direct_credentials.SetPin(pin);

  EXPECT_EQ(wifi_direct_credentials.GetPin(), kPin);
}


class WifiDirectMediumTest : public testing::TestWithParam<FeatureFlags> {
 protected:
  WifiDirectMediumTest() {
    NearbyFlags::GetInstance().OverrideBoolFlagValue(
        connections::config_package_nearby::nearby_connections_feature::
            kEnableWifiDirect,
        true);
    env_.Stop();
    env_.Start();
  }
  ~WifiDirectMediumTest() override { env_.Stop(); }

  MediumEnvironment& env_{MediumEnvironment::Instance()};
};

INSTANTIATE_TEST_SUITE_P(ParametrisedWifiDirectMediumTest, WifiDirectMediumTest,
                         testing::ValuesIn(kTestCases));

TEST_F(WifiDirectMediumTest, ConstructorDestructorWorks) {
  WifiDirectMedium wifi_direct_a;
  WifiDirectMedium wifi_direct_b;

  // Make sure we can create functional mediums.
  ASSERT_TRUE(wifi_direct_a.IsInterfaceValid());
  ASSERT_TRUE(wifi_direct_b.IsInterfaceValid());

  // Make sure we can create 2 distinct mediums.
  EXPECT_NE(&wifi_direct_a.GetImpl(), &wifi_direct_b.GetImpl());
}

TEST_F(WifiDirectMediumTest, CanStartStopWifiDirect) {
  WifiDirectMedium wifi_direct_a;

  ASSERT_TRUE(wifi_direct_a.IsInterfaceValid());
  EXPECT_TRUE(wifi_direct_a.StartWifiDirect());
  EXPECT_EQ(wifi_direct_a.GetDynamicPortRange(), std::nullopt);
  WifiDirectServerSocket server_socket = wifi_direct_a.ListenForService();
  EXPECT_TRUE(server_socket.IsValid());
  server_socket.Close();
  EXPECT_TRUE(wifi_direct_a.StopWifiDirect());
}

TEST_F(WifiDirectMediumTest, CanConnectDisconnectWifiDirect) {
  WifiDirectMedium wifi_direct_a;
  WifiDirectCredentials credentials;
  credentials.SetServiceName(std::string(kServiceName));
  credentials.SetPin(std::string(kPin));

  ASSERT_TRUE(wifi_direct_a.IsInterfaceValid());
  EXPECT_FALSE(wifi_direct_a.ConnectWifiDirect(credentials));
  EXPECT_TRUE(wifi_direct_a.DisconnectWifiDirect());
}

TEST_P(WifiDirectMediumTest, CanStartDirectGOThatOtherCanConnect) {
  FeatureFlags feature_flags = GetParam();
  env_.SetFeatureFlags(feature_flags);
  WifiDirectMedium wifi_direct_a;
  WifiDirectMedium wifi_direct_b;

  ASSERT_TRUE(wifi_direct_a.IsInterfaceValid());
  ASSERT_TRUE(wifi_direct_b.IsInterfaceValid());
  EXPECT_TRUE(wifi_direct_a.StartWifiDirect());
  WifiDirectCredentials* wifi_direct_credentials =
      wifi_direct_a.GetCredential();
  auto* medium_a =
      env_.GetWifiDirectMedium(wifi_direct_credentials->GetServiceName(), {});
  EXPECT_NE(medium_a, nullptr);
  EXPECT_TRUE(wifi_direct_b.ConnectWifiDirect(*wifi_direct_credentials));

  WifiDirectServerSocket server_socket = wifi_direct_a.ListenForService();
  EXPECT_TRUE(server_socket.IsValid());
  server_socket.PopulateWifiDirectCredentials(*wifi_direct_credentials);
  std::string wifi_direct_a_ip_addr = wifi_direct_credentials->GetGateway();
  EXPECT_FALSE(absl::EndsWith(wifi_direct_a_ip_addr, "."));
  int wifi_direct_a_port = wifi_direct_credentials->GetPort();
  EXPECT_NE(wifi_direct_a_port, 0);

  WifiDirectSocket socket_a;
  WifiDirectSocket socket_b;
  EXPECT_FALSE(socket_a.IsValid());
  EXPECT_FALSE(socket_b.IsValid());

  {
    CancellationFlag flag;
    SingleThreadExecutor server_executor;
    SingleThreadExecutor client_executor;
    client_executor.Execute(
        [&wifi_direct_b, &socket_b, &server_socket, wifi_direct_a_ip_addr,
         wifi_direct_a_port, &flag]() {
          socket_b = wifi_direct_b.ConnectToService(kIp, kPort, &flag);
          EXPECT_FALSE(socket_b.IsValid());
          socket_b = wifi_direct_b.ConnectToService(
              wifi_direct_a_ip_addr, kPort, &flag);
          EXPECT_FALSE(socket_b.IsValid());

          socket_b = wifi_direct_b.ConnectToService(
              wifi_direct_a_ip_addr, wifi_direct_a_port, &flag);
          if (!socket_b.IsValid()) {
            server_socket.Close();
          }
        });
    server_executor.Execute([&socket_a, &server_socket]() {
      socket_a = server_socket.Accept();
      if (!socket_a.IsValid()) {
        server_socket.Close();
      }
    });
  }
  EXPECT_TRUE(socket_a.IsValid());
  EXPECT_TRUE(socket_b.IsValid());
  InputStream& in_stream = socket_a.GetInputStream();
  OutputStream& out_stream = socket_b.GetOutputStream();
  EXPECT_TRUE(out_stream.Write(kData).Ok());
  ExceptionOr<ByteArray> read_data = in_stream.Read(kChunkSize);
  EXPECT_TRUE(read_data.ok());
  EXPECT_EQ(read_data.result().AsStringView(), kData);

  socket_a.Close();
  socket_b.Close();
  EXPECT_FALSE(out_stream.Write(kData).Ok());
  read_data = in_stream.Read(kChunkSize);
  EXPECT_TRUE(read_data.GetResult().Empty());

  server_socket.Close();
  EXPECT_TRUE(wifi_direct_b.DisconnectWifiDirect());
  EXPECT_TRUE(wifi_direct_a.StopWifiDirect());
  auto* medium_b =
      env_.GetWifiDirectMedium(wifi_direct_credentials->GetServiceName(), {});
  EXPECT_EQ(medium_b, nullptr);
}

TEST_P(WifiDirectMediumTest, CanStartDirectGOThatOtherCanCancelConnect) {
  FeatureFlags feature_flags = GetParam();
  env_.SetFeatureFlags(feature_flags);
  WifiDirectMedium wifi_direct_a;
  WifiDirectMedium wifi_direct_b;

  ASSERT_TRUE(wifi_direct_a.IsInterfaceValid());
  ASSERT_TRUE(wifi_direct_b.IsInterfaceValid());
  EXPECT_TRUE(wifi_direct_a.StartWifiDirect());
  WifiDirectCredentials* wifi_direct_credentials =
      wifi_direct_a.GetCredential();
  EXPECT_TRUE(wifi_direct_b.ConnectWifiDirect(*wifi_direct_credentials));

  WifiDirectServerSocket server_socket = wifi_direct_a.ListenForService();
  EXPECT_TRUE(server_socket.IsValid());
  server_socket.PopulateWifiDirectCredentials(*wifi_direct_credentials);
  std::string wifi_direct_a_ip_addr = wifi_direct_credentials->GetGateway();
  int wifi_direct_a_port = wifi_direct_credentials->GetPort();

  WifiDirectSocket socket_a;
  WifiDirectSocket socket_b;
  WifiDirectSocket socket_c;
  EXPECT_FALSE(socket_a.IsValid());
  EXPECT_FALSE(socket_b.IsValid());

  {
    CancellationFlag flag(true);
    SingleThreadExecutor server_executor;
    SingleThreadExecutor client_executor;
    client_executor.Execute(
        [&wifi_direct_b, &socket_b, &server_socket, wifi_direct_a_ip_addr,
         wifi_direct_a_port, &flag]() {
          socket_b = wifi_direct_b.ConnectToService(kIp, kPort, &flag);
          EXPECT_FALSE(socket_b.IsValid());
          socket_b = wifi_direct_b.ConnectToService(
              wifi_direct_a_ip_addr, wifi_direct_a_port, &flag);
          if (!socket_b.IsValid()) {
            server_socket.Close();
          }
        });
    server_executor.Execute([&socket_a, &server_socket]() {
      socket_a = server_socket.Accept();
      if (!socket_a.IsValid()) {
        server_socket.Close();
      }
    });
  }

  if (!feature_flags.enable_cancellation_flag) {
    EXPECT_TRUE(socket_a.IsValid());
    EXPECT_TRUE(socket_b.IsValid());
  } else {
    EXPECT_FALSE(socket_a.IsValid());
    EXPECT_FALSE(socket_b.IsValid());
  }

  server_socket.Close();
  {
    CancellationFlag flag(true);
    socket_c = wifi_direct_b.ConnectToService(wifi_direct_a_ip_addr,
                                              wifi_direct_a_port, &flag);
    EXPECT_FALSE(socket_c.IsValid());
  }

  EXPECT_TRUE(wifi_direct_b.DisconnectWifiDirect());
  EXPECT_TRUE(wifi_direct_a.StopWifiDirect());
}

TEST_F(WifiDirectMediumTest, CanStartDirectGOThatOtherFailConnect) {
  WifiDirectMedium wifi_direct_a;
  WifiDirectMedium wifi_direct_b;

  ASSERT_TRUE(wifi_direct_a.IsInterfaceValid());
  ASSERT_TRUE(wifi_direct_b.IsInterfaceValid());
  EXPECT_TRUE(wifi_direct_a.StartWifiDirect());
  WifiDirectCredentials wifi_direct_credentials;
  wifi_direct_credentials.SetServiceName(std::string(kServiceName));
  wifi_direct_credentials.SetPin(std::string(kPin));

  EXPECT_FALSE(wifi_direct_b.ConnectWifiDirect(wifi_direct_credentials));
  EXPECT_TRUE(wifi_direct_b.DisconnectWifiDirect());

  EXPECT_TRUE(wifi_direct_a.StopWifiDirect());
}

TEST_F(WifiDirectMediumTest, GetSupportedWifiDirectAuthTypes) {
  WifiDirectMedium wifi_direct_a;
  // g3 only supports WifiDirect with auth type of PIN.
  auto supported_types = wifi_direct_a.GetSupportedWifiDirectAuthTypes();
  EXPECT_EQ(supported_types.size(), 1);
  EXPECT_EQ(supported_types[0],
            location::nearby::proto::connections::
                WifiDirectAuthType::WIFI_DIRECT_WITH_PIN);
}

}  // namespace
}  // namespace nearby