// Copyright 2025 DeepMind Technologies Limited
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

// NOLINTBEGIN(whitespace/line_length)
// NOLINTBEGIN(whitespace/semicolon)

#include <emscripten.h>
#include <emscripten/bind.h>
#include <emscripten/val.h>

#include <array>
#include <cstddef>
#include <cstdint>
#include <cstdio>
#include <cstdlib>
#include <cstring>  // NOLINT
#include <memory>
#include <optional>  // NOLINT
#include <string>    // NOLINT
#include <vector>

#include <mujoco/mjmodel.h>
#include <mujoco/mjvisualize.h>
#include <mujoco/mujoco.h>
#include "engine/engine_util_errmem.h"
#include "wasm/unpack.h"

namespace mujoco::wasm {

using emscripten::enum_;
using emscripten::class_;
using emscripten::function;
using emscripten::val;
using emscripten::constant;
using emscripten::register_optional;
using emscripten::register_type;
using emscripten::register_vector;
using emscripten::return_value_policy::reference;
using emscripten::return_value_policy::take_ownership;

EMSCRIPTEN_DECLARE_VAL_TYPE(NumberArray);
EMSCRIPTEN_DECLARE_VAL_TYPE(String);

// Raises an error if the given val is null or undefined.
// A macro is used so that the error contains the name of the variable.
// TODO(matijak): Remove this when we can handle strings using UNPACK_STRING?
#define CHECK_VAL(val)                                    \
  if (val.isNull()) {                                     \
    mju_error("Invalid argument: %s is null", #val);      \
  } else if (val.isUndefined()) {                         \
    mju_error("Invalid argument: %s is undefined", #val); \
  }

void ThrowMujocoErrorToJS(const char* msg) {
  // Get a handle to the JS global Error constructor function, create a new
  // object instance and then throw the object as an exception using the
  // val::throw_() helper function.
  val(val::global("Error").new_(val("MuJoCo Error: " + std::string(msg))))
      .throw_();
}
__attribute__((constructor)) void InitMuJoCoErrorHandler() {
  mju_user_error = ThrowMujocoErrorToJS;
}

template <size_t N>
val MakeValArray(const char* (&strings)[N]) {
  val result = val::array();
  for (int i = 0; i < N; i++) {
    result.call<void>("push", val(strings[i]));
  }
  return result;
}

template <size_t N, size_t M>
val MakeValArray3(const char* (&strings)[N][M]) {
  val result = val::array();
  for (int i = 0; i < N; i++) {
    val inner = val::array();
    for (int j = 0; j < M; j++) {
      inner.call<void>("push", val(strings[i][j]));
    }
    result.call<void>("push", inner);
  }
  return result;
}

template <typename WrapperType, typename ArrayType, typename SizeType>
std::vector<WrapperType> InitWrapperArray(ArrayType* array, SizeType size) {
  std::vector<WrapperType> result;
  result.reserve(size);
  for (int i = 0; i < size; ++i) {
    result.emplace_back(&array[i]);
  }
  return result;
}

// {{ ANONYMOUS_STRUCT_TYPEDEFS }}

// {{ AUTOGENNED_STRUCTS_HEADER }}

struct MjVisual {
  MjVisual();
  explicit MjVisual(mjVisual *ptr_);
  MjVisual(const MjVisual &);
  MjVisual &operator=(const MjVisual &);
  ~MjVisual();
  std::unique_ptr<MjVisual> copy();
  mjVisual* get() const;
  void set(mjVisual* ptr);
  // INSERT-GENERATED-MjVisual-DEFINITIONS

 private:
  mjVisual* ptr_;
  bool owned_ = false;

 public:
  MjVisualGlobal global;
  MjVisualQuality quality;
  MjVisualHeadlight headlight;
  MjVisualMap map;
  MjVisualScale scale;
  MjVisualRgba rgba;
};

struct MjModel {
  explicit MjModel(mjModel *m);
  explicit MjModel(const MjModel &other);
  ~MjModel();
  std::unique_ptr<MjModel> copy();
  mjModel* get() const;
  void set(mjModel* ptr);
  // INSERT-GENERATED-MjModel-DEFINITIONS

 private:
  mjModel* ptr_;

 public:
  MjOption opt;
  MjStatistic stat;
  MjVisual vis;
};

struct MjData {
  MjData(MjModel *m);
  explicit MjData(const MjModel &, const MjData &);
  ~MjData();
  std::vector<MjContact> contact() const;
  std::unique_ptr<MjData> copy();
  mjData* get() const;
  void set(mjData* ptr);
  // INSERT-GENERATED-MjData-DEFINITIONS

 private:
  mjData* ptr_;

 public:
  mjModel *model;
  std::vector<MjSolverStat> solver;
  std::vector<MjTimerStat> timer;
  std::vector<MjWarningStat> warning;
};

struct MjvScene {
  MjvScene();
  MjvScene(MjModel *m, int maxgeom);
  // MjvScene(const MjvScene &);
  ~MjvScene();
  std::unique_ptr<MjvScene> copy();
  int GetSumFlexFaces() const;

  mjvScene* get() const;
  void set(mjvScene* ptr);

  std::vector<MjvGeom> geoms() const;

  emscripten::val geomorder() const {
    return emscripten::val(
        emscripten::typed_memory_view(ptr_->ngeom, ptr_->geomorder));
  }
  emscripten::val flexedgeadr() const {
    return emscripten::val(
        emscripten::typed_memory_view(ptr_->nflex, ptr_->flexedgeadr));
  }
  emscripten::val flexedgenum() const {
    return emscripten::val(
        emscripten::typed_memory_view(ptr_->nflex, ptr_->flexedgenum));
  }
  emscripten::val flexvertadr() const {
    return emscripten::val(
        emscripten::typed_memory_view(ptr_->nflex, ptr_->flexvertadr));
  }
  emscripten::val flexvertnum() const {
    return emscripten::val(
        emscripten::typed_memory_view(ptr_->nflex, ptr_->flexvertnum));
  }
  emscripten::val flexfaceadr() const {
    return emscripten::val(
        emscripten::typed_memory_view(ptr_->nflex, ptr_->flexfaceadr));
  }
  emscripten::val flexfacenum() const {
    return emscripten::val(
        emscripten::typed_memory_view(ptr_->nflex, ptr_->flexfacenum));
  }
  emscripten::val flexfaceused() const {
    return emscripten::val(
        emscripten::typed_memory_view(ptr_->nflex, ptr_->flexfaceused));
  }
  emscripten::val flexedge() const {
    return emscripten::val(
        emscripten::typed_memory_view(2 * model->nflexedge, ptr_->flexedge));
  }
  emscripten::val flexvert() const {
    return emscripten::val(
        emscripten::typed_memory_view(3 * model->nflexvert, ptr_->flexvert));
  }
  emscripten::val skinfacenum() const {
    return emscripten::val(
        emscripten::typed_memory_view(ptr_->nskin, ptr_->skinfacenum));
  }
  emscripten::val skinvertadr() const {
    return emscripten::val(
        emscripten::typed_memory_view(ptr_->nskin, ptr_->skinvertadr));
  }
  emscripten::val skinvertnum() const {
    return emscripten::val(
        emscripten::typed_memory_view(ptr_->nskin, ptr_->skinvertnum));
  }
  emscripten::val skinvert() const {
    return emscripten::val(
        emscripten::typed_memory_view(3 * model->nskinvert, ptr_->skinvert));
  }
  emscripten::val skinnormal() const {
    return emscripten::val(
        emscripten::typed_memory_view(3 * model->nskinvert, ptr_->skinnormal));
  }
  emscripten::val flexface() const {
    return emscripten::val(emscripten::typed_memory_view(
        9 * MjvScene::GetSumFlexFaces(), ptr_->flexface));
  }
  emscripten::val flexnormal() const {
    return emscripten::val(emscripten::typed_memory_view(
        9 * MjvScene::GetSumFlexFaces(), ptr_->flexnormal));
  }
  emscripten::val flextexcoord() const {
    return emscripten::val(emscripten::typed_memory_view(
        6 * MjvScene::GetSumFlexFaces(), ptr_->flextexcoord));
  }
  // INSERT-GENERATED-MjvScene-DEFINITIONS

 private:
  mjvScene* ptr_;
  bool owned_ = false;

 public:
  mjModel *model;
  std::vector<MjvLight> lights;
  std::vector<MjvGLCamera> camera;
};

struct MjSpec {
  MjSpec();
  explicit MjSpec(mjSpec *ptr);
  MjSpec(const MjSpec &);
  MjSpec &operator=(const MjSpec &);
  ~MjSpec();
  std::unique_ptr<MjSpec> copy();
  mjSpec* get() const;
  void set(mjSpec* ptr);
  // INSERT-GENERATED-MjSpec-DEFINITIONS

 private:
  mjSpec* ptr_;
  bool owned_ = false;

 public:
  MjOption option;
  MjVisual visual;
  MjStatistic stat;
  MjsCompiler compiler;
  MjsElement element;
};

val get_mjDISABLESTRING() { return MakeValArray(mjDISABLESTRING); }
val get_mjENABLESTRING() { return MakeValArray(mjENABLESTRING); }
val get_mjTIMERSTRING() { return MakeValArray(mjTIMERSTRING); }
val get_mjLABELSTRING() { return MakeValArray(mjLABELSTRING); }
val get_mjFRAMESTRING() { return MakeValArray(mjFRAMESTRING); }
val get_mjVISSTRING() { return MakeValArray3(mjVISSTRING); }
val get_mjRNDSTRING() { return MakeValArray3(mjRNDSTRING); }

EMSCRIPTEN_BINDINGS(mujoco_constants) {
  // from mjmodel.h
  constant("mjPI", mjPI);
  constant("mjMAXVAL", mjMAXVAL);
  constant("mjMINMU", mjMINMU);
  constant("mjMINIMP", mjMINIMP);
  constant("mjMAXIMP", mjMAXIMP);
  constant("mjMAXCONPAIR", mjMAXCONPAIR);
  constant("mjNEQDATA", mjNEQDATA);
  constant("mjNDYN", mjNDYN);
  constant("mjNGAIN", mjNGAIN);
  constant("mjNBIAS", mjNBIAS);
  constant("mjNREF", mjNREF);
  constant("mjNIMP", mjNIMP);
  constant("mjNSOLVER", mjNSOLVER);

  // from mjvisualize.h
  constant("mjNGROUP", mjNGROUP);
  constant("mjMAXLIGHT", mjMAXLIGHT);
  constant("mjMAXOVERLAY", mjMAXOVERLAY);
  constant("mjMAXLINE", mjMAXLINE);
  constant("mjMAXLINEPNT", mjMAXLINEPNT);
  constant("mjMAXPLANEGRID", mjMAXPLANEGRID);

  // from mujoco.h
  constant("mjVERSION_HEADER", mjVERSION_HEADER);

  // from mjtnum.h
  constant("mjMINVAL", mjMINVAL);

  // emscripten::constant() is designed for simple, compile-time literal values
  // (like numbers or a single string literal), complex values need to be
  // bound as functions.
  emscripten::function("get_mjDISABLESTRING", &get_mjDISABLESTRING);
  emscripten::function("get_mjENABLESTRING", &get_mjENABLESTRING);
  emscripten::function("get_mjTIMERSTRING", &get_mjTIMERSTRING);
  emscripten::function("get_mjLABELSTRING", &get_mjLABELSTRING);
  emscripten::function("get_mjFRAMESTRING", &get_mjFRAMESTRING);
  emscripten::function("get_mjVISSTRING", &get_mjVISSTRING);
  emscripten::function("get_mjRNDSTRING", &get_mjRNDSTRING);
}

EMSCRIPTEN_BINDINGS(mujoco_enums) {
// {{ ENUM_BINDINGS }}
}

// {{ AUTOGENNED_STRUCTS_SOURCE }}

// =============== MjModel =============== //
MjModel::MjModel(mjModel *m)
    : ptr_(m), opt(&m->opt), stat(&m->stat), vis(&m->vis) {}
MjModel::MjModel(const MjModel &other)
    : ptr_(mj_copyModel(nullptr, other.get())),
      opt(&ptr_->opt),
      stat(&ptr_->stat),
      vis(&ptr_->vis) {}
MjModel::~MjModel() {
  if (ptr_) {
    mj_deleteModel(ptr_);
  }
}
mjModel* MjModel::get() const { return ptr_; }
void MjModel::set(mjModel *ptr) { ptr_ = ptr; }

// TODO(manevi): Consider passing `const MjModel& m` here, mj_makeData uses a const model.
// =============== MjData =============== //
MjData::MjData(MjModel *m) {
  model = m->get();
  ptr_ = mj_makeData(model);
  if (ptr_) {
    solver = InitWrapperArray<MjSolverStat>(get()->solver, mjNSOLVER * mjNISLAND);
    timer = InitWrapperArray<MjTimerStat>(get()->timer, mjNTIMER);
    warning = InitWrapperArray<MjWarningStat>(get()->warning, mjNWARNING);
  }
}
MjData::MjData(const MjModel &model, const MjData &other)
    : ptr_(mj_copyData(nullptr, model.get(), other.get())), model(model.get()) {
  if (ptr_) {
    solver = InitWrapperArray<MjSolverStat>(get()->solver, mjNSOLVER * mjNISLAND);
    timer = InitWrapperArray<MjTimerStat>(get()->timer, mjNTIMER);
    warning = InitWrapperArray<MjWarningStat>(get()->warning, mjNWARNING);
  }
}
MjData::~MjData() {
  if (ptr_) {
    mj_deleteData(ptr_);
  }
}
mjData* MjData::get() const { return ptr_; }
void MjData::set(mjData *ptr) { ptr_ = ptr; }

std::vector<MjContact> MjData::contact() const {
  return InitWrapperArray<MjContact>(get()->contact, get()->ncon);
}

MjvScene::MjvScene() {
  owned_ = true;
  ptr_ = new mjvScene;
  mjv_defaultScene(ptr_);
  mjv_makeScene(nullptr, ptr_, 0);
  lights = InitWrapperArray<MjvLight>(ptr_->lights, mjMAXLIGHT);
  camera = InitWrapperArray<MjvGLCamera>(ptr_->camera, 2);
};

MjvScene::MjvScene(MjModel *m, int maxgeom) {
  owned_ = true;
  model = m->get();
  ptr_ = new mjvScene;
  mjv_defaultScene(ptr_);
  mjv_makeScene(model, ptr_, maxgeom);
  lights = InitWrapperArray<MjvLight>(ptr_->lights, mjMAXLIGHT);
  camera = InitWrapperArray<MjvGLCamera>(ptr_->camera, 2);
};
MjvScene::~MjvScene() {
  if (owned_ && ptr_) {
    mjv_freeScene(ptr_);
    delete ptr_;
  }
}

mjvScene* MjvScene::get() const { return ptr_; }
void MjvScene::set(mjvScene *ptr) { ptr_ = ptr; }

// Taken from the python mujoco bindings code for MjvScene Wrapper
int MjvScene::GetSumFlexFaces() const {
  int nflexface = 0;
  int flexfacenum = 0;
  for (int f = 0; f < model->nflex; f++) {
    if (model->flex_dim[f] == 0) {
      // 1D : 0
      flexfacenum = 0;
    } else if (model->flex_dim[f] == 2) {
      // 2D: 2*fragments + 2*elements
      flexfacenum = 2 * model->flex_shellnum[f] + 2 * model->flex_elemnum[f];
    } else {
      // 3D: max(fragments, 4*maxlayer)
      // find number of elements in biggest layer
      int maxlayer = 0, layer = 0, nlayer = 1;
      while (nlayer) {
        nlayer = 0;
        for (int e = 0; e < model->flex_elemnum[f]; e++) {
          if (model->flex_elemlayer[model->flex_elemadr[f] + e] == layer) {
            nlayer++;
          }
        }
        maxlayer = mjMAX(maxlayer, nlayer);
        layer++;
      }
      flexfacenum = mjMAX(model->flex_shellnum[f], 4 * maxlayer);
    }

    // accumulate over flexes
    nflexface += flexfacenum;
  }
  return nflexface;
}

std::vector<MjvGeom> MjvScene::geoms() const {
  return InitWrapperArray<MjvGeom>(ptr_->geoms, ptr_->ngeom);
}

MjSpec::MjSpec()
    : ptr_(mj_makeSpec()),
      option(&ptr_->option),
      visual(&ptr_->visual),
      stat(&ptr_->stat),
      compiler(&ptr_->compiler),
      element(ptr_->element) {
  owned_ = true;
  mjs_defaultSpec(ptr_);
};

MjSpec::MjSpec(mjSpec *ptr)
    : ptr_(ptr),
      option(&ptr_->option),
      visual(&ptr_->visual),
      stat(&ptr_->stat),
      compiler(&ptr_->compiler),
      element(ptr_->element) {}

MjSpec::MjSpec(const MjSpec &other)
    : ptr_(mj_copySpec(other.get())),
      option(&ptr_->option),
      visual(&ptr_->visual),
      stat(&ptr_->stat),
      compiler(&ptr_->compiler),
      element(ptr_->element) {
  owned_ = true;
}

MjSpec& MjSpec::operator=(const MjSpec &other) {
  if (this == &other) {
    return *this;
  }
  if (owned_ && ptr_) {
    mj_deleteSpec(ptr_);
  }
  ptr_ = mj_copySpec(other.get());
  owned_ = true;
  option.set(&ptr_->option);
  visual.set(&ptr_->visual);
  stat.set(&ptr_->stat);
  compiler.set(&ptr_->compiler);
  element.set(ptr_->element);
  return *this;
}

MjSpec::~MjSpec() {
  if (ptr_ && owned_) {
    mj_deleteSpec(ptr_);
  }
}

mjSpec *MjSpec::get() const { return ptr_; }
void MjSpec::set(mjSpec *ptr) { ptr_ = ptr; }

// ======= FACTORY AND HELPER FUNCTIONS ========= //
std::unique_ptr<MjModel> loadFromXML(std::string filename) {
  char error[1000];
  mjModel *model = mj_loadXML(filename.c_str(), nullptr, error, sizeof(error));
  if (!model) {
    printf("Loading error: %s\n", error);
    return nullptr;
  }
  return std::unique_ptr<MjModel>(new MjModel(model));
}

std::unique_ptr<MjSpec> parseXMLString(const std::string &xml) {
  char error[1000];
  mjSpec *ptr = mj_parseXMLString(xml.c_str(), nullptr, error, sizeof(error));
  if (!ptr) {
    printf("Could not create Spec from XML string: %s\n", error);
    return nullptr;
  }
  return std::unique_ptr<MjSpec>(new MjSpec(ptr));
}

EMSCRIPTEN_BINDINGS(mujoco_structs) {
  // {{ AUTOGENNED_STRUCTS_BINDINGS }}

  // TODO: should be generated in future CLs -- //
  emscripten::register_vector<MjSolverStat>("MjSolverStatVec");
  emscripten::register_vector<MjTimerStat>("MjTimerStatVec");
  emscripten::register_vector<MjWarningStat>("MjWarningStatVec");
  emscripten::register_vector<MjContact>("MjContactVec");
  emscripten::register_vector<MjvLight>("MjvLightVec");
  emscripten::register_vector<MjvGLCamera>("MjvGLCameraVec");
  emscripten::register_vector<MjvGeom>("MjvGeomVec");
}

// {{ WRAPPER_FUNCTIONS }}

void error_wrapper(const String& msg) { mju_error("%s\n", msg.as<const std::string>().data()); }

int mj_saveLastXML_wrapper(const String& filename, const MjModel& m) {
  CHECK_VAL(filename);
  std::array<char, 1024> error;
  int result = mj_saveLastXML(filename.as<const std::string>().data(), m.get(), error.data(), error.size());
  if (!result) {
    mju_error("%s", error.data());
  }
  return result;
}

int mj_setLengthRange_wrapper(const MjModel& m, const MjData& d, int index, const MjLROpt& opt) {
  std::array<char, 1024> error;
  int result = mj_setLengthRange(m.get(), d.get(), index, opt.get(), error.data(), error.size());
  if (!result) {
    mju_error("%s", error.data());
  }
  return result;
}

EMSCRIPTEN_BINDINGS(mujoco_functions) {
  function("parseXMLString", &parseXMLString, take_ownership());
  function("error", &error_wrapper);
  // {{ FUNCTION_BINDINGS }}
  class_<WasmBuffer<float>>("FloatBuffer")
      .constructor<int>()
      .class_function("FromArray", &WasmBuffer<float>::FromArray)
      .function("GetPointer", &WasmBuffer<float>::GetPointer)
      .function("GetElementCount", &WasmBuffer<float>::GetElementCount)
      .function("GetView", &WasmBuffer<float>::GetView);
  class_<WasmBuffer<double>>("DoubleBuffer")
      .constructor<int>()
      .class_function("FromArray", &WasmBuffer<double>::FromArray)
      .function("GetPointer", &WasmBuffer<double>::GetPointer)
      .function("GetElementCount", &WasmBuffer<double>::GetElementCount)
      .function("GetView", &WasmBuffer<double>::GetView);
  class_<WasmBuffer<int>>("IntBuffer")
      .constructor<int>()
      .class_function("FromArray", &WasmBuffer<int>::FromArray)
      .function("GetPointer", &WasmBuffer<int>::GetPointer)
      .function("GetElementCount", &WasmBuffer<int>::GetElementCount)
      .function("GetView", &WasmBuffer<int>::GetView);
  register_vector<std::string>("mjStringVec");
  register_vector<int>("mjIntVec");
  register_vector<mjIntVec>("mjIntVecVec");
  register_vector<float>("mjFloatVec");
  register_vector<mjFloatVec>("mjFloatVecVec");
  register_vector<double>("mjDoubleVec");
  // register_type gives better type information (val is mapped to any by default)
  register_type<NumberArray>("number[]");
  register_type<String>("string");
  register_vector<uint8_t>("mjByteVec");
  register_optional<MjsElement>();
  register_optional<MjsBody>();
  register_optional<MjsSite>();
  register_optional<MjsJoint>();
  register_optional<MjsGeom>();
  register_optional<MjsCamera>();
  register_optional<MjsLight>();
  register_optional<MjsFrame>();
  register_optional<MjsActuator>();
  register_optional<MjsSensor>();
  register_optional<MjsFlex>();
  register_optional<MjsPair>();
  register_optional<MjsExclude>();
  register_optional<MjsEquality>();
  register_optional<MjsTendon>();
  register_optional<MjsWrap>();
  register_optional<MjsNumeric>();
  register_optional<MjsText>();
  register_optional<MjsTuple>();
  register_optional<MjsKey>();
  register_optional<MjsPlugin>();
  register_optional<MjsDefault>();
  register_optional<MjsMesh>();
  register_optional<MjsHField>();
  register_optional<MjsSkin>();
  register_optional<MjsTexture>();
  register_optional<MjsMaterial>();
  register_optional<MjSpec>();
}

}  // namespace mujoco::wasm
// NOLINTEND(whitespace/semicolon)
// NOLINTEND(whitespace/line_length)