#include <jni.h>
#include <android/bitmap.h>
#include <android/log.h>
#include <cstdint>
#include <cstring>
#include <algorithm>

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

template <typename T>
static inline T clamp(T val, T minVal, T maxVal) {
    return (val < minVal) ? minVal : (val > maxVal) ? maxVal : val;
}

int quantizeTo16Levels(int gray) {
    return (gray * 15 + 127) / 255 * 17;
}

extern "C"
JNIEXPORT jobject JNICALL
Java_moe_fuquan_einkdither_DitherUtils_applyDithering(JNIEnv *env, jclass clazz, jobject bitmap) {
    if (bitmap == nullptr) return nullptr;

    AndroidBitmapInfo info;
    void* pixels = nullptr;

    if (AndroidBitmap_getInfo(env, bitmap, &info) < 0){
        LOGE("AndroidBitmap_getInfo failed");
        return bitmap;
    }
//    if (info.format != ANDROID_BITMAP_FORMAT_RGBA_8888) return bitmap;
    if (AndroidBitmap_lockPixels(env, bitmap, &pixels) < 0){
        LOGE("AndroidBitmap_getInfo failed");
        return bitmap;
    }

    int width = info.width;
    int height = info.height;
    uint32_t* line = (uint32_t*)pixels;

    // 转换为灰度数组
    int** gray = new int*[height];
    for (int y = 0; y < height; y++) {
        gray[y] = new int[width];
        for (int x = 0; x < width; x++) {
            uint32_t pixel = line[y * width + x];
            int alpha = (pixel >> 24) & 0xFF;
            int r = (pixel >> 16) & 0xFF;
            int g = (pixel >> 8) & 0xFF;
            int b = (pixel) & 0xFF;

            if (alpha == 0)
                gray[y][x] = -1;
            else
                gray[y][x] = (r * 299 + g * 587 + b * 114) / 1000;
        }
    }

    // Floyd-Steinberg 抖动
    for (int y = 0; y < height; y++) {
        for (int x = 0; x < width; x++) {
            if (gray[y][x] == -1) {
                line[y * width + x] = 0x00000000;
                continue;
            }

            int oldPixel = gray[y][x];
            int newPixel = quantizeTo16Levels(oldPixel);
            int error = oldPixel - newPixel;
            gray[y][x] = newPixel;

            if (x + 1 < width && gray[y][x + 1] != -1)
                gray[y][x + 1] = clamp(gray[y][x + 1] + error / 2, 0, 255);
            if (y + 1 < height && gray[y + 1][x] != -1)
                gray[y + 1][x] = clamp(gray[y + 1][x] + error / 2, 0, 255);

            uint8_t g = gray[y][x];
            line[y * width + x] = 0xFF000000 | (g << 16) | (g << 8) | g;
        }
    }

    // 清理
    for (int i = 0; i < height; i++) delete[] gray[i];
    delete[] gray;

    AndroidBitmap_unlockPixels(env, bitmap);
    return bitmap;
}