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

namespace TensorOps {

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

  if (info.Length() < 1) {
    Napi::TypeError::New(env, "Expected device string or options object").ThrowAsJavaScriptException();
    return env.Null();
  }

  torch::TensorOptions options;
  bool has_options = false;

  // Case 1: .to("cuda") or .to("mps")
  if (info[0].IsString()) {
    std::string device_str = info[0].As<Napi::String>().Utf8Value();
    options = torch::TensorOptions().device(TyTorchUtils::ParseDevice(device_str));
    has_options = true;

    // Case 2: .to("cuda", "float64")
    if (info.Length() > 1 && info[1].IsString()) {
      std::string dtype_str = info[1].As<Napi::String>().Utf8Value();
      options = options.dtype(TyTorchUtils::ParseDtype(dtype_str));
    }
  }
  // Case 3: .to({ device: "cuda", dtype: "float64" })
  else if (info[0].IsObject()) {
    Napi::Object opts = info[0].As<Napi::Object>();

    if (opts.Has("device")) {
      Napi::Value device_val = opts.Get("device");
      if (device_val.IsString()) {
        std::string device_str = device_val.As<Napi::String>().Utf8Value();
        options = torch::TensorOptions().device(TyTorchUtils::ParseDevice(device_str));
        has_options = true;
      }
    }

    if (opts.Has("dtype")) {
      Napi::Value dtype_val = opts.Get("dtype");
      if (dtype_val.IsString()) {
        std::string dtype_str = dtype_val.As<Napi::String>().Utf8Value();
        if (has_options) {
          options = options.dtype(TyTorchUtils::ParseDtype(dtype_str));
        } else {
          options = torch::TensorOptions().dtype(TyTorchUtils::ParseDtype(dtype_str));
          has_options = true;
        }
      }
    }
  }

  if (!has_options) {
    Napi::TypeError::New(env, "Invalid arguments to .to()").ThrowAsJavaScriptException();
    return env.Null();
  }

  torch::Tensor result = self->tensor.to(options);
  return Tensor::NewInstance(env, result);
}

}  // namespace TensorOps