#include "tensor.h"
#include "utils.h"
#include <sstream>

Napi::FunctionReference Tensor::constructor;

Napi::Object Tensor::Init(Napi::Env env, Napi::Object exports) {
  Napi::Function func = DefineClass(env, "Tensor", {
    // Arithmetic operations
    InstanceMethod("add", &Tensor::Add),
    InstanceMethod("sub", &Tensor::Sub),
    InstanceMethod("mul", &Tensor::Mul),
    InstanceMethod("div", &Tensor::Div),

    // In-place operations
    InstanceMethod("add_", &Tensor::AddInplace),
    InstanceMethod("sub_", &Tensor::SubInplace),
    InstanceMethod("mul_", &Tensor::MulInplace),
    InstanceMethod("div_", &Tensor::DivInplace),

    // Matrix operations
    InstanceMethod("matmul", &Tensor::Matmul),

    // Reductions
    InstanceMethod("sum", &Tensor::Sum),
    InstanceMethod("mean", &Tensor::Mean),

    // Conversions
    InstanceMethod("toString", &Tensor::ToString),
    InstanceMethod("toArray", &Tensor::ToArray),
    InstanceMethod("to", &Tensor::To),

    // Properties
    InstanceMethod("shape", &Tensor::Shape),
    InstanceMethod("dtype", &Tensor::Dtype),
    InstanceMethod("device", &Tensor::Device),

    // Device management
    InstanceMethod("cpu", &Tensor::Cpu),
    InstanceMethod("cuda", &Tensor::Cuda),
    InstanceMethod("mps", &Tensor::Mps),

    // Dtype shortcuts
    InstanceMethod("float", &Tensor::Float),
    InstanceMethod("double", &Tensor::Double),
    InstanceMethod("int", &Tensor::Int),
    InstanceMethod("long", &Tensor::Long),

    // Shape operations
    InstanceMethod("reshape", &Tensor::Reshape),
    InstanceMethod("flatten", &Tensor::Flatten),
    InstanceMethod("unsqueeze", &Tensor::Unsqueeze),
    InstanceMethod("squeeze", &Tensor::Squeeze),
    InstanceMethod("transpose", &Tensor::Transpose),
    InstanceMethod("permute", &Tensor::Permute),

    // Autograd operations
    InstanceMethod("getRequiresGrad", &Tensor::GetRequiresGrad),
    InstanceMethod("setRequiresGrad", &Tensor::SetRequiresGrad),
    InstanceMethod("backward", &Tensor::Backward),
    InstanceMethod("getGrad", &Tensor::GetGrad),
    InstanceMethod("zeroGrad", &Tensor::ZeroGrad),
    InstanceMethod("detach", &Tensor::Detach),

    // Activation functions
    InstanceMethod("relu", &Tensor::Relu),
    InstanceMethod("sigmoid", &Tensor::Sigmoid),
    InstanceMethod("tanh", &Tensor::Tanh),
    InstanceMethod("softmax", &Tensor::Softmax),
    InstanceMethod("log_softmax", &Tensor::LogSoftmax),

    // Loss functions
    InstanceMethod("mse_loss", &Tensor::MseLoss),
    InstanceMethod("cross_entropy", &Tensor::CrossEntropy),
    InstanceMethod("nll_loss", &Tensor::NllLoss),
    InstanceMethod("binary_cross_entropy", &Tensor::BinaryCrossEntropy),
  });

  constructor = Napi::Persistent(func);
  constructor.SuppressDestruct();

  exports.Set("Tensor", func);
  return exports;
}

Napi::Value Tensor::NewInstance(Napi::Env env, torch::Tensor tensor) {
  Napi::Object obj = constructor.New({});
  Tensor* t = Napi::ObjectWrap<Tensor>::Unwrap(obj);
  t->tensor = tensor;
  return obj;
}

