#include <napi.h>
#include "annoylib.h"
#include "kissrandom.h"
#include <string>
#include <vector>

using namespace Annoy;

template <typename Distance>
class AnnoyIndexWrapper : public Napi::ObjectWrap<AnnoyIndexWrapper<Distance>> {
public:
  using IndexType = AnnoyIndex<int32_t, float, Distance, Kiss64Random, AnnoyIndexMultiThreadedBuildPolicy>;

  static Napi::FunctionReference constructor;

  static Napi::Function Init(Napi::Env env, const char* name) {
    Napi::Function func = AnnoyIndexWrapper::DefineClass(env, name, {
      AnnoyIndexWrapper::InstanceMethod("addItem", &AnnoyIndexWrapper::AddItem),
      AnnoyIndexWrapper::InstanceMethod("build", &AnnoyIndexWrapper::Build),
      AnnoyIndexWrapper::InstanceMethod("unbuild", &AnnoyIndexWrapper::Unbuild),
      AnnoyIndexWrapper::InstanceMethod("save", &AnnoyIndexWrapper::Save),
      AnnoyIndexWrapper::InstanceMethod("load", &AnnoyIndexWrapper::Load),
      AnnoyIndexWrapper::InstanceMethod("unload", &AnnoyIndexWrapper::Unload),
      AnnoyIndexWrapper::InstanceMethod("getDistance", &AnnoyIndexWrapper::GetDistance),
      AnnoyIndexWrapper::InstanceMethod("getNnsByItem", &AnnoyIndexWrapper::GetNnsByItem),
      AnnoyIndexWrapper::InstanceMethod("getNnsByVector", &AnnoyIndexWrapper::GetNnsByVector),
      AnnoyIndexWrapper::InstanceMethod("getNItems", &AnnoyIndexWrapper::GetNItems),
      AnnoyIndexWrapper::InstanceMethod("getNTrees", &AnnoyIndexWrapper::GetNTrees),
      AnnoyIndexWrapper::InstanceMethod("getItem", &AnnoyIndexWrapper::GetItem),
      AnnoyIndexWrapper::InstanceMethod("setSeed", &AnnoyIndexWrapper::SetSeed),
      AnnoyIndexWrapper::InstanceMethod("onDiskBuild", &AnnoyIndexWrapper::OnDiskBuild),
      AnnoyIndexWrapper::InstanceMethod("verbose", &AnnoyIndexWrapper::Verbose),
      AnnoyIndexWrapper::InstanceMethod("getF", &AnnoyIndexWrapper::GetF),
    });
    constructor = Napi::Persistent(func);
    constructor.SuppressDestruct();
    return func;
  }

  AnnoyIndexWrapper(const Napi::CallbackInfo& info)
      : Napi::ObjectWrap<AnnoyIndexWrapper<Distance>>(info) {
    Napi::Env env = info.Env();
    if (info.Length() < 1 || !info[0].IsNumber()) {
      Napi::TypeError::New(env, "Expected number of dimensions as first argument").ThrowAsJavaScriptException();
      return;
    }
    int f = info[0].As<Napi::Number>().Int32Value();
    index_ = new IndexType(f);
  }

  ~AnnoyIndexWrapper() {
    delete index_;
  }

private:
  IndexType* index_;

  Napi::Value AddItem(const Napi::CallbackInfo& info) {
    Napi::Env env = info.Env();
    if (info.Length() < 2 || !info[0].IsNumber()) {
      Napi::TypeError::New(env, "Expected (index: number, vector: number[])").ThrowAsJavaScriptException();
      return env.Undefined();
    }

    int32_t item = info[0].As<Napi::Number>().Int32Value();
    std::vector<float> vec = jsArrayToVector(env, info[1]);

    char* error = nullptr;
    if (!index_->add_item(item, vec.data(), &error)) {
      std::string msg = error ? error : "Failed to add item";
      free(error);
      Napi::Error::New(env, msg).ThrowAsJavaScriptException();
      return env.Undefined();
    }
    return env.Undefined();
  }

  Napi::Value Build(const Napi::CallbackInfo& info) {
    Napi::Env env = info.Env();
    if (info.Length() < 1 || !info[0].IsNumber()) {
      Napi::TypeError::New(env, "Expected number of trees").ThrowAsJavaScriptException();
      return env.Undefined();
    }

    int q = info[0].As<Napi::Number>().Int32Value();
    int n_threads = -1;
    if (info.Length() > 1 && info[1].IsNumber()) {
      n_threads = info[1].As<Napi::Number>().Int32Value();
    }

    char* error = nullptr;
    if (!index_->build(q, n_threads, &error)) {
      std::string msg = error ? error : "Failed to build";
      free(error);
      Napi::Error::New(env, msg).ThrowAsJavaScriptException();
      return env.Undefined();
    }
    return env.Undefined();
  }

