#pragma once

#include <jsi/jsi.h>

#include <memory>
#include <unordered_map>
#include <utility>

#include "RuntimeLifecycleMonitor.h"

namespace RNJsi {

namespace jsi = facebook::jsi;

class BaseRuntimeAwareCache {
public:
  static void setMainJsRuntime(jsi::Runtime *rt) { _mainRuntime = rt; }

protected:
  static jsi::Runtime *getMainJsRuntime() {
    assert(_mainRuntime != nullptr &&
           "Expected main Javascript runtime to be set in the "
           "BaseRuntimeAwareCache class.");

    return _mainRuntime;
  }

private:
  static jsi::Runtime *_mainRuntime;
};

/**
 * Provides a way to keep data specific to a jsi::Runtime instance that gets
 * cleaned up when that runtime is destroyed. This is necessary because JSI does
 * not allow for its associated objects to be retained past the runtime
 * lifetime. If an object (e.g. jsi::Values or jsi::Function instances) is kept
 * after the runtime is torn down, its destructor (once it is destroyed
 * eventually) will result in a crash (JSI objects keep a pointer to memory
 * managed by the runtime, accessing that portion of the memory after runtime is
 * deleted is the root cause of that crash).
 *
 * In order to provide an efficient implementation that does not add an overhead
 * for the cases when only a single runtiome is used, which is the primary
 * usecase, the following assumption has been made: Only for secondary runtimes
 * we track destruction and clean up the store associated with that runtime. For
 * the first runtime we assume that the object holding the store is destroyed
 * prior to the destruction of that runtime.
 *
 * The above assumption makes it work without any overhead when only single
 * runtime is in use. Specifically, we don't perform any additional operations
 * related to tracking runtime lifecycle when only a single runtime is used.
 */
template <typename T>
class RuntimeAwareCache : public BaseRuntimeAwareCache,
                          public RuntimeLifecycleListener {

public:
  void onRuntimeDestroyed(jsi::Runtime *rt) override {
    if (getMainJsRuntime() != rt) {
      // We are removing a secondary runtime
      _secondaryRuntimeCaches.erase(rt);
    }
  }

  ~RuntimeAwareCache() {
    for (auto &cache : _secondaryRuntimeCaches) {
      RuntimeLifecycleMonitor::removeListener(
          *static_cast<jsi::Runtime *>(cache.first), this);
    }
  }

  T &get(jsi::Runtime &rt) {
    // We check if we're accessing the main runtime - this is the happy path
    // to avoid us having to lookup by runtime for caches that only has a single
    // runtime
    if (getMainJsRuntime() == &rt) {
      return _primaryCache;
    } else {
      if (_secondaryRuntimeCaches.count(&rt) == 0) {
        // we only add listener when the secondary runtime is used, this assumes
        // that the secondary runtime is terminated first. This lets us avoid
        // additional complexity for the majority of cases when objects are not
        // shared between runtimes. Otherwise we'd have to register all objecrts
        // with the RuntimeMonitor as opposed to only registering ones that are
        // used in secondary runtime. Note that we can't register listener here
        // with the primary runtime as it may run on a separate thread.
        RuntimeLifecycleMonitor::addListener(rt, this);

        T cache;
        _secondaryRuntimeCaches.emplace(&rt, std::move(cache));
      }
    }
    return _secondaryRuntimeCaches.at(&rt);
  }

private:
  std::unordered_map<void *, T> _secondaryRuntimeCaches;
  T _primaryCache;
};

} // namespace RNJsi