Tensor::Tensor(const Napi::CallbackInfo& info) : Napi::ObjectWrap<Tensor>(info) {
  Napi::Env env = info.Env();

  // Constructor for creating tensor from array with options
  if (info.Length() > 0 && info[0].IsArray()) {
    Napi::Array arr = info[0].As<Napi::Array>();
    std::vector<float> data;

    for (uint32_t i = 0; i < arr.Length(); i++) {
      Napi::Value val = arr[i];
      if (val.IsNumber()) {
        data.push_back(val.As<Napi::Number>().FloatValue());
      }
    }

    // Create tensor with default options
    torch::TensorOptions options = torch::TensorOptions().dtype(torch::kFloat32);

    // Parse options if provided
    if (info.Length() > 1 && info[1].IsObject()) {
      Napi::Object opts = info[1].As<Napi::Object>();

      // Device option
      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 = options.device(TyTorchUtils::ParseDevice(device_str));
        }
      }

      // Dtype option
      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();
          options = options.dtype(TyTorchUtils::ParseDtype(dtype_str));
        }
      }
    }

    this->tensor = torch::from_blob(data.data(), {static_cast<int64_t>(data.size())}, torch::kFloat32).clone().to(options);

    // Set requires_grad after tensor creation (cannot be set in TensorOptions)
    if (info.Length() > 1 && info[1].IsObject()) {
      Napi::Object opts = info[1].As<Napi::Object>();
      if (opts.Has("requires_grad")) {
        Napi::Value requires_grad_val = opts.Get("requires_grad");
        if (requires_grad_val.IsBoolean()) {
          bool requires_grad = requires_grad_val.As<Napi::Boolean>().Value();
          this->tensor.set_requires_grad(requires_grad);
        }
      }
    }
  } else {
    // Default: empty tensor
    this->tensor = torch::Tensor();
  }
}

// Arithmetic operations with scalar and tensor support
// Add implementation moved to ops/add.cpp

Napi::Value Tensor::Sub(const Napi::CallbackInfo& info) {
  // Implementation moved to ops/sub.cpp
  return TensorOps::Sub(this, info);
}

Napi::Value Tensor::Mul(const Napi::CallbackInfo& info) {
  // Implementation moved to ops/mul.cpp
  return TensorOps::Mul(this, info);
}

Napi::Value Tensor::Div(const Napi::CallbackInfo& info) {
  // Implementation moved to ops/div.cpp
  return TensorOps::Div(this, info);
}

// In-place operations
// AddInplace implementation moved to ops/add_.cpp

Napi::Value Tensor::SubInplace(const Napi::CallbackInfo& info) {
  // Implementation moved to ops/sub_.cpp
  return TensorOps::SubInplace(this, info);
}

Napi::Value Tensor::MulInplace(const Napi::CallbackInfo& info) {
  // Implementation moved to ops/mul_.cpp
  return TensorOps::MulInplace(this, info);
}

Napi::Value Tensor::DivInplace(const Napi::CallbackInfo& info) {
  // Implementation moved to ops/div_.cpp
  return TensorOps::DivInplace(this, info);
}

Napi::Value Tensor::Matmul(const Napi::CallbackInfo& info) {
  // Implementation moved to ops/matmul.cpp
  return TensorOps::Matmul(this, info);
}

Napi::Value Tensor::Sum(const Napi::CallbackInfo& info) {
  // Implementation moved to ops/sum.cpp
  return TensorOps::Sum(this, info);
}

Napi::Value Tensor::Mean(const Napi::CallbackInfo& info) {
  // Implementation moved to ops/mean.cpp
  return TensorOps::Mean(this, info);
}

// Conversion: .to() method
Napi::Value Tensor::To(const Napi::CallbackInfo& info) {
  // Implementation moved to ops/to.cpp
  return TensorOps::To(this, info);
}

Napi::Value Tensor::ToString(const Napi::CallbackInfo& info) {
  // Implementation moved to ops/toString.cpp
  return TensorOps::ToString(this, info);
}

Napi::Value Tensor::ToArray(const Napi::CallbackInfo& info) {
  // Implementation moved to ops/toArray.cpp
  return TensorOps::ToArray(this, info);
}

Napi::Value Tensor::Shape(const Napi::CallbackInfo& info) {
  // Implementation moved to ops/shape.cpp
  return TensorOps::Shape(this, info);
}

Napi::Value Tensor::Dtype(const Napi::CallbackInfo& info) {
  // Implementation moved to ops/dtype.cpp
  return TensorOps::Dtype(this, info);
}

Napi::Value Tensor::Device(const Napi::CallbackInfo& info) {
  // Implementation moved to ops/device.cpp
  return TensorOps::Device(this, info);
}

// Device management
Napi::Value Tensor::Cpu(const Napi::CallbackInfo& info) {
  // Implementation moved to ops/cpu.cpp
  return TensorOps::Cpu(this, info);
}

Napi::Value Tensor::Cuda(const Napi::CallbackInfo& info) {
  // Implementation moved to ops/cuda.cpp
  return TensorOps::Cuda(this, info);
}

Napi::Value Tensor::Mps(const Napi::CallbackInfo& info) {
  // Implementation moved to ops/mps.cpp
  return TensorOps::Mps(this, info);
}

// Dtype shortcuts
Napi::Value Tensor::Float(const Napi::CallbackInfo& info) {
  // Implementation moved to ops/float.cpp
  return TensorOps::Float(this, info);
}

