// Copyright 2019 The MediaPipe 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
//
//      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.

#include <map>
#include <string>
#include <vector>

#include "absl/strings/string_view.h"
#include "absl/synchronization/mutex.h"
#include "absl/time/clock.h"
#include "absl/time/time.h"
#include "mediapipe/framework/calculator_framework.h"
#include "mediapipe/framework/calculator_graph.h"
#include "mediapipe/framework/port/canonical_errors.h"
#include "mediapipe/framework/port/core_proto_inc.h"
#include "mediapipe/framework/port/gmock.h"
#include "mediapipe/framework/port/gtest.h"
#include "mediapipe/framework/port/integral_types.h"
#include "mediapipe/framework/port/logging.h"
#include "mediapipe/framework/port/status_matchers.h"
#include "mediapipe/framework/tool/sink.h"
#include "mediapipe/framework/tool/status_util.h"

namespace mediapipe {

namespace {

class CalculatorGraphEventLoopTest : public testing::Test {
 public:
  void AddThreadSafeVectorSink(const Packet& packet) {
    absl::WriterMutexLock lock(&output_packets_mutex_);
    output_packets_.push_back(packet);
  }

 protected:
  std::vector<Packet> output_packets_ ABSL_GUARDED_BY(output_packets_mutex_);
  absl::Mutex output_packets_mutex_;
};

// Allows blocking of the Process() call by locking the blocking_mutex passed to
// the input side packet. Used to force input stream queues to build up for
// testing.
class BlockingPassThroughCalculator : public CalculatorBase {
 public:
  static absl::Status GetContract(CalculatorContract* cc) {
    cc->Inputs().Index(0).SetAny();
    cc->Outputs().Index(0).SetSameAs(&cc->Inputs().Index(0));
    cc->InputSidePackets().Index(0).Set<std::unique_ptr<absl::Mutex>>();

    return absl::OkStatus();
  }

  absl::Status Open(CalculatorContext* cc) final {
    mutex_ = GetFromUniquePtr<absl::Mutex>(cc->InputSidePackets().Index(0));
    return absl::OkStatus();
  }

  absl::Status Process(CalculatorContext* cc) final {
    mutex_->Lock();
    cc->Outputs().Index(0).AddPacket(
        cc->Inputs().Index(0).Value().At(cc->InputTimestamp()));
    mutex_->Unlock();
    return absl::OkStatus();
  }

 private:
  absl::Mutex* mutex_;
};

REGISTER_CALCULATOR(BlockingPassThroughCalculator);

struct SimpleHeader {
  int width;
  int height;
};

class UsingHeaderCalculator : public CalculatorBase {
 public:
  static absl::Status GetContract(CalculatorContract* cc) {
    cc->Inputs().Index(0).SetAny();
    cc->Outputs().Index(0).SetSameAs(&cc->Inputs().Index(0));
    return absl::OkStatus();
  }

  absl::Status Open(CalculatorContext* cc) final {
    if (cc->Inputs().Index(0).Header().IsEmpty()) {
      return absl::UnknownError("No stream header present.");
    }

    const SimpleHeader& header =
        cc->Inputs().Index(0).Header().Get<SimpleHeader>();
    std::unique_ptr<SimpleHeader> output_header(new SimpleHeader);
    output_header->width = header.width;
    output_header->height = header.height;

    cc->Outputs().Index(0).SetHeader(Adopt(output_header.release()));
    return absl::OkStatus();
  }

  absl::Status Process(CalculatorContext* cc) final {
    cc->Outputs().Index(0).AddPacket(
        cc->Inputs().Index(0).Value().At(cc->InputTimestamp()));
    return absl::OkStatus();
  }
};
REGISTER_CALCULATOR(UsingHeaderCalculator);

TEST_F(CalculatorGraphEventLoopTest, WellProvisionedEventLoop) {
  CalculatorGraphConfig graph_config;
  ASSERT_TRUE(proto_ns::TextFormat::ParseFromString(
      R"(
          node {
            calculator: "PassThroughCalculator"
            input_stream: "input_numbers"
            output_stream: "output_numbers"
          }
          node {
            calculator: "CallbackCalculator"
            input_stream: "output_numbers"
            input_side_packet: "CALLBACK:callback"
          }
          input_stream: "input_numbers"
      )",
      &graph_config));

