#pragma once

#include <vector>
#include <memory>
#include <lua.hpp>

#define ASSERT(...)
#define UNUSED(x) (void)(x)

#ifndef alignof
#define alignof(T) std::alignment_of<T>::value
#endif

namespace lua {

/**
 * Base macro for declaring a Lua/C function. Any function declared will need to be exported via
 * the exportedFunctions() function in BaseLuaObject. See BaseLuaObject below for more
 * information. This macro declares a static "thunk" which checks the user data, optionally checks
 * for object death (again see BaseLuaObject below for more info), and then invokes a normal
 * object method. The actual object method needs to be implemented by the class.
 * @param Class supplies the owning class name.
 * @param Name supplies the function name.
 * @param Index supplies the stack index where "this" (Lua/C userdata) is found.
 */
#define DECLARE_LUA_FUNCTION_EX(Class, Name, Index)                                                \
  static int static_##Name(lua_State* state) {                                                     \
    Class* object = lua::alignAndCast<Class>(                \
        luaL_checkudata(state, Index, typeid(Class).name()));                                      \
    object->checkDead(state);                                                                      \
    return object->Name(state);                                                                    \
  }                                                                                                \
  int Name(lua_State* state);

/**
 * Declare a Lua function in which userdata is in stack slot 1. See DECLARE_LUA_FUNCTION_EX()
 */
#define DECLARE_LUA_FUNCTION(Class, Name) DECLARE_LUA_FUNCTION_EX(Class, Name, 1)

/**
 * Declare a Lua function in which userdata is in upvalue slot 1. See DECLARE_LUA_FUNCTION_EX()
 */
#define DECLARE_LUA_CLOSURE(Class, Name) DECLARE_LUA_FUNCTION_EX(Class, Name, lua_upvalueindex(1))

/**
 * Calculate the maximum space needed to be aligned.
 */
template <typename T> size_t maximumSpaceNeededToAlign() {
  // The allocated memory can be misaligned up to `alignof(T) - 1` bytes. Adding it to the size to
  // allocate.
  return sizeof(T) + alignof(T) - 1;
}

template <typename T> inline T* alignAndCast(void* mem) {
  size_t size = maximumSpaceNeededToAlign<T>();
  return static_cast<T*>(std::align(alignof(T), sizeof(T), mem, size));
}

/**
 * Create a new user data and assign its metatable.
 */
template <typename T> inline T* allocateLuaUserData(lua_State* state) {
  void* mem = lua_newuserdata(state, maximumSpaceNeededToAlign<T>());
  luaL_getmetatable(state, typeid(T).name());
  ASSERT(lua_istable(state, -1));
  lua_setmetatable(state, -2);

  return alignAndCast<T>(mem);
}

/**
 * This is the base class for all C++ objects that we expose out to Lua. The goal is to hide as
 * much ugliness as possible. In general, to use this, do the following:
 * 1) Make your class derive from BaseLuaObject<YourClass>
 * 2) Define methods using DECLARE_LUA_FUNCTION* macros
 * 3) Export your functions by declaring a static exportedFunctions() method in your class.
 * 4) Optionally manage "death" status on your object. (See checkDead() and markDead() below).
 * 5) Generally you will want to hold your objects inside a LuaRef or a LuaDeathRef. See below
 *    for more information on those containers.
 *
 * It's very important to understand the Lua memory model: Once an object is created, *it is
 * owned by Lua*. Lua can GC it at any time. If you want to make sure that does not happen, you
 * must hold a ref to it in C++, generally via LuaRef or LuaDeathRef.
 */
