#include "dragnn/core/compute_session_pool.h"

#include <memory>

#include <gmock/gmock.h>

#include "dragnn/core/compute_session.h"
#include "dragnn/core/test/generic.h"
#include "dragnn/core/test/mock_component.h"
#include "dragnn/core/test/mock_compute_session.h"
#include "tensorflow/core/platform/env.h"
#include "tensorflow/core/platform/test.h"

namespace syntaxnet {
namespace dragnn {

using syntaxnet::test::EqualsProto;
using testing::Return;
using testing::Invoke;
using testing::MockFunction;

class ComputeSessionPoolTestPoolAccessor {
 public:
  static void SetComponentBuilder(
      ComputeSessionPool *pool,
      std::function<std::unique_ptr<Component>(const string &component_name,
                                               const string &backend_type)>
          component_builder_function) {
    pool->SetComponentBuilder(std::move(component_builder_function));
  }

  static void SetSessionBuilder(ComputeSessionPool *pool,
                                std::function<std::unique_ptr<ComputeSession>()>
                                    session_builder_function) {
    pool->SetComputeSessionBuilder(std::move(session_builder_function));
  }
};

TEST(ComputeSessionPoolTest, DefaultConstructorWorks) {
  MasterSpec spec;
  GridPoint hyperparams;
  ComputeSessionPool pool(spec, hyperparams);
  auto request = pool.GetSession();
  EXPECT_NE(request, nullptr);
}

TEST(ComputeSessionPoolTest, ComponentBuilderInjectionWorks) {
  MasterSpec spec;
  auto component = spec.add_component();
  component->set_name("test_component_name");
  auto backend = component->mutable_backend();
  backend->set_registered_name("arbitrary_component");
  GridPoint hyperparams;

  ComputeSessionPool pool(spec, hyperparams);

  // Set up a mock component builder.
  MockFunction<std::unique_ptr<Component>(const string &component_name,
                                          const string &backend_type)>
      mock_component_builder;
  auto mock_creation_function = [](string, string) {
    return std::unique_ptr<MockComponent>(new MockComponent());
  };
  EXPECT_CALL(mock_component_builder,
              Call("test_component_name", "arbitrary_component"))
      .WillOnce(Invoke(mock_creation_function));
  ComputeSessionPoolTestPoolAccessor::SetComponentBuilder(
      &pool, mock_component_builder.AsStdFunction());

  // Now, when the session is requested, the mock component builder should see
  // the expected call.
  auto request = pool.GetSession();
  EXPECT_NE(request, nullptr);
}

TEST(ComputeSessionPoolTest, CreatesNewSessionIfNoSessionsExist) {
  // We don't need to fill these for this test.
  MasterSpec spec;
  GridPoint hyperparams;
  ComputeSessionPool pool(spec, hyperparams);

  // Create a function that will track calls to the session builder.
  MockFunction<std::unique_ptr<ComputeSession>()> mock_session_builder;

  // Initialize expectations for a request for a ComputeSession.
  std::unique_ptr<MockComputeSession> session_one(new MockComputeSession());
  MockComputeSession *session_one_ptr = session_one.get();
  auto mock_creation_function = [&session_one]() {
    return std::move(session_one);
  };
  EXPECT_CALL(mock_session_builder, Call())
      .WillOnce(Invoke(mock_creation_function))
      .RetiresOnSaturation();
  EXPECT_CALL(*session_one_ptr,
              Init(EqualsProto(spec), EqualsProto(hyperparams)));

  // Initialize expectations for another request for a ComputeSession.
  std::unique_ptr<MockComputeSession> session_two(new MockComputeSession());
  MockComputeSession *session_two_ptr = session_two.get();
  auto mock_creation_function_two = [&session_two]() {
    return std::move(session_two);
  };
  EXPECT_CALL(mock_session_builder, Call())
      .WillOnce(Invoke(mock_creation_function_two))
      .RetiresOnSaturation();
  EXPECT_CALL(*session_two_ptr,
              Init(EqualsProto(spec), EqualsProto(hyperparams)));

  // Inject the function to the pool.
  ComputeSessionPoolTestPoolAccessor::SetSessionBuilder(
      &pool, mock_session_builder.AsStdFunction());

  // The first call will recieve the second session because of how the mocks go.
  auto first_request = pool.GetSession();
  EXPECT_EQ(first_request.get(), session_two_ptr);

  auto second_request = pool.GetSession();
  EXPECT_EQ(second_request.get(), session_one_ptr);
}

TEST(ComputeSessionPoolTest, ReusesAvailableSessions) {
  // We don't need to fill these for this test.
  MasterSpec spec;
  GridPoint hyperparams;
  ComputeSessionPool pool(spec, hyperparams);

  // Create a function that will track calls to the session builder.
  MockFunction<std::unique_ptr<ComputeSession>()> mock_session_builder;

  // Initialize expectations for a request for a ComputeSession.
  std::unique_ptr<MockComputeSession> session_one(new MockComputeSession());
  MockComputeSession *session_one_ptr = session_one.get();
  auto mock_creation_function = [&session_one]() {
    return std::move(session_one);
  };
  EXPECT_CALL(mock_session_builder, Call())
      .WillOnce(Invoke(mock_creation_function))
      .RetiresOnSaturation();
  EXPECT_CALL(*session_one_ptr,
              Init(EqualsProto(spec), EqualsProto(hyperparams)));

  // Initialize expectations for another request for a ComputeSession.
  std::unique_ptr<MockComputeSession> session_two(new MockComputeSession());
  MockComputeSession *session_two_ptr = session_two.get();
  auto mock_creation_function_two = [&session_two]() {
    return std::move(session_two);
  };
  EXPECT_CALL(mock_session_builder, Call())
      .WillOnce(Invoke(mock_creation_function_two))
      .RetiresOnSaturation();
  EXPECT_CALL(*session_two_ptr,
              Init(EqualsProto(spec), EqualsProto(hyperparams)));

  // Inject the function to the pool.
  ComputeSessionPoolTestPoolAccessor::SetSessionBuilder(
      &pool, mock_session_builder.AsStdFunction());

  // The first call will recieve the second session because of how the mocks go.
  auto first_request = pool.GetSession();
  EXPECT_EQ(1, pool.num_outstanding_sessions());
  EXPECT_EQ(first_request.get(), session_two_ptr);

  // Return the first pointer. After this, the second request should get that
  // pointer.
  EXPECT_CALL(*session_two_ptr, ResetSession());
  pool.ReturnSession(std::move(first_request));
  EXPECT_EQ(0, pool.num_outstanding_sessions());
  auto second_request = pool.GetSession();
  EXPECT_EQ(1, pool.num_outstanding_sessions());
  EXPECT_EQ(second_request.get(), session_two_ptr);

  // There are now no spare sessions, so the next session request should
  // create a second session.
  auto third_request = pool.GetSession();
  EXPECT_EQ(2, pool.num_outstanding_sessions());
  EXPECT_EQ(third_request.get(), session_one_ptr);
}

TEST(ComputeSessionPoolTest, AssignsUniqueIds) {
  MasterSpec spec;
  GridPoint hyperparams;
  ComputeSessionPool pool(spec, hyperparams);
  auto session = pool.GetSession();
  auto session_2 = pool.GetSession();
  EXPECT_NE(session->Id(), session_2->Id());
}

TEST(ComputeSessionPoolTest, SupportsMultithreadedAccess) {
  MasterSpec spec;
  GridPoint hyperparams;
  ComputeSessionPool pool(spec, hyperparams);

  std::vector<std::unique_ptr<tensorflow::Thread>> request_threads;
  constexpr int kNumThreadsToTest = 100;
  for (int i = 0; i < kNumThreadsToTest; ++i) {
    request_threads.push_back(std::unique_ptr<tensorflow::Thread>(
        tensorflow::Env::Default()->StartThread(
            tensorflow::ThreadOptions(), "thread",
            [this, &pool] { auto session = pool.GetSession(); })));
  }

  // Deleting a tensorflow::Thread blocks until the thread exits,
  // so clearing the vector blocks until all threads have exited.
  request_threads.clear();

  // Make sure all the threads got their session.
  EXPECT_EQ(kNumThreadsToTest, pool.num_outstanding_sessions());
}

}  // namespace dragnn
}  // namespace syntaxnet