  // Start MediaPipe graph.
  CalculatorGraph graph(graph_config);
  MP_ASSERT_OK(graph.StartRun(
      {{"callback", MakePacket<std::function<void(const Packet&)>>(std::bind(
                        &CalculatorGraphEventLoopTest::AddThreadSafeVectorSink,
                        this, std::placeholders::_1))}}));

  // Insert 100 packets at the rate the calculator can keep up with.
  for (int i = 0; i < 100; ++i) {
    MP_ASSERT_OK(graph.AddPacketToInputStream(
        "input_numbers", Adopt(new int(i)).At(Timestamp(i))));
    // Wait for all packets to be received by the sink.
    while (true) {
      {
        absl::ReaderMutexLock lock(&output_packets_mutex_);
        if (output_packets_.size() > i) {
          break;
        }
      }
      absl::SleepFor(absl::Microseconds(1));
    }
  }
  // Check partial results.
  {
    absl::ReaderMutexLock lock(&output_packets_mutex_);
    ASSERT_EQ(100, output_packets_.size());
    for (int i = 0; i < 100; ++i) {
      EXPECT_EQ(i, output_packets_[i].Get<int>());
    }
  }

  // Insert 100 more packets at rate the graph can't keep up.
  for (int i = 100; i < 200; ++i) {
    MP_ASSERT_OK(graph.AddPacketToInputStream(
        "input_numbers", Adopt(new int(i)).At(Timestamp(i))));
  }
  // Don't wait but just close the input stream.
  MP_ASSERT_OK(graph.CloseInputStream("input_numbers"));
  // Wait properly via the API until the graph is done.
  MP_ASSERT_OK(graph.WaitUntilDone());
  // Check final results.
  {
    absl::ReaderMutexLock lock(&output_packets_mutex_);
    ASSERT_EQ(200, output_packets_.size());
    for (int i = 0; i < 200; ++i) {
      EXPECT_EQ(i, output_packets_[i].Get<int>());
    }
  }
}

// Pass-Through calculator that fails upon receiving the 10th packet.
class FailingPassThroughCalculator : public CalculatorBase {
 public:
  static absl::Status GetContract(CalculatorContract* cc) {
    cc->Inputs().Index(0).SetAny();
    cc->Outputs().Index(0).SetSameAs(&cc->Inputs().Index(0));
    return absl::OkStatus();
  }

  absl::Status Process(CalculatorContext* cc) final {
    Timestamp timestamp = cc->InputTimestamp();
    if (timestamp.Value() == 9) {
      return absl::UnknownError("Meant to fail (magicstringincludedhere).");
    }
    cc->Outputs().Index(0).AddPacket(
        cc->Inputs().Index(0).Value().At(timestamp));
    return absl::OkStatus();
  }
};
REGISTER_CALCULATOR(FailingPassThroughCalculator);

TEST_F(CalculatorGraphEventLoopTest, FailingEventLoop) {
  CalculatorGraphConfig graph_config;
  ASSERT_TRUE(proto_ns::TextFormat::ParseFromString(
      R"(
          node {
            calculator: "FailingPassThroughCalculator"
            input_stream: "input_numbers"
            output_stream: "output_numbers"
          }
          node {
            calculator: "CallbackCalculator"
            input_stream: "output_numbers"
            input_side_packet: "CALLBACK:callback"
          }
          input_stream: "input_numbers")",
      &graph_config));

  // Start MediaPipe graph.
  CalculatorGraph graph(graph_config);
  MP_ASSERT_OK(graph.StartRun(
      {{"callback", MakePacket<std::function<void(const Packet&)>>(std::bind(
                        &CalculatorGraphEventLoopTest::AddThreadSafeVectorSink,
                        this, std::placeholders::_1))}}));

