#include "requires_grad.h"
#include "../tensor.h"
#include <torch/torch.h>

namespace TensorOps {

Napi::Value GetRequiresGrad(Tensor* self, const Napi::CallbackInfo& info) {
  Napi::Env env = info.Env();

  // Call PyTorch's requires_grad() method
  bool requires_grad = self->tensor.requires_grad();

  return Napi::Boolean::New(env, requires_grad);
}

Napi::Value SetRequiresGrad(Tensor* self, const Napi::CallbackInfo& info) {
  Napi::Env env = info.Env();

  if (info.Length() < 1) {
    Napi::TypeError::New(env, "Expected boolean argument").ThrowAsJavaScriptException();
    return env.Undefined();
  }

  if (!info[0].IsBoolean()) {
    Napi::TypeError::New(env, "Argument must be a boolean").ThrowAsJavaScriptException();
    return env.Undefined();
  }

  bool requires_grad = info[0].As<Napi::Boolean>().Value();

  // Call PyTorch's set_requires_grad() method (modifies in-place)
  self->tensor.set_requires_grad(requires_grad);

  return env.Undefined();
}

}  // namespace TensorOps