/* package ndk/jni/libvpx_bridge.c */

/*********************************************************************
 * libvpx_bridge.c
 * Created at: Sun 17 May 2026 22:49:49 +0200
 * Updated at: Wed 20 May 2026 15:17:13 +0200 by Anton Afanasyeu <a.afanasieff@gmail.com>
 * Commit: 5d8e82d2e60a21fff3138d2a394ee4e8b4c6dcb8
 * Contributors:
 *    - Anton Afanasyeu <a.afanasieff@gmail.com> (6 commits, 432 lines)
 *    - Cursor Agent (project assistant)
 * Digest: SHA256	3bdffcead822ea7b00582f416ffe18d3f885f8b91ddab3e87752c23dcf326a02
**********************************************************************/
#include <jni.h>
#include <android/native_window.h>
#include <android/native_window_jni.h>
#include <stdlib.h>
#include <string.h>

#ifndef ANDROIDCAST_HAVE_LIBVPX
#define ANDROIDCAST_HAVE_LIBVPX 0
#endif

#if ANDROIDCAST_HAVE_LIBVPX
#include <vpx/vpx_encoder.h>
#include <vpx/vp8cx.h>
#include <vpx/vpx_decoder.h>
#include <vpx/vp8dx.h>
#endif

#define VPX_ENC_MAGIC 0x56505845L /* VPXE */
#define VPX_DEC_MAGIC 0x56505844L /* VPXD */

#if ANDROIDCAST_HAVE_LIBVPX

typedef struct {
    unsigned long magic;
    vpx_codec_ctx_t codec;
    vpx_image_t img;
    int width;
    int height;
    int frame_count;
    int force_keyframe;
} VpxEncoderCtx;

static vpx_codec_iface_t *pick_encoder_iface(const char *mime) {
    if (mime != NULL && strstr(mime, "VP9") != NULL) {
        return vpx_codec_vp9_cx();
    }
    return vpx_codec_vp8_cx();
}

static vpx_codec_iface_t *pick_decoder_iface(const char *mime) {
    if (mime != NULL && strstr(mime, "VP9") != NULL) {
        return vpx_codec_vp9_dx();
    }
    return vpx_codec_vp8_dx();
}

static vpx_img_fmt_t img_fmt(void) {
    return VPX_IMG_FMT_I420;
}

static void copy_i420_to_image(const uint8_t *yuv, vpx_image_t *img, int width, int height) {
    int y_stride = img->stride[VPX_PLANE_Y];
    int uv_stride = img->stride[VPX_PLANE_U];
    const uint8_t *src_y = yuv;
    const uint8_t *src_u = yuv + width * height;
    const uint8_t *src_v = src_u + (width / 2) * (height / 2);
    for (int row = 0; row < height; row++) {
        memcpy(img->planes[VPX_PLANE_Y] + row * y_stride, src_y + row * width, (size_t) width);
    }
    int chroma_h = height / 2;
    int chroma_w = width / 2;
    for (int row = 0; row < chroma_h; row++) {
        memcpy(img->planes[VPX_PLANE_U] + row * uv_stride, src_u + row * chroma_w, (size_t) chroma_w);
        memcpy(img->planes[VPX_PLANE_V] + row * uv_stride, src_v + row * chroma_w, (size_t) chroma_w);
    }
}

static VpxEncoderCtx *enc_from_handle(jlong handle) {
    VpxEncoderCtx *ctx = (VpxEncoderCtx *) (intptr_t) handle;
    if (ctx == NULL || ctx->magic != VPX_ENC_MAGIC) {
        return NULL;
    }
    return ctx;
}

typedef struct {
    unsigned long magic;
    vpx_codec_ctx_t codec;
} VpxDecoderCtx;

static VpxDecoderCtx *dec_from_handle(jlong handle) {
    VpxDecoderCtx *ctx = (VpxDecoderCtx *) (intptr_t) handle;
    if (ctx == NULL || ctx->magic != VPX_DEC_MAGIC) {
        return NULL;
    }
    return ctx;
}

