#include "jni.hpp"

#include <algorithm>
#include <optional>

#ifdef ANDROID
#include <android/log.h>
#else
#include <iostream>
#endif

namespace {

JavaVM *g_vm;

const char *_jni_error(jint err) {
  switch (err) {
    case JNI_OK:
      return "OK";
    case JNI_ERR:
      return "Unknown error";
    case JNI_EDETACHED:
      return "Thread detached from the VM";
    case JNI_EVERSION:
      return "JNI version error";
    case JNI_ENOMEM:
      return "Not enough memory";
    case JNI_EEXIST:
      return "VM already created";
    case JNI_EINVAL:
      return "Invalid arguments";
    default:
      return "Unexpected error";
  }
}

const char* u8_read(const char* str, uint32_t& out) {
  // Assume valid UTF-8 for speed and so it can be used to read modified utf-8 as well
  auto* const d = reinterpret_cast<const uint8_t*>(str);
  switch (d[0] >> 4) {
    case 0xf: // 4 byte
      out = (static_cast<uint32_t>(d[0] & 0x7) << 18) |
            (static_cast<uint32_t>(d[1] & 0x3f) << 12) |
            (static_cast<uint32_t>(d[2] & 0x3f) << 6) |
            static_cast<uint8_t>(d[3] & 0x3f);
      return str + 4;
    case 0xe: // 3 byte
      out = (static_cast<uint32_t>(d[0] & 0xf) << 12) |
            (static_cast<uint32_t>(d[1] & 0x3f) << 6) |
            static_cast<uint8_t>(d[2] & 0x3f);
      return str + 3;
    case 0xd:
    case 0xc: // 2 byte
      out = (static_cast<uint32_t>(d[0] & 0x1f) << 6) |
          static_cast<uint8_t>(d[1] & 0x3f);
      return str + 2;
    default: // 1 byte
      out = static_cast<uint8_t>(d[0]);
      return str + 1;
  }
}

void u8_write(std::string& ret, uint32_t c) {
  if (c < 0x80) {
    ret.push_back(static_cast<char>(c));
  } else if (c < 0x800) {
    ret.push_back(static_cast<char>(0xc0 | (c >> 6)));
    ret.push_back(static_cast<char>(0x80 | (c & 0x3f)));
  } else if (c < 0x10000) {
    ret.push_back(static_cast<char>(0xe0 | (c >> 12)));
    ret.push_back(static_cast<char>(0x80 | ((c >> 6) & 0x3f)));
    ret.push_back(static_cast<char>(0x80 | (c & 0x3f)));
  } else {
    ret.push_back(static_cast<char>(0xf0 | (c >> 18)));
    ret.push_back(static_cast<char>(0x80 | ((c >> 12) & 0x3f)));
    ret.push_back(static_cast<char>(0x80 | ((c >> 6) & 0x3f)));
    ret.push_back(static_cast<char>(0x80 | (c & 0x3f)));
  }
}

std::string MakeModifiedUTF8(const char* str, const char* prev, const char* ptr, uint32_t c) {
  std::string ret;
  while (true) {
    ret.append(str, prev - str);
    str = ptr;
    u8_write(ret, 0xd800 + ((c & 0xffff) >> 10));
    u8_write(ret, 0xdc00 + (c & 0x3ff));
    if (!*ptr) break;
    do {
      prev = ptr;
      ptr = u8_read(prev, c);
      if (c > 0xffff)
        break;
    } while (*ptr);
  }
  return ret;
}

std::optional<std::string> MakeModifiedUTF8IfNeeded(const char* str) {
  auto* ptr = str;
  while (*ptr) {
    uint32_t c;
    auto* next = u8_read(ptr, c);
    if (c > 0xffff) {
      return MakeModifiedUTF8(str, ptr, next, c);
    }
    ptr = next;
  }
  return std::nullopt;
}

const char* splice(std::string& str, const char* start, const char* end,
                   const char* insert, size_t size) {
  auto pos = start - str.c_str();
  str.replace(pos, end - start, insert, size);
  return str.c_str() + pos + size;
}

void UnmodifyUTF8(std::string& str) {
  auto* ptr = str.c_str();
  while (*ptr) {
    uint32_t u;
    auto* next = u8_read(ptr, u);
    if (u == 0) {
      next = splice(str, ptr, next, "\0", 1);
    } else if (u >= 0xd800 && u <= 0xdfff) {
      uint32_t v;
      next = u8_read(next, v);
      u = 0x10000 | ((u - 0xd800) << 10) | (v - 0xdc00);
      std::string tmp;
      u8_write(tmp, u);
      next = splice(str, ptr, next, tmp.data(), tmp.size());
    }
    ptr = next;
  }
}

}  // namespace