  // Insert packets.
  absl::Status status;
  for (int i = 0; true; ++i) {
    status = graph.AddPacketToInputStream("input_numbers",
                                          Adopt(new int(i)).At(Timestamp(i)));
    if (!status.ok()) {
      ASSERT_TRUE(graph.HasError());  // Graph failed.
      ASSERT_THAT(
          status.message(),
          testing::HasSubstr("Meant to fail (magicstringincludedhere)."));
      break;
    }
  }
  MP_ASSERT_OK(graph.CloseInputStream("input_numbers"));
  status = graph.WaitUntilDone();
  ASSERT_THAT(status.message(),
              testing::HasSubstr("Meant to fail (magicstringincludedhere)."));
}

// Test the step by step mode.
TEST_F(CalculatorGraphEventLoopTest, StepByStepSchedulerLoop) {
  CalculatorGraphConfig graph_config;
  ASSERT_TRUE(proto_ns::TextFormat::ParseFromString(
      R"(
          node {
            calculator: "PassThroughCalculator"
            input_stream: "input_numbers"
            output_stream: "output_numbers"
          }
          node {
            calculator: "CallbackCalculator"
            input_stream: "output_numbers"
            input_side_packet: "CALLBACK:callback"
          }
          input_stream: "input_numbers"
      )",
      &graph_config));

  // Start MediaPipe graph.
  CalculatorGraph graph(graph_config);
  MP_ASSERT_OK(graph.StartRun(
      {{"callback", MakePacket<std::function<void(const Packet&)>>(std::bind(
                        &CalculatorGraphEventLoopTest::AddThreadSafeVectorSink,
                        this, std::placeholders::_1))}}));

  // Add packet one at a time, we should be able to syncrhonize the output for
  // each addition in the step by step mode.
  for (int i = 0; i < 100; ++i) {
    MP_ASSERT_OK(graph.AddPacketToInputStream(
        "input_numbers", Adopt(new int(i)).At(Timestamp(i))));
    MP_ASSERT_OK(graph.WaitUntilIdle());
    absl::ReaderMutexLock lock(&output_packets_mutex_);
    ASSERT_EQ(i + 1, output_packets_.size());
  }
  // Don't wait but just close the input stream.
  MP_ASSERT_OK(graph.CloseInputStream("input_numbers"));
  // Wait properly via the API until the graph is done.
  MP_ASSERT_OK(graph.WaitUntilDone());
}

// Test setting the stream header.
TEST_F(CalculatorGraphEventLoopTest, SetStreamHeader) {
  CalculatorGraphConfig graph_config;
  ASSERT_TRUE(proto_ns::TextFormat::ParseFromString(
      R"(
          node {
            calculator: "UsingHeaderCalculator"
            input_stream: "input_numbers"
            output_stream: "output_numbers"
          }
          node {
            calculator: "CallbackCalculator"
            input_stream: "output_numbers"
            input_side_packet: "CALLBACK:callback"
          }
          input_stream: "input_numbers"
      )",
      &graph_config));

  CalculatorGraph graph(graph_config);
  MP_ASSERT_OK(graph.StartRun(
      {{"callback", MakePacket<std::function<void(const Packet&)>>(std::bind(
                        &CalculatorGraphEventLoopTest::AddThreadSafeVectorSink,
                        this, std::placeholders::_1))}}));

  absl::Status status = graph.WaitUntilIdle();
  // Expect to fail if header not set.
  ASSERT_FALSE(status.ok());
  EXPECT_EQ(status.code(), absl::StatusCode::kUnknown);
  EXPECT_THAT(status.message(),
              testing::HasSubstr("No stream header present."));

  CalculatorGraph graph2(graph_config);
  std::unique_ptr<SimpleHeader> header(new SimpleHeader);
  header->width = 320;
  header->height = 240;
  // With stream header set, the StartRun should succeed.
  MP_ASSERT_OK(graph2.StartRun(
      {{"callback", MakePacket<std::function<void(const Packet&)>>(std::bind(
                        &CalculatorGraphEventLoopTest::AddThreadSafeVectorSink,
                        this, std::placeholders::_1))}},
      {{"input_numbers", Adopt(header.release())}}));
  // Don't wait but just close the input stream.
  MP_ASSERT_OK(graph2.CloseInputStream("input_numbers"));
  // Wait properly via the API until the graph is done.
  MP_ASSERT_OK(graph2.WaitUntilDone());
}