static void render_i420_to_window(ANativeWindow *window, const vpx_image_t *img) {
    if (window == NULL || img == NULL) {
        return;
    }
    const int width = (int) img->d_w;
    const int height = (int) img->d_h;
    if (width <= 0 || height <= 0) {
        return;
    }
    ANativeWindow_setBuffersGeometry(window, width, height, WINDOW_FORMAT_RGBA_8888);
    ANativeWindow_Buffer buffer;
    if (ANativeWindow_lock(window, &buffer, NULL) != 0) {
        return;
    }
    uint8_t *dst = (uint8_t *) buffer.bits;
    const int dst_stride = buffer.stride * 4;
    const uint8_t *y_plane = img->planes[VPX_PLANE_Y];
    const uint8_t *u_plane = img->planes[VPX_PLANE_U];
    const uint8_t *v_plane = img->planes[VPX_PLANE_V];
    const int y_stride = img->stride[VPX_PLANE_Y];
    const int u_stride = img->stride[VPX_PLANE_U];
    const int v_stride = img->stride[VPX_PLANE_V];

    for (int row = 0; row < height; row++) {
        uint8_t *row_dst = dst + row * dst_stride;
        for (int col = 0; col < width; col++) {
            int y = y_plane[row * y_stride + col];
            int u = u_plane[(row / 2) * u_stride + (col / 2)] - 128;
            int v = v_plane[(row / 2) * v_stride + (col / 2)] - 128;
            int r = y + ((351 * v) >> 8);
            int g = y - ((179 * v + 86 * u) >> 8);
            int b = y + ((443 * u) >> 8);
            if (r < 0) {
                r = 0;
            } else if (r > 255) {
                r = 255;
            }
            if (g < 0) {
                g = 0;
            } else if (g > 255) {
                g = 255;
            }
            if (b < 0) {
                b = 0;
            } else if (b > 255) {
                b = 255;
            }
            row_dst[col * 4] = (uint8_t) r;
            row_dst[col * 4 + 1] = (uint8_t) g;
            row_dst[col * 4 + 2] = (uint8_t) b;
            row_dst[col * 4 + 3] = 255;
        }
    }
    ANativeWindow_unlockAndPost(window);
}

#endif /* ANDROIDCAST_HAVE_LIBVPX */

JNIEXPORT jboolean JNICALL
Java_com_foxx_androidcast_media_codec_jni_NativeCodecBridge_nativeIsLibvpxAvailable(
        JNIEnv *env, jclass clazz) {
    (void) env;
    (void) clazz;
    return ANDROIDCAST_HAVE_LIBVPX ? JNI_TRUE : JNI_FALSE;
}

JNIEXPORT jlong JNICALL
Java_com_foxx_androidcast_media_codec_jni_NativeCodecBridge_nativeVpxEncoderCreate(
        JNIEnv *env, jclass clazz, jstring mime, jint width, jint height, jint bitrateKbps) {
    (void) clazz;
#if !ANDROIDCAST_HAVE_LIBVPX
    (void) env;
    (void) mime;
    (void) width;
    (void) height;
    (void) bitrateKbps;
    return 0L;
#else
    if (width <= 0 || height <= 0 || bitrateKbps <= 0) {
        return 0L;
    }
    const char *mime_utf = (*env)->GetStringUTFChars(env, mime, NULL);
    VpxEncoderCtx *ctx = (VpxEncoderCtx *) calloc(1, sizeof(VpxEncoderCtx));
    if (ctx == NULL) {
        (*env)->ReleaseStringUTFChars(env, mime, mime_utf);
        return 0L;
    }
    ctx->magic = VPX_ENC_MAGIC;
    ctx->width = width;
    ctx->height = height;

    vpx_codec_enc_cfg_t cfg;
    vpx_codec_iface_t *iface = pick_encoder_iface(mime_utf);
    if (vpx_codec_enc_config_default(iface, &cfg, 0) != VPX_CODEC_OK) {
        free(ctx);
        (*env)->ReleaseStringUTFChars(env, mime, mime_utf);
        return 0L;
    }
    cfg.g_w = (unsigned int) width;
    cfg.g_h = (unsigned int) height;
    cfg.rc_target_bitrate = bitrateKbps;
    cfg.g_timebase.num = 1;
    cfg.g_timebase.den = 30;
    cfg.g_error_resilient = VPX_ERROR_RESILIENT_DEFAULT;
    cfg.g_lag_in_frames = 0;

    if (vpx_codec_enc_init(&ctx->codec, iface, &cfg, 0) != VPX_CODEC_OK) {
        free(ctx);
        (*env)->ReleaseStringUTFChars(env, mime, mime_utf);
        return 0L;
    }
    /* VP8 and VP9 both use VP8E_SET_CPUUSED (see libvpx tests / vp9_cx_iface.c). */
    vpx_codec_control(&ctx->codec, VP8E_SET_CPUUSED, 8);

    if (!vpx_img_alloc(&ctx->img, img_fmt(), width, height, 1)) {
        vpx_codec_destroy(&ctx->codec);
        free(ctx);
        (*env)->ReleaseStringUTFChars(env, mime, mime_utf);
        return 0L;
    }
    (*env)->ReleaseStringUTFChars(env, mime, mime_utf);
    return (jlong) (intptr_t) ctx;
#endif
}

