#include "reshape.h"
#include "../tensor.h"
#include "../utils.h"
#include <torch/torch.h>

namespace TensorOps {

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

  if (info.Length() < 1) {
    Napi::TypeError::New(env, "Expected shape array as argument").ThrowAsJavaScriptException();
    return env.Null();
  }

  if (!info[0].IsArray()) {
    Napi::TypeError::New(env, "Expected shape to be an array of numbers").ThrowAsJavaScriptException();
    return env.Null();
  }

  Napi::Array shape_array = info[0].As<Napi::Array>();
  std::vector<int64_t> new_shape;

  for (uint32_t i = 0; i < shape_array.Length(); i++) {
    Napi::Value val = shape_array[i];
    if (!val.IsNumber()) {
      Napi::TypeError::New(env, "Shape array must contain only numbers").ThrowAsJavaScriptException();
      return env.Null();
    }
    new_shape.push_back(val.As<Napi::Number>().Int64Value());
  }

  // Use PyTorch's reshape method
  torch::Tensor result = self->tensor.reshape(new_shape);
  return Tensor::NewInstance(env, result);
}

}  // namespace TensorOps