/*
 * Copyright (c) Meta Platforms, Inc. and affiliates.
 *
 * This source code is licensed under the MIT license found in the
 * LICENSE file in the root directory of this source tree.
 */

#pragma once

#include <react/bridging/Base.h>
#include <react/bridging/CallbackWrapper.h>

#include <ReactCommon/CallInvoker.h>
#include <ReactCommon/SchedulerPriority.h>

namespace facebook::react {

template <typename F>
class SyncCallback;

template <typename... Args>
class AsyncCallback {
 public:
  AsyncCallback(
      jsi::Runtime& runtime,
      jsi::Function function,
      std::shared_ptr<CallInvoker> jsInvoker)
      : callback_(std::make_shared<SyncCallback<void(Args...)>>(
            runtime,
            std::move(function),
            std::move(jsInvoker))) {}

  void operator()(Args... args) const noexcept {
    call(std::forward<Args>(args)...);
  }

  void call(Args... args) const noexcept {
    callWithArgs(std::nullopt, std::forward<Args>(args)...);
  }

  void callWithPriority(SchedulerPriority priority, Args... args)
      const noexcept {
    callWithArgs(priority, std::forward<Args>(args)...);
  }

  void call(std::function<void(jsi::Runtime&, jsi::Function&)>&& callImpl)
      const noexcept {
    callWithFunction(std::nullopt, std::move(callImpl));
  }

  void callWithPriority(
      SchedulerPriority priority,
      std::function<void(jsi::Runtime&, jsi::Function&)>&& callImpl)
      const noexcept {
    callWithFunction(priority, std::move(callImpl));
  }

 private:
  friend Bridging<AsyncCallback>;

  std::shared_ptr<SyncCallback<void(Args...)>> callback_;

  void callWithArgs(std::optional<SchedulerPriority> priority, Args... args)
      const noexcept {
    if (auto wrapper = callback_->wrapper_.lock()) {
      auto fn = [callback = callback_,
                 argsPtr = std::make_shared<std::tuple<Args...>>(
                     std::make_tuple(std::forward<Args>(args)...))](
                    jsi::Runtime&) { callback->apply(std::move(*argsPtr)); };

      auto& jsInvoker = wrapper->jsInvoker();
      if (priority) {
        jsInvoker.invokeAsync(*priority, std::move(fn));
      } else {
        jsInvoker.invokeAsync(std::move(fn));
      }
    }
  }

  void callWithFunction(
      std::optional<SchedulerPriority> priority,
      std::function<void(jsi::Runtime&, jsi::Function&)>&& callImpl)
      const noexcept {
    if (auto wrapper = callback_->wrapper_.lock()) {
      // Capture callback_ and not wrapper_. If callback_ is deallocated or the
      // JSVM is shutdown before the async task is scheduled, the underlying
      // function will have been deallocated.
      auto fn = [callback = callback_,
                 callImpl = std::move(callImpl)](jsi::Runtime& rt) {
        if (auto wrapper2 = callback->wrapper_.lock()) {
          callImpl(rt, wrapper2->callback());
        }
      };

      auto& jsInvoker = wrapper->jsInvoker();
      if (priority) {
        jsInvoker.invokeAsync(*priority, std::move(fn));
      } else {
        jsInvoker.invokeAsync(std::move(fn));
      }
    }
  }
};

// You must ensure that when invoking this you're located on the JS thread, or
// have exclusive control of the JS VM context. If you cannot ensure this, use
// AsyncCallback instead.
template <typename R, typename... Args>
class SyncCallback<R(Args...)> {
 public:
  SyncCallback(
      jsi::Runtime& rt,
      jsi::Function function,
      std::shared_ptr<CallInvoker> jsInvoker)
      : wrapper_(CallbackWrapper::createWeak(
            std::move(function),
            rt,
            std::move(jsInvoker))) {}

  // Disallow copying, as we can no longer safely destroy the callback
  // from the destructor if there's multiple copies
  SyncCallback(const SyncCallback&) = delete;
  SyncCallback& operator=(const SyncCallback&) = delete;

  // Allow move
  SyncCallback(SyncCallback&& other) noexcept
      : wrapper_(std::move(other.wrapper_)) {}

  SyncCallback& operator=(SyncCallback&& other) noexcept {
    wrapper_ = std::move(other.wrapper_);
    return *this;
  }

  ~SyncCallback() {
    if (auto wrapper = wrapper_.lock()) {
      wrapper->destroy();
    }
  }

  R operator()(Args... args) const {
    return call(std::forward<Args>(args)...);
  }

