#pragma once

#include <c10/core/impl/InlineDeviceGuard.h>

namespace c10 {
namespace impl {

/**
 * A StreamGuard is an RAII class that changes the current device
 * to the device corresponding to some stream, and changes the
 * default stream on that device to be this stream.
 *
 * InlineStreamGuard is a helper class for implementing StreamGuards.
 * See InlineDeviceGuard for guidance on how to use this class.
 */
template <typename T>
class InlineStreamGuard : private InlineDeviceGuard<T> {
public:
  /// No default constructor, see Note [Omitted default constructor from RAII]
  explicit InlineStreamGuard() = delete;

  /// Set the current device to the device associated with the passed stream,
  /// and set the current stream on that device to the passed stream.
  explicit InlineStreamGuard(Stream stream)
    : InlineDeviceGuard<T>(stream.device())
    , original_stream_of_original_device_(this->impl_.getStream(original_device()))
    , original_stream_of_current_device_(this->impl_.exchangeStream(stream))
    , current_stream_(stream)
    {}

  /// This constructor exists purely for testing
  template <typename U=T, typename=typename std::enable_if<std::is_same<U, VirtualGuardImpl>::value>::type>
  explicit InlineStreamGuard(Stream stream, const DeviceGuardImplInterface* impl)
    : InlineDeviceGuard<T>(stream.device(), impl ? impl : getDeviceGuardImpl(stream.device_type()))
    , original_stream_of_original_device_(this->impl_.getStream(original_device()))
    , original_stream_of_current_device_(this->impl_.exchangeStream(stream))
    , current_stream_(stream)
    {}

  /// Copy is disallowed
  InlineStreamGuard(const InlineStreamGuard<T>&) = delete;
  InlineStreamGuard<T>& operator=(const InlineStreamGuard<T>&) = delete;

  /// Move is disallowed, as StreamGuard does not have an uninitialized state,
  /// which is required for moves on types with nontrivial destructors.
  InlineStreamGuard(InlineStreamGuard<T>&& other) = delete;
  InlineStreamGuard& operator=(InlineStreamGuard<T>&& other) = delete;

  ~InlineStreamGuard() {
    this->impl_.exchangeStream(original_stream_of_current_device_);
  }

  /// Resets the currently set stream to the original stream and
  /// the currently set device to the original device.  Then,
  /// set the current device to the device associated with the passed stream,
  /// and set the current stream on that device to the passed stream.
  ///
  /// NOTE: this implementation may skip some stream/device setting if
  /// it can prove that it is unnecessary.
  ///
  /// WARNING: reset_stream does NOT preserve previously set streams on
  /// different devices.  If you need to set streams on multiple devices
  /// on CUDA, use CUDAMultiStreamGuard instead.
  void reset_stream(Stream stream) {
    // TODO: make a version that takes an impl argument.  Unfortunately,
    // that will require SFINAE because impl is only valid for the
    // VirtualGuardImpl specialization.
    if (stream.device() == this->current_device()) {
      this->impl_.exchangeStream(stream);
      current_stream_ = stream;
    } else {
      // Destruct and reconstruct the StreamGuard in-place
      this->impl_.exchangeStream(original_stream_of_current_device_);
      this->reset_device(stream.device());
      original_stream_of_current_device_ = this->impl_.exchangeStream(stream);
      current_stream_ = stream;
    }
  }

  // It's not clear if set_device should also reset the current stream
  // if the device is unchanged; therefore, we don't provide it.
  // The situation is somewhat clearer with reset_device, but it's still
  // a pretty weird thing to do, so haven't added this either.

  /// Returns the stream of the original device prior to this guard.  Subtly,
  /// the stream returned here is the original stream of the *original*
  /// device; i.e., it's the stream that your computation *would* have
  /// been put on, if it hadn't been for this meddling stream guard.
  /// This is usually what you want.
  Stream original_stream() const {
    return original_stream_of_original_device_;
  }

  /// Returns the most recent stream that was set using this device guard,
  /// either from construction, or via set_stream.
  Stream current_stream() const {
    return current_stream_;
  }

  /// Returns the most recent device that was set using this device guard,
  /// either from construction, or via set_device/reset_device/set_index.
  Device current_device() const {
    return InlineDeviceGuard<T>::current_device();
  }

  /// Returns the device that was set at the most recent reset_stream(),
  /// or otherwise the device at construction time.
  Device original_device() const {
    return InlineDeviceGuard<T>::original_device();
  }

private:
  Stream original_stream_of_original_device_; // what the user probably cares about
  Stream original_stream_of_current_device_; // what we need to restore
  Stream current_stream_;
};

/**
 * An OptionalStreamGuard is an RAII class that sets a device to some value on
 * initialization, and resets the device to its original value on destruction.
 * See InlineOptionalDeviceGuard for more guidance on how to use this class.
 */
template <typename T>
class InlineOptionalStreamGuard {
public:
  /// Creates an uninitialized stream guard.
  explicit InlineOptionalStreamGuard()
    : guard_() // See Note [Explicit initialization of optional fields]
    {}

  /// Set the current device to the device associated with the passed stream,
  /// and set the current stream on that device to the passed stream,
  /// if the passed stream is not nullopt.
  explicit InlineOptionalStreamGuard(optional<Stream> stream_opt)
    : guard_() {
    if (stream_opt.has_value()) {
      guard_.emplace(stream_opt.value());
    }
  }

  /// All constructors of StreamGuard are valid for OptionalStreamGuard
  template <typename... Args>
  explicit InlineOptionalStreamGuard(Args&&... args)
    : guard_(in_place, std::forward<Args>(args)...) {}

  // See Note [Move construction for RAII guards is tricky]
  InlineOptionalStreamGuard(InlineOptionalStreamGuard<T>&& other) = delete;

  // See Note [Move assignment for RAII guards is tricky]
  InlineOptionalStreamGuard& operator=(InlineOptionalStreamGuard&& other) = delete;

  /// Resets the currently set stream to the original stream and
  /// the currently set device to the original device.  Then,
  /// set the current device to the device associated with the passed stream,
  /// and set the current stream on that device to the passed stream.
  /// Initializes the OptionalStreamGuard if it was not previously initialized.
  void reset_stream(Stream stream) {
    if (guard_.has_value()) {
      guard_->reset_stream(stream);
    } else {
      guard_.emplace(stream);
    }
  }

  /// Returns the stream that was set at the time the guard was most recently
  /// initialized, or nullopt if the guard is uninitialized.
  optional<Stream> original_stream() const {
    return guard_.has_value() ? make_optional(guard_->original_stream()) : nullopt;
  }

  /// Returns the most recent stream that was set using this stream guard,
  /// either from construction, or via reset_stream, if the guard is initialized,
  /// or nullopt if the guard is uninitialized.
  optional<Stream> current_stream() const {
    return guard_.has_value() ? make_optional(guard_->current_stream()) : nullopt;
  }

  /// Restore the original device and stream, resetting this guard to uninitialized state.
  void reset() {
    guard_.reset();
  }

private:
  optional<InlineStreamGuard<T>> guard_;
};

}} // namespace c10::impl