namespace jni {

namespace internal {

void _abort_if_not_ok(const char *file, int line, jint ret) {
  if (ret == JNI_OK) [[likely]] return;
#ifdef ANDROID
  __android_log_assert(nullptr, "jni", "JNI error: %s", _jni_error(ret));
#else
  std::cerr << "JNI error: " << _jni_error(ret) << std::endl;
  abort();
#endif
}

void _abort_with_exception(const char* file, int line, JNIEnv* env) {
  auto throwable = jni::ExceptionOccurred(env);
  env->ExceptionClear();
  if (throwable) {
    auto throwable_class = jni::FindClass(env, "java/lang/Throwable");
    if (throwable_class) {
      auto throwable_toString = env->GetMethodID(throwable_class.get(),
                                                 "toString",
                                                 "()Ljava/lang/String;");
      if (throwable_toString) {
        auto description = jni::CallObjectMethod<jstring>(env, throwable,
                                                          throwable_toString);
        auto str = jni::StringToUTF8(env, description);
#ifdef ANDROID
        __android_log_assert(nullptr, "jni", "JNI error: %s", str.c_str());
#else
        std::cerr << "JNI error: " << str << std::endl;
        abort();
#endif
      }
    }
    env->ExceptionClear();
#ifdef ANDROID
    __android_log_assert(nullptr, "jni",
                         "Unexpected NULL but no exception");
#else
    std::cerr << "Unexpected NULL but no exception" << std::endl;
    abort();
#endif
  }
}

JNIEnv* NonFatalAttachCurrentThread() {
#ifdef ANDROID
  JNIEnv* env;
  auto ret = g_vm->AttachCurrentThread(&env, nullptr);
  if (ret == JNI_OK)
    return env;
  return nullptr;
#else
  void* v_env;
  auto ret = g_vm->AttachCurrentThread(&v_env, nullptr);
  if (ret == JNI_OK)
    return reinterpret_cast<JNIEnv*>(v_env);
  return nullptr;
#endif
}

}  // namespace internal

JNIEnv* AttachCurrentThread() {
#ifdef ANDROID
  JNIEnv* env;
  auto ret = g_vm->AttachCurrentThread(&env, nullptr);
  ABORT_IF_NOT_OK(ret);
  return env;
#else
  void* v_env;
  auto ret = g_vm->AttachCurrentThread(&v_env, nullptr);
  ABORT_IF_NOT_OK(ret);
  return reinterpret_cast<JNIEnv*>(v_env);
#endif
}

JNIEnv* OnLoad(JavaVM* vm) {
  void* v_env;
  auto ret = vm->GetEnv(&v_env, JNI_VERSION);
  ABORT_IF_NOT_OK(ret);
  g_vm = vm;
  return reinterpret_cast<JNIEnv*>(v_env);
}

LocalRef<jclass> FindClass(JNIEnv *env, const char *name) {
  return {env, env->FindClass(name)};
}

LocalRef<jthrowable> ExceptionOccurred(JNIEnv* env) {
  return {env, env->ExceptionOccurred()};
}

std::string StringToUTF8(JNIEnv* env, const Ref<jstring>& str) {
  if (!str) return "null";
  auto size = env->GetStringUTFLength(str.get());
  auto len = env->GetStringLength(str.get());
  std::string ret(size, ' ');
  env->GetStringUTFRegion(str.get(), 0, len, ret.data());
  UnmodifyUTF8(ret);
  return ret;
}

LocalRef<jstring> UTF8ToString(JNIEnv* env, const std::string& str) {
  auto ret = MakeModifiedUTF8IfNeeded(str.c_str());
  if (ret.has_value()) {
    return {env, env->NewStringUTF(ret->c_str())};
  }
  return {env, env->NewStringUTF(str.c_str())};
}

LocalRef<jstring> UTF8ToString(JNIEnv* env, const char* str) {
  if (str == nullptr) return nullptr;
  auto ret = MakeModifiedUTF8IfNeeded(str);
  if (ret.has_value()) {
    return {env, env->NewStringUTF(ret->c_str())};
  }
  return {env, env->NewStringUTF(str)};
}

LocalRef<jbyteArray> VectorToByteArray(JNIEnv* env, const std::vector<uint8_t>& data) {
  auto len = static_cast<jsize>(data.size());
  auto ret = LocalRef<jbyteArray>(env, env->NewByteArray(len));
  ABORT_IF_NULL(env, ret);
  auto* ptr = reinterpret_cast<jbyte*>(env->GetPrimitiveArrayCritical(ret.get(), nullptr));
  std::copy_n(data.data(), data.size(), ptr);
  env->ReleasePrimitiveArrayCritical(ret.get(), ptr, JNI_COMMIT);
  return ret;
}

std::vector<uint8_t> ByteArrayToVector(JNIEnv* env, const Ref<jbyteArray>& data) {
  if (!data) return {};
  auto len = env->GetArrayLength(data.get());
  std::vector<uint8_t> ret(len);
  static_assert(sizeof(jbyte) == sizeof(uint8_t));
  env->GetByteArrayRegion(data.get(), 0, len, reinterpret_cast<jbyte*>(ret.data()));
  return ret;
}

LocalRef<jobjectArray> CreateArray(JNIEnv* env, const Ref<jclass>& element_class, std::vector<LocalRef<jobject>> objects) {
  auto ret = LocalRef<jobjectArray>(env, env->NewObjectArray(static_cast<jsize>(objects.size()), element_class.get(), nullptr));
  ABORT_IF_NULL(env, ret);
  jsize i = 0;
  for (auto&& obj : objects) {
    env->SetObjectArrayElement(ret.get(), i++, obj.release());
  }
  return ret;
}

}  // namespace jni