// Tencent is pleased to support the open source community by making ncnn available.
//
// Copyright (C) 2025 THL A29 Limited, a Tencent company. All rights reserved.
//
// Licensed under the BSD 3-Clause License (the "License"); you may not use this file except
// in compliance with the License. You may obtain a copy of the License at
//
// https://opensource.org/licenses/BSD-3-Clause
//
// 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.

#include <android/asset_manager_jni.h>
#include <android/bitmap.h>
#include <android/native_window_jni.h>
#include <android/native_window.h>

#include <android/log.h>

#include <jni.h>

#include <string>
#include <vector>
#include <cstring>

#include <platform.h>
#include <benchmark.h>

#include "rvm.h"

#include "ndkcamera.h"

#include <opencv2/core/core.hpp>
#include <opencv2/imgproc/imgproc.hpp>

#if __ARM_NEON
#include <arm_neon.h>
#endif // __ARM_NEON

static int draw_unsupported(cv::Mat& rgb)
{
    const char text[] = "unsupported";

    int baseLine = 0;
    cv::Size label_size = cv::getTextSize(text, cv::FONT_HERSHEY_SIMPLEX, 1.0, 1, &baseLine);

    int y = (rgb.rows - label_size.height) / 2;
    int x = (rgb.cols - label_size.width) / 2;

    cv::rectangle(rgb, cv::Rect(cv::Point(x, y), cv::Size(label_size.width, label_size.height + baseLine)),
                    cv::Scalar(255, 255, 255), -1);

    cv::putText(rgb, text, cv::Point(x, y + label_size.height),
                cv::FONT_HERSHEY_SIMPLEX, 1.0, cv::Scalar(0, 0, 0));

    return 0;
}

static int draw_fps(cv::Mat& rgb)
{
    // resolve moving average
    float avg_fps = 0.f;
    {
        static double t0 = 0.f;
        static float fps_history[10] = {0.f};

        double t1 = ncnn::get_current_time();
        if (t0 == 0.f)
        {
            t0 = t1;
            return 0;
        }

        float fps = 1000.f / (t1 - t0);
        t0 = t1;

        for (int i = 9; i >= 1; i--)
        {
            fps_history[i] = fps_history[i - 1];
        }
        fps_history[0] = fps;

        if (fps_history[9] == 0.f)
        {
            return 0;
        }

        for (int i = 0; i < 10; i++)
        {
            avg_fps += fps_history[i];
        }
        avg_fps /= 10.f;
    }

    char text[32];
    sprintf(text, "FPS=%.2f", avg_fps);

    int baseLine = 0;
    cv::Size label_size = cv::getTextSize(text, cv::FONT_HERSHEY_SIMPLEX, 0.5, 1, &baseLine);

    int y = 0;
    int x = rgb.cols - label_size.width;

    cv::rectangle(rgb, cv::Rect(cv::Point(x, y), cv::Size(label_size.width, label_size.height + baseLine)),
                    cv::Scalar(255, 255, 255), -1);

    cv::putText(rgb, text, cv::Point(x, y + label_size.height),
                cv::FONT_HERSHEY_SIMPLEX, 0.5, cv::Scalar(0, 0, 0));

    return 0;
}

static RVM* g_rvm = 0;
static ncnn::Mutex lock;
static InterFeatures g_feats;

class MyNdkCamera : public NdkCameraWindow
{
public:
    virtual void on_image_render(cv::Mat& rgb) const;
};

void MyNdkCamera::on_image_render(cv::Mat& rgb) const
{
    // rvm
    {
        ncnn::MutexLockGuard g(lock);

        if (g_rvm)
        {
            cv::Mat fgr;
            cv::Mat pha;
            cv::Mat seg;
            g_rvm->detect(rgb, g_feats, fgr, pha, seg);

            g_rvm->draw(rgb, fgr, pha, seg);
        }
        else
        {
            draw_unsupported(rgb);
        }
    }

    // overlay fps
    draw_fps(rgb);

    // enforce target output resolution 640x640 and 180-degree rotation
    {
        // resize to 640x640 if needed
        if (rgb.cols != 640 || rgb.rows != 640)
        {
            cv::Mat resized;
            cv::resize(rgb, resized, cv::Size(640, 640), 0, 0, cv::INTER_LINEAR);
            resized.copyTo(rgb);
        }

        // rotate 180 degrees
        cv::Mat rotated;
        cv::rotate(rgb, rotated, cv::ROTATE_180);
        rotated.copyTo(rgb);
    }
}

