#pragma once

#include <ATen/core/dispatch/DispatchTable.h>
#include <ATen/core/dispatch/OperatorOptions.h>
#include <ATen/core/dispatch/RegistrationHandleRAII.h>
#include <list>

namespace c10 {
namespace impl {
  class OperatorEntry;
}

namespace impl {

// This is a private class used inside the Dispatcher to represent an operator
// and its dispatch table. This is not part of the public API.
class OperatorEntry final {
public:
  explicit OperatorEntry(FunctionSchema&& schema, OperatorOptions&& options);

  OperatorEntry(const OperatorEntry&) = delete;
  OperatorEntry(OperatorEntry&&) noexcept = delete;
  OperatorEntry& operator=(const OperatorEntry&) = delete;
  OperatorEntry& operator=(OperatorEntry&&) noexcept = delete;

  const FunctionSchema& schema() const {
    return schema_;
  }

  template<class Return, class... Args>
  Return callUnboxed(TensorTypeId dispatchKey, Args... args) const {
    // TODO Remove dispatchKey argument and instead infer dispatchKey from args...
    #if !defined(__clang__) && !defined(_MSC_VER) && defined(__GNUC__) && __GNUC__ < 5
      // GCC 4 has issues with parameter packs inside lambdas, let's instead
      // return the KernelFunction from the lambda. Note: This copies the
      // KernelFunction and is slow, but it's the only way to make it work for
      // GCC 4.
      KernelFunction func = dispatchTable_.read([&] (const DispatchTable& dispatchTable) -> KernelFunction {
          return dispatchTable.lookup(dispatchKey);
      });
      return func.callUnboxed<Return, Args...>(std::forward<Args>(args)...);
    #else
      // For all other compilers and newer GCC, let's do it right
      return dispatchTable_.read([&] (const DispatchTable& dispatchTable) -> Return {
          return dispatchTable.lookup(dispatchKey)
              .callUnboxed<Return, Args...>(std::forward<Args>(args)...);
      });
    #endif
  }

  template<class Return, class... Args>
  Return callUnboxedOnly(TensorTypeId dispatchKey, Args... args) const {
    // TODO Remove dispatchKey argument and instead infer dispatchKey from args...
    #if !defined(__clang__) && !defined(_MSC_VER) && defined(__GNUC__) && __GNUC__ < 5
      // GCC 4 has issues with parameter packs inside lambdas, let's instead
      // return the KernelFunction from the lambda. Note: This copies the
      // KernelFunction and is slow, but it's the only way to make it work for
      // GCC 4.
      KernelFunction func = dispatchTable_.read([&] (const DispatchTable& dispatchTable) -> KernelFunction {
          return dispatchTable.lookup(dispatchKey);
      });
      return func.callUnboxedOnly<Return, Args...>(std::forward<Args>(args)...);
    #else
      // For all other compilers and newer GCC, let's do it right
      return dispatchTable_.read([&] (const DispatchTable& dispatchTable) -> Return {
          return dispatchTable.lookup(dispatchKey)
              .callUnboxedOnly<Return, Args...>(std::forward<Args>(args)...);
      });
    #endif
  }

  void callBoxed(Stack* stack) const {
    return dispatchTable_.read([&] (const DispatchTable& dispatchTable) {
        dispatchTable.lookup(stack).callBoxed(stack);
    });
  }

  void prepareForDeregistration();

  RegistrationHandleRAII registerKernel(TensorTypeId dispatch_key, KernelFunction kernel);
  RegistrationHandleRAII registerCatchallKernel(KernelFunction kernel);

  const OperatorOptions& options() {
    return options_;
  }

private:
  void deregisterKernel_(TensorTypeId dispatch_key, std::list<KernelFunction>::iterator kernel);
  void deregisterCatchallKernel_(std::list<KernelFunction>::iterator kernel);

  FunctionSchema schema_;

  // The dispatchTable stores the current kernel for each dispatch key
  LeftRight<DispatchTable> dispatchTable_;

  // kernels_ stores all registered kernels for the corresponding dispatch key
  // and catchAllKernels_ stores the catch-all kernels.
  // If an operator library gets loaded that overwrites an already existing kernel,
  // both kernels will be in that list but only the newer one will be in
  // dispatchTable. If any of the kernels go away (say the library gets
  // unloaded), we remove the kernel from this list and update the
  // dispatchTable if necessary.
  // Kernels in the list are ordered by registration time descendingly,
  // newer registrations are before older registrations.
  // We do not combine dispatchTable and kernels into one hash map because
  // kernels is a larger data structure and accessed quite infrequently
  // while dispatchTable is accessed often and should be kept small to fit
  // into CPU caches.
  // Invariants:
  //  - dispatchTable[dispatch_key] == kernels_[dispatch_key].front()
  //  - dispatchTable[dispatch_key] does not exist if and only if
  //    kernels_[dispatch_key] does not exist
  //  - If kernels_[dispatch_key] exists, then it has elements.
  //    It is never an empty list.
  // Analogous invariants for catchAllKernels_.
  //
  // Why do we do that?
  // -----
  // We mostly do this to enable Jupyter notebooks where a cell registering
  // a kernel could be executed multiple times and the later execution
  // should overwrite the earlier one. Note that this still fails when the
  // function schema changed between the executions, but it works as long
  // as the function schema didn't change. A better solution would be to
  // unload the old extension library from the Jupyter cell when the cell is
  // re-executed and then only allow one kernel here, i.e. error if a kernel
  // is already registered, but that's a lot of effort to implement and
  // currently not high-pri.
  ska::flat_hash_map<TensorTypeId, std::list<KernelFunction>> kernels_;
  std::list<KernelFunction> catchAllKernels_;

  // Some metadata about the operator
  OperatorOptions options_;

  std::mutex kernelsMutex_; // protects kernels_

  // This function re-establishes the invariant that dispatchTable
  // contains the front element from the kernels list for a given dispatch key.
  void updateDispatchTable_(TensorTypeId dispatch_key);
  void updateCatchallDispatchTable_();
};

}
}