// Copyright (c) 2019 PaddlePaddle Authors. All Rights Reserved.
//
// 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.

#pragma once

#if defined _WIN32 || defined __APPLE__
#else
#define _LINUX
#endif

#include <glog/logging.h>
#include <algorithm>
#include <condition_variable>  // NOLINT
#include <deque>
#include <limits>
#include <memory>
#include <mutex>  // NOLINT
#include <utility>
#include <vector>
#include "paddle/fluid/framework/expect.h"

namespace paddle {
namespace framework {

template <class T>
class ChannelObject {
 public:
  ChannelObject() {}

  // capacity can be zero
  explicit ChannelObject(size_t capacity) {
    capacity_ = std::min(MaxCapacity(), capacity);
  }

  void Clear() {
    std::unique_lock<std::mutex> lock(mutex_);
    data_.clear();
    data_.shrink_to_fit();
  }

  size_t Capacity() {
    return capacity_;  // atomic
  }

  void SetCapacity(size_t x) {  // capacity can be zero
    std::lock_guard<std::mutex> lock(mutex_);
    capacity_ = std::min(MaxCapacity(), x);
    Notify();
  }

  size_t BlockSize() {
    return block_size_;  // atomic
  }

  void SetBlockSize(size_t x) {
    CHECK(x >= 1) << "block size must be >= 1";
    std::lock_guard<std::mutex> lock(mutex_);
    block_size_ = x;
  }

  template <class U>
  void InheritFrom(const std::shared_ptr<ChannelObject<U>>& other) {
    std::lock_guard<std::mutex> lock(mutex_);
    capacity_ = other->Capacity();
    block_size_ = other->BlockSize();
  }

  bool Closed() {
    return closed_;  // atomic
  }

  // open channel, then data can be write() to channel
  void Open() {
    std::lock_guard<std::mutex> lock(mutex_);
    closed_ = false;
    Notify();
  }

  // close channel, then no more data can be write() to channel
  void Close() {
    std::lock_guard<std::mutex> lock(mutex_);
    closed_ = true;
    Notify();
  }

  size_t Size() {
    std::lock_guard<std::mutex> lock(mutex_);
    return data_.size();
  }

  bool Empty() {
    std::lock_guard<std::mutex> lock(mutex_);
    return EmptyUnlocked();
  }

  // blocking operation
  bool Get(T& val) { return Read(1, &val) != 0; }  // NOLINT

  // blocking operation
  // returns 0 if the channel is closed and empty
  size_t Read(size_t n, T* p) {
    if (n == 0) {
      return 0;
    }

    std::unique_lock<std::mutex> lock(mutex_);
    size_t finished = Read(n, p, lock);
    Notify();
    return finished;
  }

  // blocking operation
  bool Put(T&& val) { return WriteMove(1, &val) != 0; }

  // blocking operation
  bool Put(const T& val) { return Write(1, &val) != 0; }

  // blocking operation
  // returns value less than n if the channel is closed
  size_t Write(size_t n, const T* p) {
    if (n == 0) {
      return 0;
    }
    std::unique_lock<std::mutex> lock(mutex_);
    size_t finished = Write(n, p, lock);
    Notify();
    return finished;
  }

  // WriteMove() will clear original contents of input array
  size_t WriteMove(size_t n, T* p) {
    if (n == 0) {
      return 0;
    }
    std::unique_lock<std::mutex> lock(mutex_);
    size_t finished = WriteMove(n, p, lock);
    Notify();
    return finished;
  }

  // read data of block size from channel to vector
  size_t Read(std::vector<T>& p) {  // NOLINT
    p.resize(block_size_);
    size_t finished = Read(p.size(), &p[0]);
    p.resize(finished);
    return finished;
  }

  size_t ReadAll(std::vector<T>& p) {  // NOLINT
    p.clear();
    size_t finished = 0;
    size_t n = 0;
    do {
      // _block_size may change anytime
      n = block_size_;
      p.resize(finished + n);
      n = Read(n, &p[finished]);
      finished += n;
    } while (n != 0);
    p.resize(finished);
    return finished;
  }

  // write data from vector to channel
  size_t Write(const std::vector<T>& p) { return Write(p.size(), &p[0]); }

