#include <jni.h>
#include <android/bitmap.h>
#include <android/log.h>

#include "RetinaFaceEngineRKNN.h"

#define LOG_TAG "RetinaFaceJNI"
#define LOGE(...) __android_log_print(ANDROID_LOG_ERROR, LOG_TAG, __VA_ARGS__)

extern "C" {

JNIEXPORT jlong JNICALL
Java_com_digitalperson_engine_RetinaFaceEngineRKNN_createEngineNative(JNIEnv* env, jobject thiz) {
    auto* engine = new RetinaFaceEngineRKNN();
    if (engine == nullptr) {
        return 0;
    }
    return reinterpret_cast<jlong>(engine);
}

JNIEXPORT jint JNICALL
Java_com_digitalperson_engine_RetinaFaceEngineRKNN_initNative(
        JNIEnv* env,
        jobject thiz,
        jlong ptr,
        jstring modelPath,
        jint inputSize,
        jfloat scoreThreshold,
        jfloat nmsThreshold) {
    auto* engine = reinterpret_cast<RetinaFaceEngineRKNN*>(ptr);
    if (engine == nullptr || modelPath == nullptr) {
        return -1;
    }
    const char* model = env->GetStringUTFChars(modelPath, nullptr);
    if (model == nullptr) {
        return -1;
    }
    int ret = engine->init(model, static_cast<int>(inputSize), scoreThreshold, nmsThreshold);
    env->ReleaseStringUTFChars(modelPath, model);
    return ret;
}

JNIEXPORT jfloatArray JNICALL
Java_com_digitalperson_engine_RetinaFaceEngineRKNN_detectNative(
        JNIEnv* env,
        jobject thiz,
        jlong ptr,
        jobject bitmapObj) {
    auto* engine = reinterpret_cast<RetinaFaceEngineRKNN*>(ptr);
    if (engine == nullptr || bitmapObj == nullptr) {
        return env->NewFloatArray(0);
    }

    AndroidBitmapInfo info{};
    if (AndroidBitmap_getInfo(env, bitmapObj, &info) < 0) {
        LOGE("AndroidBitmap_getInfo failed");
        return env->NewFloatArray(0);
    }
    if (info.format != ANDROID_BITMAP_FORMAT_RGBA_8888) {
        LOGE("Unsupported bitmap format: %d", info.format);
        return env->NewFloatArray(0);
    }

    void* pixels = nullptr;
    if (AndroidBitmap_lockPixels(env, bitmapObj, &pixels) < 0 || pixels == nullptr) {
        LOGE("AndroidBitmap_lockPixels failed");
        return env->NewFloatArray(0);
    }

    std::vector<float> result = engine->detect(
            reinterpret_cast<uint32_t*>(pixels),
            static_cast<int>(info.width),
            static_cast<int>(info.height),
            static_cast<int>(info.stride));

    AndroidBitmap_unlockPixels(env, bitmapObj);

    jfloatArray out = env->NewFloatArray(static_cast<jsize>(result.size()));
    if (out == nullptr) {
        return env->NewFloatArray(0);
    }
    if (!result.empty()) {
        env->SetFloatArrayRegion(out, 0, static_cast<jsize>(result.size()), result.data());
    }
    return out;
}

JNIEXPORT void JNICALL
Java_com_digitalperson_engine_RetinaFaceEngineRKNN_releaseNative(
        JNIEnv* env,
        jobject thiz,
        jlong ptr) {
    auto* engine = reinterpret_cast<RetinaFaceEngineRKNN*>(ptr);
    if (engine != nullptr) {
        engine->release();
        delete engine;
    }
}

}  // extern "C"