//
// Created by dizi on 2025/2/17.
//
#include "include/ssc_lib.h"

void* handle;

/**
 * 初始化
 */
extern "C"
JNIEXPORT jint JNICALL
Java_com_wmdigit_wmpos_ai_ssc_SnackSameCheckEngine_initSSC(JNIEnv *env,
                                                           jobject thiz,
                                                           jboolean use_gpu,
                                                           jstring j_cache_path) {
    // GPU缓存地址转char*
    const char* cache_path = env->GetStringUTFChars(j_cache_path, nullptr);
    LOGD("%s", cache_path);
    CLSRETRI_SSC_INIT_INFO init_info;
    init_info.cache_dir = cache_path;
    init_info.precision_level = 2;
    if (use_gpu){
        init_info.device_type = CLSRETRI_SSC_DEVICE::CLSRETRI_SSC_OPENGL;
    }
    else{
        init_info.device_type = CLSRETRI_SSC_DEVICE::CLSRETRI_SSC_CPU;
    }
    // 初始化
    int ret = CLSRETRI_SSC_Init(init_info,&handle);
    // 释放资源
    env->ReleaseStringUTFChars(j_cache_path, cache_path);
    LOGI("SSC init result:%x", ret);
    return ret;
}

/**
 * 获取当前运算平台
 */
extern "C"
JNIEXPORT jint JNICALL
Java_com_wmdigit_wmpos_ai_ssc_SnackSameCheckEngine_checkPlatform(JNIEnv *env, jobject thiz) {
    CLSRETRI_SSC_DEVICE device_info;
    int ret = CLSRETRI_SSC_GetInfo(&handle, "running_device", &device_info);
    if (ret != 0){
        LOGE("Failed to checkPlatform");
        return ret;
    }
    else{
        LOGI("当前使用平台为%d", device_info);
        return device_info;
    }
}

extern "C"
JNIEXPORT jint JNICALL
Java_com_wmdigit_wmpos_ai_ssc_SnackSameCheckEngine_process(JNIEnv *env, jobject thiz,
                                                           jobject bitmap) {
    int result = 0;
    // 保存图片像素数据
    void *bitmapPixels;
    // 保存图片参数
    AndroidBitmapInfo bitmapInfo;
    // 获取图片参数
    if(AndroidBitmap_getInfo(env, bitmap, &bitmapInfo) < 0){
        return result;
    }
    // 只支持 ARGB_8888 和 RGB_565
    if(bitmapInfo.format != ANDROID_BITMAP_FORMAT_RGBA_8888 && bitmapInfo.format != ANDROID_BITMAP_FORMAT_RGB_565){
        return result;
    }
    // 获取图片像素（锁定内存块）
    if(AndroidBitmap_lockPixels(env, bitmap, &bitmapPixels) < 0){
        return result;
    }
    int type;
    if (bitmapInfo.format == ANDROID_BITMAP_FORMAT_RGBA_8888) {
        type = CV_8UC4;
    } else {
        type = CV_8UC2;
    }
    cv::Mat mat(bitmapInfo.height, bitmapInfo.width, type, bitmapPixels);

    cv::Mat dst_mat;
    cv::cvtColor(mat, dst_mat, CV_RGBA2BGR);
    // 解锁
    AndroidBitmap_unlockPixels(env, bitmap);
    CLSRETRI_SSC_INPUT input;
    CLSRETRI_SSC_OUTPUT output;
    memset(&output, 0, sizeof(CLSRETRI_SSC_OUTPUT));
    input.img = dst_mat;
    // 推理
    int hr = CLSRETRI_SSC_Process(input, &output, handle);
    if (hr != 0)
    {
        LOGE("CLSRETRI_SSC_Process error code:%x", hr);
        return 0;
    }
    float maxProb = 0;
    // 遍历推理结果
    for (int i = 0; i < CLSRETRI_SSC_CLS_TOP_N; ++i) {
        if (output.cls[i].prob > maxProb){
            maxProb = output.cls[i].prob;
            result  = output.cls[i].label;
        }
    }
    LOGI("result = %d, maxProb = %f", result, maxProb);
//    if (result == 1 && maxProb < 0.99){
//        result = 0;
//    }
    return result;
}

/**
 * 资源释放
 */
extern "C"
JNIEXPORT void JNICALL
Java_com_wmdigit_wmpos_ai_ssc_SnackSameCheckEngine_release(JNIEnv *env, jobject thiz) {
    CLSRETRI_SSC_Release(&handle);
}