  Napi::Value Unbuild(const Napi::CallbackInfo& info) {
    Napi::Env env = info.Env();
    char* error = nullptr;
    if (!index_->unbuild(&error)) {
      std::string msg = error ? error : "Failed to unbuild";
      free(error);
      Napi::Error::New(env, msg).ThrowAsJavaScriptException();
    }
    return env.Undefined();
  }

  Napi::Value Save(const Napi::CallbackInfo& info) {
    Napi::Env env = info.Env();
    if (info.Length() < 1 || !info[0].IsString()) {
      Napi::TypeError::New(env, "Expected filename").ThrowAsJavaScriptException();
      return env.Undefined();
    }

    std::string filename = info[0].As<Napi::String>().Utf8Value();
    bool prefault = false;
    if (info.Length() > 1 && info[1].IsBoolean()) {
      prefault = info[1].As<Napi::Boolean>().Value();
    }

    char* error = nullptr;
    if (!index_->save(filename.c_str(), prefault, &error)) {
      std::string msg = error ? error : "Failed to save";
      free(error);
      Napi::Error::New(env, msg).ThrowAsJavaScriptException();
    }
    return env.Undefined();
  }

  Napi::Value Load(const Napi::CallbackInfo& info) {
    Napi::Env env = info.Env();
    if (info.Length() < 1 || !info[0].IsString()) {
      Napi::TypeError::New(env, "Expected filename").ThrowAsJavaScriptException();
      return env.Undefined();
    }

    std::string filename = info[0].As<Napi::String>().Utf8Value();
    bool prefault = false;
    if (info.Length() > 1 && info[1].IsBoolean()) {
      prefault = info[1].As<Napi::Boolean>().Value();
    }

    char* error = nullptr;
    if (!index_->load(filename.c_str(), prefault, &error)) {
      std::string msg = error ? error : "Failed to load";
      free(error);
      Napi::Error::New(env, msg).ThrowAsJavaScriptException();
    }
    return env.Undefined();
  }

  Napi::Value Unload(const Napi::CallbackInfo& info) {
    index_->unload();
    return info.Env().Undefined();
  }

  Napi::Value GetDistance(const Napi::CallbackInfo& info) {
    Napi::Env env = info.Env();
    if (info.Length() < 2 || !info[0].IsNumber() || !info[1].IsNumber()) {
      Napi::TypeError::New(env, "Expected (i: number, j: number)").ThrowAsJavaScriptException();
      return env.Undefined();
    }

    int32_t i = info[0].As<Napi::Number>().Int32Value();
    int32_t j = info[1].As<Napi::Number>().Int32Value();
    return Napi::Number::New(env, index_->get_distance(i, j));
  }

  Napi::Value GetNnsByItem(const Napi::CallbackInfo& info) {
    Napi::Env env = info.Env();
    if (info.Length() < 2 || !info[0].IsNumber() || !info[1].IsNumber()) {
      Napi::TypeError::New(env, "Expected (item: number, n: number, searchK?: number, includeDistances?: boolean)")
          .ThrowAsJavaScriptException();
      return env.Undefined();
    }

    int32_t item = info[0].As<Napi::Number>().Int32Value();
    size_t n = info[1].As<Napi::Number>().Uint32Value();
    int search_k = -1;
    bool include_distances = false;

    if (info.Length() > 2 && info[2].IsNumber()) {
      search_k = info[2].As<Napi::Number>().Int32Value();
    }
    if (info.Length() > 3 && info[3].IsBoolean()) {
      include_distances = info[3].As<Napi::Boolean>().Value();
    }

    std::vector<int32_t> result;
    std::vector<float> distances;
    index_->get_nns_by_item(item, n, search_k,
                            &result,
                            include_distances ? &distances : nullptr);

    if (include_distances) {
      Napi::Object obj = Napi::Object::New(env);
      obj.Set("neighbors", vectorToJsArray(env, result));
      obj.Set("distances", floatVectorToJsArray(env, distances));
      return obj;
    }
    return vectorToJsArray(env, result);
  }

  Napi::Value GetNnsByVector(const Napi::CallbackInfo& info) {
    Napi::Env env = info.Env();
    if (info.Length() < 2 || !info[1].IsNumber()) {
      Napi::TypeError::New(env, "Expected (vector: number[], n: number, searchK?: number, includeDistances?: boolean)")
          .ThrowAsJavaScriptException();
      return env.Undefined();
    }

    std::vector<float> vec = jsArrayToVector(env, info[0]);
    size_t n = info[1].As<Napi::Number>().Uint32Value();
    int search_k = -1;
    bool include_distances = false;

    if (info.Length() > 2 && info[2].IsNumber()) {
      search_k = info[2].As<Napi::Number>().Int32Value();
    }
    if (info.Length() > 3 && info[3].IsBoolean()) {
      include_distances = info[3].As<Napi::Boolean>().Value();
    }

    std::vector<int32_t> result;
    std::vector<float> distances;
    index_->get_nns_by_vector(vec.data(), n, search_k,
                              &result,
                              include_distances ? &distances : nullptr);

    if (include_distances) {
      Napi::Object obj = Napi::Object::New(env);
      obj.Set("neighbors", vectorToJsArray(env, result));
      obj.Set("distances", floatVectorToJsArray(env, distances));
      return obj;
    }
    return vectorToJsArray(env, result);
  }