  R call(Args... args) const {
    auto wrapper = wrapper_.lock();

    // If the wrapper has been deallocated, we can no longer provide a return
    // value consistently, so our only option is to throw
    if (!wrapper) {
      if constexpr (std::is_void_v<R>) {
        return;
      } else {
        throw std::runtime_error("Failed to call invalidated sync callback");
      }
    }

    auto& callback = wrapper->callback();
    auto& rt = wrapper->runtime();
    auto jsInvoker = wrapper->jsInvokerPtr();

    if constexpr (std::is_void_v<R>) {
      callback.call(
          rt, bridging::toJs(rt, std::forward<Args>(args), jsInvoker)...);
    } else {
      return bridging::fromJs<R>(
          rt,
          callback.call(
              rt, bridging::toJs(rt, std::forward<Args>(args), jsInvoker)...),
          jsInvoker);
    }
  }

 private:
  friend AsyncCallback<Args...>;
  friend Bridging<SyncCallback>;

  R apply(std::tuple<Args...>&& args) const {
    return apply(std::move(args), std::index_sequence_for<Args...>{});
  }

  template <size_t... Index>
  R apply(std::tuple<Args...>&& args, std::index_sequence<Index...>) const {
    return call(std::move(std::get<Index>(args))...);
  }

  // Held weakly so lifetime is managed by LongLivedObjectCollection.
  std::weak_ptr<CallbackWrapper> wrapper_;
};

template <typename... Args>
struct Bridging<AsyncCallback<Args...>> {
  static AsyncCallback<Args...> fromJs(
      jsi::Runtime& rt,
      jsi::Function&& value,
      const std::shared_ptr<CallInvoker>& jsInvoker) {
    return AsyncCallback<Args...>(rt, std::move(value), jsInvoker);
  }

  static jsi::Function toJs(
      jsi::Runtime& rt,
      const AsyncCallback<Args...>& value) {
    return value.callback_->function_.getFunction(rt);
  }
};

template <typename R, typename... Args>
struct Bridging<SyncCallback<R(Args...)>> {
  static SyncCallback<R(Args...)> fromJs(
      jsi::Runtime& rt,
      jsi::Function&& value,
      const std::shared_ptr<CallInvoker>& jsInvoker) {
    return SyncCallback<R(Args...)>(rt, std::move(value), jsInvoker);
  }

  static jsi::Function toJs(
      jsi::Runtime& rt,
      const SyncCallback<R(Args...)>& value) {
    return value.function_.getFunction(rt);
  }
};

template <typename R, typename... Args>
struct Bridging<std::function<R(Args...)>> {
  using Func = std::function<R(Args...)>;
  using IndexSequence = std::index_sequence_for<Args...>;

  static constexpr size_t kArgumentCount = sizeof...(Args);

  static jsi::Function toJs(
      jsi::Runtime& rt,
      Func fn,
      const std::shared_ptr<CallInvoker>& jsInvoker) {
    return jsi::Function::createFromHostFunction(
        rt,
        jsi::PropNameID::forAscii(rt, "BridgedFunction"),
        kArgumentCount,
        [fn = std::make_shared<Func>(std::move(fn)), jsInvoker](
            jsi::Runtime& rt,
            const jsi::Value&,
            const jsi::Value* args,
            size_t count) -> jsi::Value {
          if (count < kArgumentCount) {
            throw jsi::JSError(rt, "Incorrect number of arguments");
          }

          if constexpr (std::is_void_v<R>) {
            callFromJs(*fn, rt, args, jsInvoker, IndexSequence{});
            return jsi::Value();
          } else {
            return bridging::toJs(
                rt,
                callFromJs(*fn, rt, args, jsInvoker, IndexSequence{}),
                jsInvoker);
          }
        });
  }

 private:
  template <size_t... Index>
  static R callFromJs(
      Func& fn,
      jsi::Runtime& rt,
      const jsi::Value* args,
      const std::shared_ptr<CallInvoker>& jsInvoker,
      std::index_sequence<Index...>) {
    return fn(bridging::fromJs<Args>(rt, args[Index], jsInvoker)...);
  }
};

template <typename R, typename... Args>
struct Bridging<
    std::function<R(Args...)>,
    std::enable_if_t<
        !std::is_same_v<std::function<R(Args...)>, std::function<R(Args...)>>>>
    : Bridging<std::function<R(Args...)>> {};

template <typename R, typename... Args>
struct Bridging<R(Args...)> : Bridging<std::function<R(Args...)>> {};

template <typename R, typename... Args>
struct Bridging<R (*)(Args...)> : Bridging<std::function<R(Args...)>> {};

} // namespace facebook::react