static MyNdkCamera* g_camera = 0;

extern "C" {

JNIEXPORT jint JNI_OnLoad(JavaVM* vm, void* reserved)
{
    __android_log_print(ANDROID_LOG_DEBUG, "ncnn", "JNI_OnLoad");

    g_camera = new MyNdkCamera;

    ncnn::create_gpu_instance();

    return JNI_VERSION_1_4;
}

JNIEXPORT void JNI_OnUnload(JavaVM* vm, void* reserved)
{
    __android_log_print(ANDROID_LOG_DEBUG, "ncnn", "JNI_OnUnload");

    {
        ncnn::MutexLockGuard g(lock);

        delete g_rvm;
        g_rvm = 0;
    }

    ncnn::destroy_gpu_instance();

    delete g_camera;
    g_camera = 0;
}

// public native boolean loadModel(AssetManager mgr, int modelid, int sizeid, int intrainterid, int postprocid, int cpugpu);
JNIEXPORT jboolean JNICALL Java_io_livekit_android_track_processing_video_RVMNcnn_loadModel(JNIEnv* env, jobject thiz, jobject assetManager, jint modelid, jint sizeid, jint intrainterid, jint postprocid, jint cpugpu)
{
    if (modelid < 0 || modelid > 1 || sizeid < 0 || sizeid > 6 || intrainterid < 0 || intrainterid > 1 || postprocid < 0 || postprocid > 2 || cpugpu < 0 || cpugpu > 2)
    {
        return JNI_FALSE;
    }

    AAssetManager* mgr = AAssetManager_fromJava(env, assetManager);

    __android_log_print(ANDROID_LOG_DEBUG, "ncnn", "loadModel %p", mgr);

    const char* modeltypes[2] =
    {
        "mobilenetv3",
        "resnet50"
    };

    const int sizetypes[7] =
    {
        256,
        320,
        384,
        448,
        512,
        576,
        640
    };

    std::string parampath = std::string("rvm_") + modeltypes[(int)modelid] + ".ncnn.param";
    std::string modelpath = std::string("rvm_") + modeltypes[(int)modelid] + ".ncnn.bin";
    bool use_gpu = (int)cpugpu == 1;
    bool use_turnip = (int)cpugpu == 2;

    // reload
    {
        ncnn::MutexLockGuard g(lock);

        {
            // reset inter feats
            g_feats.r1.release();
            g_feats.r2.release();
            g_feats.r3.release();
            g_feats.r4.release();

            static int old_modelid = 0;
            static int old_cpugpu = 0;
            if (modelid != old_modelid || cpugpu != old_cpugpu)
            {
                // model or cpugpu changed
                delete g_rvm;
                g_rvm = 0;
            }
            old_modelid = modelid;
            old_cpugpu = cpugpu;

            ncnn::destroy_gpu_instance();

            if (use_turnip)
            {
                ncnn::create_gpu_instance("libvulkan_freedreno.so");
            }
            else if (use_gpu)
            {
                ncnn::create_gpu_instance();
            }

            if (!g_rvm)
            {
                g_rvm = new RVM;

                g_rvm->load(mgr, parampath.c_str(), modelpath.c_str(), use_gpu || use_turnip);
            }
            g_rvm->set_model_type((int)modelid);
            g_rvm->set_target_size(sizetypes[(int)sizeid]);
            g_rvm->set_intra_inter((int)intrainterid);

            // deep
            // fast
            // seg
            if (postprocid == 0)
                g_rvm->set_postproc_mode(false, true, false);
            if (postprocid == 1)
                g_rvm->set_postproc_mode(false, false, true);
            if (postprocid == 2)
                g_rvm->set_postproc_mode(true, false, false);
        }
    }

    return JNI_TRUE;
}

// duplicate loadModel for OpencvVideoProcessor
JNIEXPORT jboolean JNICALL Java_io_livekit_android_track_processing_video_OpencvVideoProcessor_loadModel(JNIEnv* env, jobject thiz, jobject assetManager, jint modelid, jint sizeid, jint intrainterid, jint postprocid, jint cpugpu)
{
    if (modelid < 0 || modelid > 1 || sizeid < 0 || sizeid > 6 || intrainterid < 0 || intrainterid > 1 || postprocid < 0 || postprocid > 2 || cpugpu < 0 || cpugpu > 2)
    {
        return JNI_FALSE;
    }

    AAssetManager* mgr = AAssetManager_fromJava(env, assetManager);

    __android_log_print(ANDROID_LOG_DEBUG, "ncnn", "loadModel %p (OpencvVideoProcessor)", mgr);

    const char* modeltypes[2] =
    {
        "mobilenetv3",
        "resnet50"
    };

    const int sizetypes[7] =
    {
        256,
        320,
        384,
        448,
        512,
        576,
        640
    };

    std::string parampath = std::string("rvm_") + modeltypes[(int)modelid] + ".ncnn.param";
    std::string modelpath = std::string("rvm_") + modeltypes[(int)modelid] + ".ncnn.bin";
    bool use_gpu = (int)cpugpu == 1;
    bool use_turnip = (int)cpugpu == 2;

    {
        ncnn::MutexLockGuard g(lock);

        {
            // reset inter feats
            g_feats.r1.release();
            g_feats.r2.release();
            g_feats.r3.release();
            g_feats.r4.release();

            static int old_modelid = 0;
            static int old_cpugpu = 0;
            if (modelid != old_modelid || cpugpu != old_cpugpu)
            {
                delete g_rvm;
                g_rvm = 0;
            }
            old_modelid = modelid;
            old_cpugpu = cpugpu;

            ncnn::destroy_gpu_instance();

            if (use_turnip)
            {
                ncnn::create_gpu_instance("libvulkan_freedreno.so");
            }
            else if (use_gpu)
            {
                ncnn::create_gpu_instance();
            }

            if (!g_rvm)
            {
                g_rvm = new RVM;

                g_rvm->load(mgr, parampath.c_str(), modelpath.c_str(), use_gpu || use_turnip);
            }
            g_rvm->set_model_type((int)modelid);
            g_rvm->set_target_size(sizetypes[(int)sizeid]);
            g_rvm->set_intra_inter((int)intrainterid);

            if (postprocid == 0)
                g_rvm->set_postproc_mode(false, true, false);
            if (postprocid == 1)
                g_rvm->set_postproc_mode(false, false, true);
            if (postprocid == 2)
                g_rvm->set_postproc_mode(true, false, false);
        }
    }

    return JNI_TRUE;
}

// public native boolean openCamera(int facing);
JNIEXPORT jboolean JNICALL Java_io_livekit_android_track_processing_video_RVMNcnn_openCamera(JNIEnv* env, jobject thiz, jint facing)
{
    if (facing < 0 || facing > 1)
        return JNI_FALSE;

    __android_log_print(ANDROID_LOG_DEBUG, "ncnn", "openCamera %d", facing);

    g_camera->open((int)facing);

    return JNI_TRUE;
}

// public native boolean closeCamera();
JNIEXPORT jboolean JNICALL Java_io_livekit_android_track_processing_video_RVMNcnn_closeCamera(JNIEnv* env, jobject thiz)
{
    __android_log_print(ANDROID_LOG_DEBUG, "ncnn", "closeCamera");

    g_camera->close();

    return JNI_TRUE;
}

// public native boolean setOutputWindow(Surface surface);
JNIEXPORT jboolean JNICALL Java_io_livekit_android_track_processing_video_RVMNcnn_setOutputWindow(JNIEnv* env, jobject thiz, jobject surface)
{
    ANativeWindow* win = ANativeWindow_fromSurface(env, surface);

    __android_log_print(ANDROID_LOG_DEBUG, "ncnn", "setOutputWindow %p", win);

    // Set buffer geometry to 640x640, keep current format (0)
    ANativeWindow_setBuffersGeometry(win, 640, 640, 0);

    g_camera->set_window(win);

    return JNI_TRUE;
}

// public native boolean setBackgroundImage(Bitmap bitmap);
JNIEXPORT jboolean JNICALL Java_io_livekit_android_track_processing_video_RVMNcnn_setBackgroundImage(JNIEnv* env, jobject thiz, jobject bitmap)
{
    if (!bitmap)
    {
        ncnn::MutexLockGuard g(lock);
        if (g_rvm)
        {
            g_rvm->clear_background_image();
        }
        return JNI_TRUE;
    }

    AndroidBitmapInfo info;
    if (AndroidBitmap_getInfo(env, bitmap, &info) < 0)
    {
        return JNI_FALSE;
    }

    if (info.format != ANDROID_BITMAP_FORMAT_RGBA_8888 && info.format != ANDROID_BITMAP_FORMAT_RGB_565)
    {
        return JNI_FALSE;
    }

    void* pixels;
    if (AndroidBitmap_lockPixels(env, bitmap, &pixels) < 0)
    {
        return JNI_FALSE;
    }

    cv::Mat mat;
    if (info.format == ANDROID_BITMAP_FORMAT_RGBA_8888)
    {
        mat = cv::Mat(info.height, info.width, CV_8UC4, pixels);
    }
    else if (info.format == ANDROID_BITMAP_FORMAT_RGB_565)
    {
        mat = cv::Mat(info.height, info.width, CV_8UC2, pixels);
        cv::cvtColor(mat, mat, cv::COLOR_BGR5652BGR);
    }

    cv::Mat background_bgr;
    if (mat.channels() == 4)
    {
        cv::cvtColor(mat, background_bgr, cv::COLOR_RGBA2BGR);
    }
    else
    {
        background_bgr = mat.clone();
    }

    AndroidBitmap_unlockPixels(env, bitmap);

    {
        ncnn::MutexLockGuard g(lock);
        if (g_rvm)
        {
            g_rvm->set_background_image(background_bgr);
        }
    }

    return JNI_TRUE;
}



// public native boolean processFrame(long rgbaAddr, int width, int height, int rotation);
JNIEXPORT jboolean JNICALL Java_io_livekit_android_track_processing_video_RVMNcnn_processFrame(JNIEnv* env, jobject thiz, jlong rgbaAddr, jint width, jint height, jint rotation)
{
    cv::Mat* rgba = (cv::Mat*)rgbaAddr;
    if (rgba == nullptr || rgba->empty()) {
        return JNI_FALSE;
    }

    ncnn::MutexLockGuard g(lock);

    if (g_rvm)
    {
        cv::Mat input = *rgba;
        
        // 根据旋转角度处理图像
        cv::Mat rotated;
        if (rotation == 90) {
            cv::rotate(input, rotated, cv::ROTATE_90_CLOCKWISE);
        } else if (rotation == 180) {
            cv::rotate(input, rotated, cv::ROTATE_180);
        } else if (rotation == 270) {
            cv::rotate(input, rotated, cv::ROTATE_90_COUNTERCLOCKWISE);
        } else {
            rotated = input;
        }

        cv::Mat fgr;
        cv::Mat pha;
        cv::Mat seg;
        
        // 处理图像
        g_rvm->detect(rotated, g_feats, fgr, pha, seg);
        
        // 应用虚拟背景（这里使用绿色背景作为示例）
        cv::Mat result;
        g_rvm->draw(rotated, fgr, pha, seg);
        
        // 将结果复制回原矩阵
        if (rotation == 90) {
            cv::rotate(rotated, result, cv::ROTATE_90_COUNTERCLOCKWISE);
            result.copyTo(*rgba);
        } else if (rotation == 180) {
            cv::rotate(rotated, result, cv::ROTATE_180);
            result.copyTo(*rgba);
        } else if (rotation == 270) {
            cv::rotate(rotated, result, cv::ROTATE_90_CLOCKWISE);
            result.copyTo(*rgba);
        } else {
            rotated.copyTo(*rgba);
        }

        return JNI_TRUE;
    }

    return JNI_FALSE;
}

// process I420 in/out without Java-side OpenCV
// signature: Java_io_livekit_android_track_processing_video_OpencvVideoProcessor_processI420ToI420
JNIEXPORT jboolean JNICALL Java_io_livekit_android_track_processing_video_OpencvVideoProcessor_processI420ToI420(
    JNIEnv* env, jclass,
    jobject yBuf, jint yStride,
    jobject uBuf, jint uStride,
    jobject vBuf, jint vStride,
    jint width, jint height, jint rotation,
    jobject outYBuf, jint outYStride,
    jobject outUBuf, jint outUStride,
    jobject outVBuf, jint outVStride)
{
    if (!yBuf || !uBuf || !vBuf || !outYBuf || !outUBuf || !outVBuf || width <= 0 || height <= 0)
        return JNI_FALSE;

    uint8_t* yPtr = (uint8_t*)env->GetDirectBufferAddress(yBuf);
    uint8_t* uPtr = (uint8_t*)env->GetDirectBufferAddress(uBuf);
    uint8_t* vPtr = (uint8_t*)env->GetDirectBufferAddress(vBuf);
    uint8_t* outYPtr = (uint8_t*)env->GetDirectBufferAddress(outYBuf);
    uint8_t* outUPtr = (uint8_t*)env->GetDirectBufferAddress(outUBuf);
    uint8_t* outVPtr = (uint8_t*)env->GetDirectBufferAddress(outVBuf);

    if (!yPtr || !uPtr || !vPtr || !outYPtr || !outUPtr || !outVPtr)
        return JNI_FALSE;

    // Pack input planes with stride into a contiguous I420 buffer
    const int yH = height;
    const int uvH = height / 2;
    const int yW = width;
    const int uvW = width / 2;

    const int ySize = yW * yH;
    const int uSize = uvW * uvH;
    const int vSize = uvW * uvH;

    std::vector<uint8_t> i420_in(ySize + uSize + vSize);
    uint8_t* inY = i420_in.data();
    uint8_t* inU = inY + ySize;
    uint8_t* inV = inU + uSize;

    for (int r = 0; r < yH; ++r) {
        memcpy(inY + r * yW, yPtr + r * yStride, yW);
    }
    for (int r = 0; r < uvH; ++r) {
        memcpy(inU + r * uvW, uPtr + r * uStride, uvW);
        memcpy(inV + r * uvW, vPtr + r * vStride, uvW);
    }

    // Wrap as a single-channel Mat (H + H/2) x W and convert to BGR
    cv::Mat i420_mat(height + height / 2, width, CV_8UC1, i420_in.data());
    cv::Mat rgb;
    cv::cvtColor(i420_mat, rgb, cv::COLOR_YUV2RGB_I420);
    // Rotate to upright orientation for the model
    if (rotation == 90) {
        cv::rotate(rgb, rgb, cv::ROTATE_90_CLOCKWISE);
    } else if (rotation == 180) {
        cv::rotate(rgb, rgb, cv::ROTATE_180);
    } else if (rotation == 270) {
        cv::rotate(rgb, rgb, cv::ROTATE_90_COUNTERCLOCKWISE);
    }

    // Process with RVM
    {
        ncnn::MutexLockGuard g(lock);
        if (g_rvm) {
            cv::Mat fgr, pha, seg;
            g_rvm->detect(rgb, g_feats, fgr, pha, seg);
            g_rvm->draw(rgb, fgr, pha, seg);
        } else {
            draw_unsupported(rgb);
        }
    }

    // Rotate back to original orientation before returning to I420
    if (rotation == 90) {
        cv::rotate(rgb, rgb, cv::ROTATE_90_COUNTERCLOCKWISE);
    } else if (rotation == 180) {
        cv::rotate(rgb, rgb, cv::ROTATE_180);
    } else if (rotation == 270) {
        cv::rotate(rgb, rgb, cv::ROTATE_90_CLOCKWISE);
    }

    // Convert back to I420
    cv::Mat i420_out;
    cv::cvtColor(rgb, i420_out, cv::COLOR_RGB2YUV_I420);
    if (i420_out.empty() || i420_out.cols != width || i420_out.rows != height + height / 2)
        return JNI_FALSE;

    const uint8_t* outBase = i420_out.ptr<uint8_t>(0);
    const uint8_t* srcY = outBase;
    const uint8_t* srcU = srcY + ySize;
    const uint8_t* srcV = srcU + uSize;

    // Write back to output planes honoring strides
    for (int r = 0; r < yH; ++r) {
        memcpy(outYPtr + r * outYStride, srcY + r * yW, yW);
    }
    for (int r = 0; r < uvH; ++r) {
        memcpy(outUPtr + r * outUStride, srcU + r * uvW, uvW);
        memcpy(outVPtr + r * outVStride, srcV + r * uvW, uvW);
    }

    // We ignore input 'rotation' here and unify to 180 at Java metadata level
    (void)rotation;

    return JNI_TRUE;
}
}