// Test ADD_IF_NOT_FULL mode for graph input streams (by creating more packets
// than the queue will support). At least some of these attempts should fail.
TEST_F(CalculatorGraphEventLoopTest, TryToAddPacketToInputStream) {
  CalculatorGraphConfig graph_config;
  ASSERT_TRUE(proto_ns::TextFormat::ParseFromString(
      R"(
          node {
            calculator: "BlockingPassThroughCalculator"
            input_stream: "input_numbers"
            output_stream: "output_numbers"
            input_side_packet: "blocking_mutex"
          }
          node {
            calculator: "CallbackCalculator"
            input_stream: "output_numbers"
            input_side_packet: "CALLBACK:callback"
          }
          input_stream: "input_numbers"
          num_threads: 2
          max_queue_size: 1
      )",
      &graph_config));

  absl::Mutex* mutex = new absl::Mutex();
  Packet mutex_side_packet = AdoptAsUniquePtr(mutex);

  CalculatorGraph graph(graph_config);
  graph.SetGraphInputStreamAddMode(
      CalculatorGraph::GraphInputStreamAddMode::ADD_IF_NOT_FULL);

  // Start MediaPipe graph.
  MP_ASSERT_OK(graph.StartRun(
      {{"callback", MakePacket<std::function<void(const Packet&)>>(std::bind(
                        &CalculatorGraphEventLoopTest::AddThreadSafeVectorSink,
                        this, std::placeholders::_1))},
       {"blocking_mutex", mutex_side_packet}}));

  constexpr int kNumInputPackets = 20;
  constexpr int kMaxQueueSize = 1;

  // Lock the mutex so that the BlockingPassThroughCalculator cannot read any of
  // these packets.
  mutex->Lock();
  int fail_count = 0;
  // Expect at least kNumInputPackets - kMaxQueueSize - 1 attempts to add
  // packets to fail since the queue builds up. The -1 is because our throttling
  // mechanism could be off by 1 at most due to the order of acquisition of
  // locks.
  for (int i = 0; i < kNumInputPackets; ++i) {
    absl::Status status = graph.AddPacketToInputStream(
        "input_numbers", Adopt(new int(i)).At(Timestamp(i)));
    if (!status.ok()) {
      ++fail_count;
    }
  }
  mutex->Unlock();

  EXPECT_GE(fail_count, kNumInputPackets - kMaxQueueSize - 1);
  // Don't wait but just close the input stream.
  MP_ASSERT_OK(graph.CloseInputStream("input_numbers"));
  // Wait properly via the API until the graph is done.
  MP_ASSERT_OK(graph.WaitUntilDone());
}

// Verify that "max_queue_size: -1" disables throttling of graph-input-streams.
TEST_F(CalculatorGraphEventLoopTest, ThrottlingDisabled) {
  CalculatorGraphConfig graph_config;
  ASSERT_TRUE(proto_ns::TextFormat::ParseFromString(
      R"(
          node {
            calculator: "BlockingPassThroughCalculator"
            input_stream: "input_numbers"
            output_stream: "output_numbers"
            input_side_packet: "blocking_mutex"
          }
          input_stream: "input_numbers"
          max_queue_size: -1
      )",
      &graph_config));

  absl::Mutex* mutex = new absl::Mutex();
  Packet mutex_side_packet = AdoptAsUniquePtr(mutex);

  CalculatorGraph graph(graph_config);
  graph.SetGraphInputStreamAddMode(
      CalculatorGraph::GraphInputStreamAddMode::ADD_IF_NOT_FULL);

  // Start MediaPipe graph.
  MP_ASSERT_OK(graph.StartRun({{"blocking_mutex", mutex_side_packet}}));

  // Lock the mutex so that the BlockingPassThroughCalculator cannot read any
  // of these packets.
  mutex->Lock();
  for (int i = 0; i < 10; ++i) {
    MP_EXPECT_OK(graph.AddPacketToInputStream(
        "input_numbers", Adopt(new int(i)).At(Timestamp(i))));
  }
  mutex->Unlock();
  MP_EXPECT_OK(graph.CloseInputStream("input_numbers"));
  MP_EXPECT_OK(graph.WaitUntilDone());
}

