#include <jsi/jsi.h>
#include <worklets/AnimationFrameQueue/AnimationFrameBatchinator.h>
#include <worklets/SharedItems/Serializable.h>
#include <worklets/WorkletRuntime/RuntimeBindings.h>

#include <atomic>
#include <functional>
#include <memory>
#include <mutex>
#include <utility>
#include <vector>

namespace worklets {

void AnimationFrameBatchinator::addToBatch(jsi::Function &&callback) {
  {
    std::lock_guard<std::mutex> lock(callbacksMutex_);
    callbacks_.push_back(std::make_shared<const jsi::Function>(std::move(callback)));
  }
  flush();
}

AnimationFrameBatchinator::JsiRequestAnimationFrame AnimationFrameBatchinator::getJsiRequestAnimationFrame() {
  return [weakThis = weak_from_this()](facebook::jsi::Runtime &rt, const facebook::jsi::Value &callback) {
    const auto strongThis = weakThis.lock();
    if (!strongThis) {
      return;
    }

    strongThis->addToBatch(callback.asObject(rt).asFunction(rt));
  };
}

void AnimationFrameBatchinator::flush() {
  if (flushRequested_.exchange(true)) {
    return;
  }

  requestAnimationFrame_([weakThis = weak_from_this()](double timestampMs) {
    const auto strongThis = weakThis.lock();
    if (!strongThis) {
      return;
    }

    auto callbacks = strongThis->pullCallbacks();
    strongThis->flushRequested_ = false;

    auto &uiWorkletRuntime = strongThis->uiWorkletRuntime_;
    for (const auto &callback : callbacks) {
      uiWorkletRuntime->runSync(*callback, timestampMs);
    }
  });
}

std::vector<std::shared_ptr<const jsi::Function>> AnimationFrameBatchinator::pullCallbacks() {
  std::lock_guard<std::mutex> lock(callbacksMutex_);
  return std::move(callbacks_);
}

AnimationFrameBatchinator::AnimationFrameBatchinator(
    const std::shared_ptr<WorkletRuntime> &uiWorkletRuntime,
    RuntimeBindings::RequestAnimationFrame requestAnimationFrame)
    : uiWorkletRuntime_(uiWorkletRuntime), requestAnimationFrame_(std::move(requestAnimationFrame)) {}

} // namespace worklets