template <class T> class BaseLuaObject {
public:
  using ExportedFunctions = std::vector<std::pair<const char*, lua_CFunction>>;

  virtual ~BaseLuaObject() = default;

  /**
   * Create a new object of this type, owned by Lua. This type must have previously been registered
   * via the registerType() routine below.
   * @param state supplies the owning Lua state.
   * @param args supplies the variadic constructor arguments for the object.
   * @return a pair containing a pointer to the new object and the state it was created with. (This
   *         is done for convenience when passing a created object to a LuaRef or a LuaDeathRef.
   */
  template <typename... ConstructorArgs>
  static std::pair<T*, lua_State*> create(lua_State* state, ConstructorArgs&&... args) {
    // Memory is allocated via Lua and it is raw. We use placement new to run the constructor.
    T* mem = allocateLuaUserData<T>(state);
    return {new (mem) T(std::forward<ConstructorArgs>(args)...), state};
  }

  /**
   * Register a type with Lua.
   * @param state supplies the state to register with.
   */
  static void registerType(lua_State* state) {
    std::vector<luaL_Reg> to_register;

    // Fetch all of the functions to be exported to Lua so that we can register them in the
    // metatable.
    ExportedFunctions functions = T::exportedFunctions();
    for (auto function : functions) {
      to_register.push_back({function.first, function.second});
    }

    // Always register a __gc method so that we can run the object's destructor. We do this
    // manually because the memory is raw and was allocated by Lua.
    to_register.push_back(
        {"__gc", [](lua_State* state) {
           T* object = alignAndCast<T>(luaL_checkudata(state, 1, typeid(T).name()));
           object->~T();
           return 0;
         }});

    // Add the sentinel.
    to_register.push_back({nullptr, nullptr});

    // Register the type by creating a new metatable, setting __index to itself, and then
    // performing the register.
    luaL_newmetatable(state, typeid(T).name());

    lua_pushvalue(state, -1);
    lua_setfield(state, -2, "__index");
    luaL_register(state, nullptr, to_register.data());
  }

  /**
   * This function is called as part of the DECLARE_LUA_FUNCTION* macros. The idea here is that
   * we cannot control when Lua destroys things. However, we may expose wrappers to a script that
   * should not be used after some event. This allows us to mark objects as dead so that if they
   * are used again they will throw a Lua error and not reach our code.
   * @param state supplies the calling LuaState.
   */
  int checkDead(lua_State* state) {
    if (dead_) {
      return luaL_error(state, "object used outside of proper scope");
    }
    return 0;
  }

  /**
   * Mark an object as dead so that a checkDead() call will throw an error. See checkDead().
   */
  void markDead() {
    dead_ = true;
    onMarkDead();
  }

  /**
   * Mark an object as live so that a checkDead() call will not throw an error. See checkDead().
   */
  void markLive() {
    dead_ = false;
    onMarkLive();
  }

protected:
  /**
   * Called from markDead() when an object is marked dead. This is effectively a C++ destructor for
   * Lua/C objects. Objects can perform inline cleanup or mark other objects as dead if needed. It
   * can also be used to protect objects from use if they get assigned to a global variable and
   * used across coroutines.
   */
  virtual void onMarkDead() {}

  /**
   * Called from markLive() when an object is marked live. This is a companion to onMarkDead(). See
   * the comments there.
   */
  virtual void onMarkLive() {}

private:
  bool dead_{};
};

/**
 * This is basically a Lua smart pointer. The idea is that given a Lua object, if we want to
 * guarantee that Lua won't destroy it, we need to reference it. This wraps the reference
 * functionality. While a LuaRef owns an object it's guaranteed that Lua will not GC it.
 * TODO(mattklein123): Add dedicated unit tests. This will require mocking a Lua state.
 */
template <typename T> class LuaRef {
public:
  /**
   * Create an empty LuaRef.
   */
  LuaRef() { reset(); }

  /**
   * Create a LuaRef from an object.
   * @param object supplies the object. Generally this is the return value from a Object::create()
   *        call. The object must be at the top of the Lua stack.
   * @param leave_on_stack supplies whether to leave the object on the stack or not when the ref
   *        is constructed.
   */
  LuaRef(const std::pair<T*, lua_State*>& object, bool leave_on_stack) {
    reset(object, leave_on_stack);
  }

  ~LuaRef() { unref(); }
  T* get() { return object_.first; }

  /**
   * Same as the LuaRef non-default constructor, but post-construction.
   */
  void reset(const std::pair<T*, lua_State*>& object, bool leave_on_stack) {
    unref();

    if (leave_on_stack) {
      lua_pushvalue(object.second, -1);
    }

    object_ = object;
    ref_ = luaL_ref(object_.second, LUA_REGISTRYINDEX);
    ASSERT(ref_ != LUA_REFNIL);
  }

  /**
   * Return a LuaRef to its default/empty state.
   */
  void reset() {
    unref();
    object_ = std::pair<T*, lua_State*>{};
    ref_ = LUA_NOREF;
  }

  /**
   * Push the referenced object back onto the stack.
   */
  void pushStack() {
    ASSERT(object_.first);
    lua_rawgeti(object_.second, LUA_REGISTRYINDEX, ref_);
  }

protected:
  void unref() {
    if (object_.second != nullptr) {
      luaL_unref(object_.second, LUA_REGISTRYINDEX, ref_);
    }
  }

  std::pair<T*, lua_State*> object_;
  int ref_;
};

/**
 * This is a variant of LuaRef which also marks an object as dead during destruction. This is
 * useful if an object should not be used after the scope of the pcall() or resume().
 * TODO(mattklein123): Add dedicated unit tests. This will require mocking a Lua state.
 */
template <typename T> class LuaDeathRef : public LuaRef<T> {
public:
  ~LuaDeathRef() { markDead(); }

  void markDead() {
    if (this->object_.first) {
      this->object_.first->markDead();
    }
  }

  void markLive() {
    if (this->object_.first) {
      this->object_.first->markLive();
    }
  }

  void reset(const std::pair<T*, lua_State*>& object, bool leave_on_stack) {
    markDead();
    LuaRef<T>::reset(object, leave_on_stack);
  }

  void reset() {
    markDead();
    LuaRef<T>::reset();
  }
};

}