// Verify that the graph input stream throttling code still works if we run the
// graph twice.
TEST_F(CalculatorGraphEventLoopTest, ThrottleGraphInputStreamTwice) {
  CalculatorGraphConfig graph_config;
  ASSERT_TRUE(proto_ns::TextFormat::ParseFromString(
      R"(
          node {
            calculator: "BlockingPassThroughCalculator"
            input_stream: "input_numbers"
            output_stream: "output_numbers"
            input_side_packet: "blocking_mutex"
          }
          input_stream: "input_numbers"
          max_queue_size: 1
      )",
      &graph_config));

  absl::Mutex* mutex = new absl::Mutex();
  Packet mutex_side_packet = AdoptAsUniquePtr(mutex);

  CalculatorGraph graph(graph_config);
  graph.SetGraphInputStreamAddMode(
      CalculatorGraph::GraphInputStreamAddMode::ADD_IF_NOT_FULL);

  // Run the graph twice.
  for (int i = 0; i < 2; ++i) {
    // Start MediaPipe graph.
    MP_ASSERT_OK(graph.StartRun({{"blocking_mutex", mutex_side_packet}}));

    // Lock the mutex so that the BlockingPassThroughCalculator cannot read any
    // of these packets.
    mutex->Lock();
    absl::Status status = absl::OkStatus();
    for (int i = 0; i < 10; ++i) {
      status = graph.AddPacketToInputStream("input_numbers",
                                            Adopt(new int(i)).At(Timestamp(i)));
      if (!status.ok()) {
        break;
      }
    }
    mutex->Unlock();
    ASSERT_FALSE(status.ok());
    EXPECT_EQ(status.code(), absl::StatusCode::kUnavailable);
    EXPECT_THAT(status.message(), testing::HasSubstr("Graph is throttled."));
    MP_ASSERT_OK(graph.CloseInputStream("input_numbers"));
    MP_ASSERT_OK(graph.WaitUntilDone());
  }
}

// Test WAIT_TILL_NOT_FULL mode (default mode) for graph input streams (by
// creating more packets than the queue will support). All packets sent to the
// graph should be processed.
TEST_F(CalculatorGraphEventLoopTest, WaitToAddPacketToInputStream) {
  CalculatorGraphConfig graph_config;
  ASSERT_TRUE(proto_ns::TextFormat::ParseFromString(
      R"(
          node {
            calculator: "PassThroughCalculator"
            input_stream: "input_numbers"
            output_stream: "output_numbers"
          }
          node {
            calculator: "CallbackCalculator"
            input_stream: "output_numbers"
            input_side_packet: "CALLBACK:callback"
          }
          input_stream: "input_numbers"
          num_threads: 2
          max_queue_size: 10
      )",
      &graph_config));

  // Start MediaPipe graph.
  CalculatorGraph graph(graph_config);
  MP_ASSERT_OK(graph.StartRun(
      {{"callback", MakePacket<std::function<void(const Packet&)>>(std::bind(
                        &CalculatorGraphEventLoopTest::AddThreadSafeVectorSink,
                        this, std::placeholders::_1))}}));

  constexpr int kNumInputPackets = 20;
  // All of these packets should be accepted by the graph.
  int fail_count = 0;
  for (int i = 0; i < kNumInputPackets; ++i) {
    absl::Status status = graph.AddPacketToInputStream(
        "input_numbers", Adopt(new int(i)).At(Timestamp(i)));
    if (!status.ok()) {
      ++fail_count;
    }
  }

  EXPECT_EQ(0, fail_count);

  // Don't wait but just close the input stream.
  MP_ASSERT_OK(graph.CloseInputStream("input_numbers"));
  // Wait properly via the API until the graph is done.
  MP_ASSERT_OK(graph.WaitUntilDone());

  absl::ReaderMutexLock lock(&output_packets_mutex_);
  ASSERT_EQ(kNumInputPackets, output_packets_.size());
}