Napi::Value Tensor::Double(const Napi::CallbackInfo& info) {
  // Implementation moved to ops/double.cpp
  return TensorOps::Double(this, info);
}

Napi::Value Tensor::Int(const Napi::CallbackInfo& info) {
  // Implementation moved to ops/int.cpp
  return TensorOps::Int(this, info);
}

Napi::Value Tensor::Long(const Napi::CallbackInfo& info) {
  // Implementation moved to ops/long.cpp
  return TensorOps::Long(this, info);
}

// Shape operations
Napi::Value Tensor::Reshape(const Napi::CallbackInfo& info) {
  // Implementation moved to ops/reshape.cpp
  return TensorOps::Reshape(this, info);
}

Napi::Value Tensor::Flatten(const Napi::CallbackInfo& info) {
  // Implementation moved to ops/flatten.cpp
  return TensorOps::Flatten(this, info);
}

Napi::Value Tensor::Unsqueeze(const Napi::CallbackInfo& info) {
  // Implementation moved to ops/unsqueeze.cpp
  return TensorOps::Unsqueeze(this, info);
}

Napi::Value Tensor::Squeeze(const Napi::CallbackInfo& info) {
  // Implementation moved to ops/squeeze.cpp
  return TensorOps::Squeeze(this, info);
}

Napi::Value Tensor::Transpose(const Napi::CallbackInfo& info) {
  // Implementation moved to ops/transpose.cpp
  return TensorOps::Transpose(this, info);
}

Napi::Value Tensor::Permute(const Napi::CallbackInfo& info) {
  // Implementation moved to ops/permute.cpp
  return TensorOps::Permute(this, info);
}

// Autograd operations
Napi::Value Tensor::GetRequiresGrad(const Napi::CallbackInfo& info) {
  // Implementation moved to ops/requiresGrad.cpp
  return TensorOps::GetRequiresGrad(this, info);
}

Napi::Value Tensor::SetRequiresGrad(const Napi::CallbackInfo& info) {
  // Implementation moved to ops/requires_grad.cpp
  return TensorOps::SetRequiresGrad(this, info);
}

Napi::Value Tensor::Backward(const Napi::CallbackInfo& info) {
  // Implementation moved to ops/backward.cpp
  return TensorOps::Backward(this, info);
}

Napi::Value Tensor::GetGrad(const Napi::CallbackInfo& info) {
  // Implementation moved to ops/grad.cpp
  return TensorOps::GetGrad(this, info);
}

Napi::Value Tensor::ZeroGrad(const Napi::CallbackInfo& info) {
  // Implementation moved to ops/zero_grad.cpp
  return TensorOps::ZeroGrad(this, info);
}

Napi::Value Tensor::Detach(const Napi::CallbackInfo& info) {
  // Implementation moved to ops/detach.cpp
  return TensorOps::Detach(this, info);
}

// Activation functions
Napi::Value Tensor::Relu(const Napi::CallbackInfo& info) {
  // Implementation moved to ops/relu.cpp
  return TensorOps::Relu(this, info);
}

Napi::Value Tensor::Sigmoid(const Napi::CallbackInfo& info) {
  // Implementation moved to ops/sigmoid.cpp
  return TensorOps::Sigmoid(this, info);
}

Napi::Value Tensor::Tanh(const Napi::CallbackInfo& info) {
  // Implementation moved to ops/tanh.cpp
  return TensorOps::Tanh(this, info);
}

Napi::Value Tensor::Softmax(const Napi::CallbackInfo& info) {
  // Implementation moved to ops/softmax.cpp
  return TensorOps::Softmax(this, info);
}

Napi::Value Tensor::LogSoftmax(const Napi::CallbackInfo& info) {
  // Implementation moved to ops/log_softmax.cpp
  return TensorOps::LogSoftmax(this, info);
}

Napi::Value Tensor::MseLoss(const Napi::CallbackInfo& info) {
  // Implementation moved to ops/mse_loss.cpp
  return TensorOps::MseLoss(this, info);
}

Napi::Value Tensor::CrossEntropy(const Napi::CallbackInfo& info) {
  // Implementation moved to ops/cross_entropy.cpp
  return TensorOps::CrossEntropy(this, info);
}

Napi::Value Tensor::NllLoss(const Napi::CallbackInfo& info) {
  // Implementation moved to ops/nll_loss.cpp
  return TensorOps::NllLoss(this, info);
}

Napi::Value Tensor::BinaryCrossEntropy(const Napi::CallbackInfo& info) {
  // Implementation moved to ops/binary_cross_entropy.cpp
  return TensorOps::BinaryCrossEntropy(this, info);
}