JNIEXPORT jobject JNICALL
Java_com_foxx_androidcast_media_codec_jni_NativeCodecBridge_nativeVpxEncodeYuvFrame(
        JNIEnv *env, jclass clazz, jlong handle, jbyteArray yuv, jlong ptsUs, jboolean forceKeyframe) {
    (void) clazz;
#if !ANDROIDCAST_HAVE_LIBVPX
    (void) env;
    (void) handle;
    (void) yuv;
    (void) ptsUs;
    (void) forceKeyframe;
    return NULL;
#else
    VpxEncoderCtx *ctx = enc_from_handle(handle);
    if (ctx == NULL || yuv == NULL) {
        return NULL;
    }
    jsize yuv_len = (*env)->GetArrayLength(env, yuv);
    int need = ctx->width * ctx->height * 3 / 2;
    if (yuv_len < need) {
        return NULL;
    }
    jbyte *yuv_bytes = (*env)->GetByteArrayElements(env, yuv, NULL);
    if (yuv_bytes == NULL) {
        return NULL;
    }
    copy_i420_to_image((const uint8_t *) yuv_bytes, &ctx->img, ctx->width, ctx->height);
    (*env)->ReleaseByteArrayElements(env, yuv, yuv_bytes, JNI_ABORT);

    vpx_enc_frame_flags_t flags = 0;
    if (forceKeyframe || ctx->force_keyframe) {
        flags |= VPX_EFLAG_FORCE_KF;
        ctx->force_keyframe = 0;
    }

    if (vpx_codec_encode(&ctx->codec, &ctx->img, (vpx_codec_pts_t) ptsUs, 1, flags,
                         VPX_DL_REALTIME) != VPX_CODEC_OK) {
        return NULL;
    }

    const vpx_codec_cx_pkt_t *pkt;
    vpx_codec_iter_t iter = NULL;
    while ((pkt = vpx_codec_get_cx_data(&ctx->codec, &iter)) != NULL) {
        if (pkt->kind == VPX_CODEC_CX_FRAME_PKT) {
            jbyteArray out = (*env)->NewByteArray(env, (jsize) pkt->data.frame.sz);
            if (out == NULL) {
                return NULL;
            }
            (*env)->SetByteArrayRegion(env, out, 0, (jsize) pkt->data.frame.sz,
                                       (const jbyte *) pkt->data.frame.buf);
            jboolean key = (pkt->data.frame.flags & VPX_FRAME_IS_KEY) ? JNI_TRUE : JNI_FALSE;
            jclass frame_cls = (*env)->FindClass(
                    env, "com/foxx/androidcast/media/codec/jni/VpxEncodedFrame");
            jmethodID ctor = (*env)->GetMethodID(env, frame_cls, "<init>", "([BZ)V");
            jobject frame = (*env)->NewObject(env, frame_cls, ctor, out, key);
            ctx->frame_count++;
            return frame;
        }
    }
    return NULL;
#endif
}

JNIEXPORT jbyteArray JNICALL
Java_com_foxx_androidcast_media_codec_jni_NativeCodecBridge_nativeVpxEncodeYuv(
        JNIEnv *env, jclass clazz, jlong handle, jbyteArray yuv, jlong ptsUs, jboolean forceKeyframe) {
#if !ANDROIDCAST_HAVE_LIBVPX
    (void) env;
    (void) clazz;
    (void) handle;
    (void) yuv;
    (void) ptsUs;
    (void) forceKeyframe;
    return NULL;
#else
    jobject frame = Java_com_foxx_androidcast_media_codec_jni_NativeCodecBridge_nativeVpxEncodeYuvFrame(
            env, clazz, handle, yuv, ptsUs, forceKeyframe);
    if (frame == NULL) {
        return NULL;
    }
    jclass frame_cls = (*env)->GetObjectClass(env, frame);
    jfieldID data_f = (*env)->GetFieldID(env, frame_cls, "data", "[B");
    return (jbyteArray) (*env)->GetObjectField(env, frame, data_f);
#endif
}

JNIEXPORT void JNICALL
Java_com_foxx_androidcast_media_codec_jni_NativeCodecBridge_nativeVpxEncoderRelease(
        JNIEnv *env, jclass clazz, jlong handle) {
    (void) env;
    (void) clazz;
#if ANDROIDCAST_HAVE_LIBVPX
    VpxEncoderCtx *ctx = enc_from_handle(handle);
    if (ctx == NULL) {
        return;
    }
    ctx->magic = 0;
    vpx_img_free(&ctx->img);
    vpx_codec_destroy(&ctx->codec);
    free(ctx);
#endif
}