  Napi::Value GetNItems(const Napi::CallbackInfo& info) {
    return Napi::Number::New(info.Env(), index_->get_n_items());
  }

  Napi::Value GetNTrees(const Napi::CallbackInfo& info) {
    return Napi::Number::New(info.Env(), index_->get_n_trees());
  }

  Napi::Value GetItem(const Napi::CallbackInfo& info) {
    Napi::Env env = info.Env();
    if (info.Length() < 1 || !info[0].IsNumber()) {
      Napi::TypeError::New(env, "Expected item index").ThrowAsJavaScriptException();
      return env.Undefined();
    }

    int32_t item = info[0].As<Napi::Number>().Int32Value();
    int f = index_->get_f();
    std::vector<float> vec(f);
    index_->get_item(item, vec.data());
    return floatVectorToJsArray(env, vec);
  }

  Napi::Value SetSeed(const Napi::CallbackInfo& info) {
    Napi::Env env = info.Env();
    if (info.Length() < 1 || !info[0].IsNumber()) {
      Napi::TypeError::New(env, "Expected seed number").ThrowAsJavaScriptException();
      return env.Undefined();
    }
    index_->set_seed(info[0].As<Napi::Number>().Int64Value());
    return env.Undefined();
  }

  Napi::Value OnDiskBuild(const Napi::CallbackInfo& info) {
    Napi::Env env = info.Env();
    if (info.Length() < 1 || !info[0].IsString()) {
      Napi::TypeError::New(env, "Expected filename").ThrowAsJavaScriptException();
      return env.Undefined();
    }

    std::string filename = info[0].As<Napi::String>().Utf8Value();
    char* error = nullptr;
    if (!index_->on_disk_build(filename.c_str(), &error)) {
      std::string msg = error ? error : "Failed to set on-disk build";
      free(error);
      Napi::Error::New(env, msg).ThrowAsJavaScriptException();
    }
    return env.Undefined();
  }

  Napi::Value Verbose(const Napi::CallbackInfo& info) {
    Napi::Env env = info.Env();
    if (info.Length() < 1 || !info[0].IsBoolean()) {
      Napi::TypeError::New(env, "Expected boolean").ThrowAsJavaScriptException();
      return env.Undefined();
    }
    index_->verbose(info[0].As<Napi::Boolean>().Value());
    return env.Undefined();
  }

  Napi::Value GetF(const Napi::CallbackInfo& info) {
    return Napi::Number::New(info.Env(), index_->get_f());
  }

  // Helpers
  static std::vector<float> jsArrayToVector(Napi::Env env, Napi::Value val) {
    if (!val.IsArray()) {
      Napi::TypeError::New(env, "Expected an array of numbers").ThrowAsJavaScriptException();
      return {};
    }
    Napi::Array arr = val.As<Napi::Array>();
    uint32_t len = arr.Length();
    std::vector<float> vec(len);
    for (uint32_t i = 0; i < len; i++) {
      vec[i] = arr.Get(i).As<Napi::Number>().FloatValue();
    }
    return vec;
  }

  static Napi::Array vectorToJsArray(Napi::Env env, const std::vector<int32_t>& vec) {
    Napi::Array arr = Napi::Array::New(env, vec.size());
    for (size_t i = 0; i < vec.size(); i++) {
      arr.Set(i, Napi::Number::New(env, vec[i]));
    }
    return arr;
  }

  static Napi::Array floatVectorToJsArray(Napi::Env env, const std::vector<float>& vec) {
    Napi::Array arr = Napi::Array::New(env, vec.size());
    for (size_t i = 0; i < vec.size(); i++) {
      arr.Set(i, Napi::Number::New(env, vec[i]));
    }
    return arr;
  }
};

template <typename Distance>
Napi::FunctionReference AnnoyIndexWrapper<Distance>::constructor;

Napi::Object Init(Napi::Env env, Napi::Object exports) {
  exports.Set("AnnoyIndexAngular",
              AnnoyIndexWrapper<Angular>::Init(env, "AnnoyIndexAngular"));
  exports.Set("AnnoyIndexEuclidean",
              AnnoyIndexWrapper<Euclidean>::Init(env, "AnnoyIndexEuclidean"));
  exports.Set("AnnoyIndexManhattan",
              AnnoyIndexWrapper<Manhattan>::Init(env, "AnnoyIndexManhattan"));
  exports.Set("AnnoyIndexDotProduct",
              AnnoyIndexWrapper<DotProduct>::Init(env, "AnnoyIndexDotProduct"));
  return exports;
}

NODE_API_MODULE(annoy, Init)