// Captures log messages during testing.
class TextMessageLogSink : public LogSink {
 public:
  std::vector<std::string> messages;
  void Send(const LogEntry& entry) {
    messages.push_back(std::string(entry.text_message()));
  }
};

// Verifies that CalculatorGraph::UnthrottleSources does not run repeatedly
// in a "busy-loop" while the graph is throttled due to a graph-output stream.
TEST_F(CalculatorGraphEventLoopTest, UnthrottleSources) {
  CalculatorGraphConfig graph_config;
  ASSERT_TRUE(proto_ns::TextFormat::ParseFromString(
      R"(
          node {
            calculator: "PassThroughCalculator"
            input_stream: "input_numbers"
            output_stream: "output_numbers"
          }
          input_stream: "input_numbers"
          output_stream: "output_numbers"
          num_threads: 2
          max_queue_size: 5
      )",
      &graph_config));
  constexpr int kQueueSize = 5;

  // Initialize and start the mediapipe graph.
  CalculatorGraph graph;
  MP_ASSERT_OK(graph.Initialize(graph_config));
  graph.SetGraphInputStreamAddMode(
      CalculatorGraph::GraphInputStreamAddMode::ADD_IF_NOT_FULL);
  auto poller_status = graph.AddOutputStreamPoller("output_numbers");
  MP_ASSERT_OK(poller_status.status());
  mediapipe::OutputStreamPoller& poller = poller_status.value();
  poller.SetMaxQueueSize(kQueueSize);
  MP_ASSERT_OK(graph.StartRun({}));

  // Lambda that adds a packet to the calculator graph.
  auto add_packet = [&graph](std::string s, int i) {
    return graph.AddPacketToInputStream(s, MakePacket<int>(i).At(Timestamp(i)));
  };

  // Start capturing VLOG messages from the mediapipe::Scheduler.
  TextMessageLogSink log_listener;
  mediapipe::AddLogSink(&log_listener);
  SetVLOGLevel("scheduler", 3);

  // Add just enough packets to fill the output stream queue.
  std::vector<Packet> out_packets;
  for (int i = 0; i < kQueueSize; ++i) {
    MP_EXPECT_OK(add_packet("input_numbers", i));
    MP_EXPECT_OK(graph.WaitUntilIdle());
  }

  // The graph is throttled due to the full output stream.
  EXPECT_FALSE(add_packet("input_numbers", kQueueSize).ok());

  // CalculatorGraph::UnthrottleSources should be called just one time.
  absl::SleepFor(absl::Milliseconds(100));

  // Read all packets from the output stream queue and close the graph.
  for (int i = 0; i < kQueueSize; ++i) {
    Packet packet;
    EXPECT_TRUE(poller.Next(&packet));
    out_packets.push_back(packet);
  }
  MP_EXPECT_OK(graph.CloseAllInputStreams());
  MP_EXPECT_OK(graph.WaitUntilDone());
  EXPECT_EQ(kQueueSize, out_packets.size());

  // Stop capturing VLOG messages.
  SetVLOGLevel("scheduler", 0);
  mediapipe::RemoveLogSink(&log_listener);

  // Count and validate the calls to UnthrottleSources.
  int loop_count = 0;
  for (auto& message : log_listener.messages) {
    loop_count += (message == "HandleIdle: unthrottling") ? 1 : 0;
  }
  EXPECT_GE(loop_count, 1);
  EXPECT_LE(loop_count, 2);
}

}  // namespace
}  // namespace mediapipe