#include <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 */

#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_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;
}

#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 jbyteArray JNICALL
Java_com_foxx_androidcast_media_codec_jni_NativeCodecBridge_nativeVpxEncodeYuv(
        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) {
                (*env)->SetByteArrayRegion(env, out, 0, (jsize) pkt->data.frame.sz,
                        (const jbyte *) pkt->data.frame.buf);
            }
            ctx->frame_count++;
            return out;
        }
    }
    return NULL;
#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
}