JNIEXPORT void JNICALL
Java_com_foxx_androidcast_media_codec_jni_NativeCodecBridge_nativeVpxRequestKeyframe(
        JNIEnv *env, jclass clazz, jlong handle) {
    (void) env;
    (void) clazz;
#if ANDROIDCAST_HAVE_LIBVPX
    VpxEncoderCtx *ctx = enc_from_handle(handle);
    if (ctx != NULL) {
        ctx->force_keyframe = 1;
    }
#endif
}

JNIEXPORT jlong JNICALL
Java_com_foxx_androidcast_media_codec_jni_NativeCodecBridge_nativeVpxDecoderCreate(
        JNIEnv *env, jclass clazz, jstring mime) {
    (void) clazz;
#if !ANDROIDCAST_HAVE_LIBVPX
    (void) env;
    (void) mime;
    return 0L;
#else
    const char *mime_utf = (*env)->GetStringUTFChars(env, mime, NULL);
    VpxDecoderCtx *ctx = (VpxDecoderCtx *) calloc(1, sizeof(VpxDecoderCtx));
    if (ctx == NULL) {
        (*env)->ReleaseStringUTFChars(env, mime, mime_utf);
        return 0L;
    }
    ctx->magic = VPX_DEC_MAGIC;
    vpx_codec_iface_t *iface = pick_decoder_iface(mime_utf);
    vpx_codec_dec_cfg_t cfg = {0};
    cfg.threads = 1;
    /* vpxdec.c: dec_flags |= VPX_CODEC_USE_ERROR_CONCEALMENT when --error-concealment */
    int dec_flags = 0;
    if (vpx_codec_get_caps(iface) & VPX_CODEC_CAP_ERROR_CONCEALMENT) {
        dec_flags |= VPX_CODEC_USE_ERROR_CONCEALMENT;
    }
    if (vpx_codec_dec_init(&ctx->codec, iface, &cfg, dec_flags) != VPX_CODEC_OK) {
        free(ctx);
        (*env)->ReleaseStringUTFChars(env, mime, mime_utf);
        return 0L;
    }
    (*env)->ReleaseStringUTFChars(env, mime, mime_utf);
    return (jlong) (intptr_t) ctx;
#endif
}

JNIEXPORT jboolean JNICALL
Java_com_foxx_androidcast_media_codec_jni_NativeCodecBridge_nativeVpxDecodeFrame(
        JNIEnv *env, jclass clazz, jlong handle, jbyteArray data, jlong ptsUs, jobject surface) {
    (void) clazz;
    (void) ptsUs;
#if !ANDROIDCAST_HAVE_LIBVPX
    (void) env;
    (void) handle;
    (void) data;
    (void) surface;
    return JNI_FALSE;
#else
    VpxDecoderCtx *ctx = dec_from_handle(handle);
    if (ctx == NULL || data == NULL || surface == NULL) {
        return JNI_FALSE;
    }
    jsize len = (*env)->GetArrayLength(env, data);
    if (len <= 0) {
        return JNI_FALSE;
    }
    jbyte *bytes = (*env)->GetByteArrayElements(env, data, NULL);
    if (bytes == NULL) {
        return JNI_FALSE;
    }
    vpx_codec_err_t err = vpx_codec_decode(&ctx->codec, (const uint8_t *) bytes, (unsigned int) len,
                                            NULL, 0);
    (*env)->ReleaseByteArrayElements(env, data, bytes, JNI_ABORT);
    /* With error concealment, vpx_codec_decode may warn yet still output a frame (see vpxdec.c). */
    (void) err;

    ANativeWindow *window = ANativeWindow_fromSurface(env, surface);
    if (window == NULL) {
        return JNI_FALSE;
    }

    vpx_codec_iter_t iter = NULL;
    vpx_image_t *img;
    jboolean rendered = JNI_FALSE;
    while ((img = vpx_codec_get_frame(&ctx->codec, &iter)) != NULL) {
        render_i420_to_window(window, img);
        rendered = JNI_TRUE;
    }
    ANativeWindow_release(window);
    return rendered;
#endif
}

JNIEXPORT void JNICALL
Java_com_foxx_androidcast_media_codec_jni_NativeCodecBridge_nativeVpxDecoderRelease(
        JNIEnv *env, jclass clazz, jlong handle) {
    (void) env;
    (void) clazz;
#if ANDROIDCAST_HAVE_LIBVPX
    VpxDecoderCtx *ctx = dec_from_handle(handle);
    if (ctx == NULL) {
        return;
    }
    ctx->magic = 0;
    vpx_codec_destroy(&ctx->codec);
    free(ctx);
#endif
}