#include "no_grad.h"
#include <torch/torch.h>

namespace TensorOps {

Napi::Value NoGrad(const Napi::CallbackInfo& info) {
  Napi::Env env = info.Env();

  try {
    // Expect a callback function as the argument
    if (info.Length() < 1 || !info[0].IsFunction()) {
      Napi::TypeError::New(env, "Expected a callback function").ThrowAsJavaScriptException();
      return env.Undefined();
    }

    Napi::Function callback = info[0].As<Napi::Function>();

    // Create NoGradGuard scope
    {
      torch::NoGradGuard no_grad;
      // Execute the callback within the no_grad scope
      callback.Call(env.Global(), {});
    }
    // Guard is destroyed here, re-enabling gradient tracking

  } catch (const std::exception& e) {
    Napi::Error::New(env, e.what()).ThrowAsJavaScriptException();
    return env.Undefined();
  }

  return env.Undefined();
}

} // namespace TensorOps