  // write data from vector to channel
  size_t Write(std::vector<T>&& p) { return WriteMove(p.size(), &p[0]); }

 private:
  size_t capacity_ = MaxCapacity();
  size_t block_size_ = 1024;
  bool closed_ = false;
  std::mutex mutex_;
  // use deque to store data
  std::deque<T> data_;
  size_t reading_count_ = 0;
  int empty_waiters_ = 0;
  int full_waiters_ = 0;
  std::condition_variable empty_cond_;
  std::condition_variable full_cond_;

  static constexpr size_t MaxCapacity() {
    return std::numeric_limits<size_t>::max() / 2;
  }

  void Notify() {
    if (empty_waiters_ != 0 && (!EmptyUnlocked() || closed_)) {
      empty_cond_.notify_one();
    }
    if (full_waiters_ != 0 && (!FullUnlocked() || closed_)) {
      full_cond_.notify_one();
    }
  }

  bool EmptyUnlocked() { return data_.empty(); }

  bool FullUnlocked() { return data_.size() >= capacity_ + reading_count_; }

  bool WaitForRead(std::unique_lock<std::mutex>& lock) {  // NOLINT
#ifdef _LINUX
    while (unlikely(EmptyUnlocked() && !closed_)) {
#else
    while (EmptyUnlocked() && !closed_) {
#endif
      if (full_waiters_ != 0) {
        full_cond_.notify_one();
      }
      empty_waiters_++;
      empty_cond_.wait(lock);
      empty_waiters_--;
    }
    return !EmptyUnlocked();
  }

  bool WaitForWrite(std::unique_lock<std::mutex>& lock) {  // NOLINT
#ifdef _LINUX
    while (unlikely(FullUnlocked() && !closed_)) {
#else
    while (FullUnlocked() && !closed_) {
#endif
      if (empty_waiters_ != 0) {
        empty_cond_.notify_one();
      }
      full_waiters_++;
      full_cond_.wait(lock);
      full_waiters_--;
    }
    return !closed_;
  }

  size_t Read(size_t n, T* p, std::unique_lock<std::mutex>& lock) {  // NOLINT
    size_t finished = 0;
    CHECK(n <= MaxCapacity() - reading_count_);
    reading_count_ += n;
    while (finished < n && WaitForRead(lock)) {
      size_t m = std::min(n - finished, data_.size());
      for (size_t i = 0; i < m; i++) {
        p[finished++] = std::move(data_.front());
        data_.pop_front();
      }
      reading_count_ -= m;
    }
    reading_count_ -= n - finished;
    return finished;
  }

  size_t Write(size_t n,
               const T* p,                            // NOLINT
               std::unique_lock<std::mutex>& lock) {  // NOLINT
    size_t finished = 0;
    while (finished < n && WaitForWrite(lock)) {
      size_t m =
          std::min(n - finished, capacity_ + reading_count_ - data_.size());
      for (size_t i = 0; i < m; i++) {
        data_.push_back(p[finished++]);
      }
    }
    return finished;
  }

  size_t WriteMove(size_t n,
                   T* p,                                  // NOLINT
                   std::unique_lock<std::mutex>& lock) {  // NOLINT
    size_t finished = 0;
    while (finished < n && WaitForWrite(lock)) {
      size_t m =
          std::min(n - finished, capacity_ + reading_count_ - data_.size());
      for (size_t i = 0; i < m; i++) {
        data_.push_back(std::move(p[finished++]));
      }
    }
    return finished;
  }
};  // NOLINT

template <class T>
using Channel = std::shared_ptr<ChannelObject<T>>;

template <class T>
Channel<T> MakeChannel(size_t capacity = std::numeric_limits<size_t>::max()) {
  return std::make_shared<ChannelObject<T>>(capacity);
}

template <class T, class U>
Channel<T> MakeChannel(const Channel<U>& other) {
  CHECK(other != nullptr) << "channel can not be NULL";
  Channel<T> chan = std::make_shared<ChannelObject<T>>();
  chan->InheritFrom(other);
  return chan;
}

// NOTE: ChannelReader is a wrapper for quick read channel with a buffer. It
// will read a block data from channel, but user can get data one by one. So it
// is important to notice that user must call operator>> until false, or call
// get_buffer_remain until false to make sure the buffered data all readed.
template <class T>
class ChannelReader {
 public:
  explicit ChannelReader(ChannelObject<T>* channel = nullptr) {
    Reset(channel);
  }

  ~ChannelReader() { CHECK(cursor_ == 0) << "Forgot to read buffer data"; }

  ChannelObject<T>* channel() { return channel_; }

  void Reset(ChannelObject<T>* channel) {
    CHECK(channel != nullptr) << "Channel can not be nullptr";
    channel_ = channel;
    cursor_ = 0;
    failed_ = !channel;
  }

  // whether there were read failed
  operator bool() { return !failed_; }

  ChannelReader<T>& operator>>(T& val) {
    if (failed_) {
      return *this;
    }
    if (cursor_ >= buffer_.size()) {
      cursor_ = 0;
      if (channel_->read(buffer_) == 0) {
        failed_ = true;
        return *this;
      }
    }
    val = std::move(buffer_[cursor_++]);
    return *this;
  }

  bool GetBufferRemain(T& val) {  // NOLINT
    if (cursor_ >= buffer_.size()) {
      cursor_ = 0;
      return false;
    }
    val = std::move(buffer_[cursor_++]);
    return true;
  }

 private:
  ChannelObject<T>* channel_ = nullptr;
  std::vector<T> buffer_;
  size_t cursor_ = 0;
  bool failed_ = true;
};  // NOLINT

template <class T>
class ChannelWriter {
 public:
  explicit ChannelWriter(ChannelObject<T>* channel = nullptr) {
    Reset(channel);
  }

  ~ChannelWriter() { CHECK(buffer_.empty()) << "Forgot to flush"; }

  ChannelObject<T>* channel() { return channel_; }

  void Reset(ChannelObject<T>* channel) {
    CHECK(buffer_.empty()) << "Forgot to flush";
    CHECK(channel != nullptr) << "Channel can not be nullptr";
    channel_ = channel;
    buffer_.clear();
    failed_ = !channel;
  }

  // whether there were write failed
  operator bool() { return !failed_; }

  ChannelWriter<T>& operator<<(T&& val) {
    if (failed_) {
      return *this;
    }
    buffer_.push_back(std::move(val));
    if (buffer_.size() >= channel_->BlockSize()) {
      Flush();
    }
    return *this;
  }

  ChannelWriter<T>& operator<<(const T& val) {
    if (failed_) {
      return *this;
    }
    buffer_.push_back(val);
    if (buffer_.size() >= channel_->BlockSize()) {
      Flush();
    }
    return *this;
  }

  void Flush() {
    if (failed_ || buffer_.empty()) {
      buffer_.clear();
      return;
    }
    failed_ |=
        channel_->WriteMove(buffer_.size(), &buffer_[0]) != buffer_.size();
    buffer_.clear();
  }

 private:
  ChannelObject<T>* channel_ = nullptr;
  std::vector<T> buffer_;
  bool failed_ = true;
};  // NOLINT

// only used for range-for loop
// for (auto& x : chan) {...}
template <class T>
struct ChannelIterator {
  std::shared_ptr<ChannelReader<T>> reader_;
  T data_;

  void operator++() {
    CHECK(reader_ != nullptr) << "reader can not be NULL";
    if (!(*reader_ >> data_)) {
      reader_ = nullptr;
    }
  }

  T& operator*() { return data_; }

  friend bool operator==(const ChannelIterator<T>& a,
                         const ChannelIterator<T>& b) {
    return a.reader_ == b.reader_;
  }

  friend bool operator!=(const ChannelIterator<T>& a,
                         const ChannelIterator<T>& b) {
    return a.reader_ != b.reader_;
  }
};  // NOLINT

template <class T>
ChannelIterator<T> begin(ChannelObject<T>* chan) {
  ChannelIterator<T> it{std::make_shared<ChannelReader<T>>(chan), T()};
  ++it;
  return it;
}

template <class T>
ChannelIterator<T> end(ChannelObject<T>* chan) {
  return {nullptr, T()};
}

}  // namespace framework
}  // namespace paddle