/* package ndk/jni/speex_bridge.c */

/*********************************************************************
 * speex_bridge.c
 * Created at: Sun 17 May 2026 22:49:49 +0200
 * Updated at: Sun 28 Jun 2026 22:00:00 +0200 by Anton Afanasyeu <a.afanasieff@gmail.com>
 * Contributors:
 *    - Anton Afanasyeu <a.afanasieff@gmail.com>
 *    - Cursor Agent (project assistant)
 **********************************************************************/
#include <jni.h>
#include <stdlib.h>
#include <string.h>
#include <android/log.h>

#define LOG_TAG "SpeexBridge"
#define LOGE(...) __android_log_print(ANDROID_LOG_ERROR, LOG_TAG, __VA_ARGS__)
#define LOGI(...) __android_log_print(ANDROID_LOG_INFO,  LOG_TAG, __VA_ARGS__)

#ifndef ANDROIDCAST_HAVE_SPEEX
#define ANDROIDCAST_HAVE_SPEEX 0
#endif

#if ANDROIDCAST_HAVE_SPEEX
#include <speex/speex.h>
#include <speex/speex_header.h>

#define SPEEX_ENC_MAGIC 0x53505845L /* SPXE */
#define SPEEX_DEC_MAGIC 0x53505844L /* SPXD */

/* Speex frame sizes per mode:
 *   narrowband  (8 kHz): 160 samples/frame
 *   wideband   (16 kHz): 320 samples/frame
 *   ultra-wide (32 kHz): 640 samples/frame */

typedef struct {
    unsigned long  magic;
    void          *state;
    SpeexBits      bits;
    int            sample_rate;
    int            frame_size;     /* samples per frame */
    int            encoded_bytes;  /* bytes per encoded frame */
} SpeexEncCtx;

typedef struct {
    unsigned long  magic;
    void          *state;
    SpeexBits      bits;
    int            sample_rate;
    int            frame_size;
} SpeexDecCtx;

static SpeexEncCtx *enc_from_handle(jlong h) {
    if (h == 0) return NULL;
    SpeexEncCtx *ctx = (SpeexEncCtx *)(intptr_t)h;
    if (ctx->magic != SPEEX_ENC_MAGIC) return NULL;
    return ctx;
}

static SpeexDecCtx *dec_from_handle(jlong h) {
    if (h == 0) return NULL;
    SpeexDecCtx *ctx = (SpeexDecCtx *)(intptr_t)h;
    if (ctx->magic != SPEEX_DEC_MAGIC) return NULL;
    return ctx;
}

static const SpeexMode *mode_for_rate(int sr) {
    if (sr <= 8000)  return &speex_nb_mode;
    if (sr <= 16000) return &speex_wb_mode;
    return &speex_uwb_mode;
}
#endif /* ANDROIDCAST_HAVE_SPEEX */

/* ── availability probe ───────────────────────────────────────────── */

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

/* ── encoder ────────────────────────────────────────────────────────── */

JNIEXPORT jlong JNICALL
Java_com_foxx_androidcast_media_codec_jni_NativeCodecBridge_nativeSpeexEncoderCreate(
        JNIEnv *env, jclass clazz, jint sampleRate, jint quality) {
    (void)env; (void)clazz;
#if ANDROIDCAST_HAVE_SPEEX
    const SpeexMode *mode = mode_for_rate(sampleRate);
    void *state = speex_encoder_init(mode);
    if (!state) { LOGE("speex_encoder_init failed"); return 0L; }

    int q = (quality >= 1 && quality <= 10) ? quality : 8;
    speex_encoder_ctl(state, SPEEX_SET_QUALITY, &q);

    /* quality 8 wideband → ~15 kbps */
    int frame_size = 0;
    speex_encoder_ctl(state, SPEEX_GET_FRAME_SIZE, &frame_size);

    SpeexEncCtx *ctx = (SpeexEncCtx *)calloc(1, sizeof(SpeexEncCtx));
    if (!ctx) { speex_encoder_destroy(state); return 0L; }
    ctx->magic       = SPEEX_ENC_MAGIC;
    ctx->state       = state;
    ctx->sample_rate = sampleRate;
    ctx->frame_size  = frame_size;
    speex_bits_init(&ctx->bits);
    /* Estimate encoded size: ~38 bytes for quality-8 wideband frame */
    ctx->encoded_bytes = 64;
    LOGI("encoder created sr=%d fs=%d q=%d", sampleRate, frame_size, q);
    return (jlong)(intptr_t)ctx;
#else
    (void)sampleRate; (void)quality;
    return 0L;
#endif
}

JNIEXPORT jbyteArray JNICALL
Java_com_foxx_androidcast_media_codec_jni_NativeCodecBridge_nativeSpeexEncode(
        JNIEnv *env, jclass clazz,
        jlong handle, jbyteArray pcm16le, jint length) {
    (void)clazz;
#if ANDROIDCAST_HAVE_SPEEX
    SpeexEncCtx *ctx = enc_from_handle(handle);
    if (!ctx || pcm16le == NULL || length <= 0) return NULL;

    jbyte *pcm_bytes = (*env)->GetByteArrayElements(env, pcm16le, NULL);
    if (!pcm_bytes) return NULL;

    speex_bits_reset(&ctx->bits);
    speex_encode_int(ctx->state, (spx_int16_t *)pcm_bytes, &ctx->bits);
    (*env)->ReleaseByteArrayElements(env, pcm16le, pcm_bytes, JNI_ABORT);

    int nbytes = speex_bits_nbytes(&ctx->bits);
    unsigned char *buf = (unsigned char *)malloc(nbytes);
    if (!buf) return NULL;
    speex_bits_write(&ctx->bits, (char *)buf, nbytes);

    jbyteArray result = (*env)->NewByteArray(env, nbytes);
    if (result) {
        (*env)->SetByteArrayRegion(env, result, 0, nbytes, (jbyte *)buf);
    }
    free(buf);
    return result;
#else
    (void)handle; (void)pcm16le; (void)length;
    return NULL;
#endif
}

JNIEXPORT void JNICALL
Java_com_foxx_androidcast_media_codec_jni_NativeCodecBridge_nativeSpeexEncoderRelease(
        JNIEnv *env, jclass clazz, jlong handle) {
    (void)env; (void)clazz;
#if ANDROIDCAST_HAVE_SPEEX
    SpeexEncCtx *ctx = enc_from_handle(handle);
    if (!ctx) return;
    speex_bits_destroy(&ctx->bits);
    speex_encoder_destroy(ctx->state);
    ctx->magic = 0;
    free(ctx);
#else
    (void)handle;
#endif
}

JNIEXPORT jint JNICALL
Java_com_foxx_androidcast_media_codec_jni_NativeCodecBridge_nativeSpeexEncoderGetFrameSize(
        JNIEnv *env, jclass clazz, jlong handle) {
    (void)env; (void)clazz;
#if ANDROIDCAST_HAVE_SPEEX
    SpeexEncCtx *ctx = enc_from_handle(handle);
    return ctx ? ctx->frame_size : 0;
#else
    (void)handle; return 0;
#endif
}

/* ── decoder ────────────────────────────────────────────────────────── */

JNIEXPORT jlong JNICALL
Java_com_foxx_androidcast_media_codec_jni_NativeCodecBridge_nativeSpeexDecoderCreate(
        JNIEnv *env, jclass clazz, jint sampleRate) {
    (void)env; (void)clazz;
#if ANDROIDCAST_HAVE_SPEEX
    const SpeexMode *mode = mode_for_rate(sampleRate);
    void *state = speex_decoder_init(mode);
    if (!state) { LOGE("speex_decoder_init failed"); return 0L; }

    int enh = 1;
    speex_decoder_ctl(state, SPEEX_SET_ENH, &enh);

    int frame_size = 0;
    speex_decoder_ctl(state, SPEEX_GET_FRAME_SIZE, &frame_size);

    SpeexDecCtx *ctx = (SpeexDecCtx *)calloc(1, sizeof(SpeexDecCtx));
    if (!ctx) { speex_decoder_destroy(state); return 0L; }
    ctx->magic       = SPEEX_DEC_MAGIC;
    ctx->state       = state;
    ctx->sample_rate = sampleRate;
    ctx->frame_size  = frame_size;
    speex_bits_init(&ctx->bits);
    LOGI("decoder created sr=%d fs=%d", sampleRate, frame_size);
    return (jlong)(intptr_t)ctx;
#else
    (void)sampleRate;
    return 0L;
#endif
}

JNIEXPORT jbyteArray JNICALL
Java_com_foxx_androidcast_media_codec_jni_NativeCodecBridge_nativeSpeexDecode(
        JNIEnv *env, jclass clazz,
        jlong handle, jbyteArray speex_data, jint length) {
    (void)clazz;
#if ANDROIDCAST_HAVE_SPEEX
    SpeexDecCtx *ctx = dec_from_handle(handle);
    if (!ctx) return NULL;

    jbyte *in_bytes = NULL;
    if (speex_data != NULL && length > 0) {
        in_bytes = (*env)->GetByteArrayElements(env, speex_data, NULL);
        speex_bits_read_from(&ctx->bits, (char *)in_bytes, length);
    } else {
        /* PLC: feed 0 bytes to trigger packet loss concealment */
        speex_bits_reset(&ctx->bits);
    }

    spx_int16_t pcm_buf[640]; /* max uwb frame */
    speex_decode_int(ctx->state, in_bytes ? &ctx->bits : NULL, pcm_buf);

    if (in_bytes) (*env)->ReleaseByteArrayElements(env, speex_data, in_bytes, JNI_ABORT);

    int pcm_bytes = ctx->frame_size * 2; /* 16-bit mono */
    jbyteArray result = (*env)->NewByteArray(env, pcm_bytes);
    if (result) {
        (*env)->SetByteArrayRegion(env, result, 0, pcm_bytes, (jbyte *)pcm_buf);
    }
    return result;
#else
    (void)handle; (void)speex_data; (void)length;
    return NULL;
#endif
}

JNIEXPORT void JNICALL
Java_com_foxx_androidcast_media_codec_jni_NativeCodecBridge_nativeSpeexDecoderRelease(
        JNIEnv *env, jclass clazz, jlong handle) {
    (void)env; (void)clazz;
#if ANDROIDCAST_HAVE_SPEEX
    SpeexDecCtx *ctx = dec_from_handle(handle);
    if (!ctx) return;
    speex_bits_destroy(&ctx->bits);
    speex_decoder_destroy(ctx->state);
    ctx->magic = 0;
    free(ctx);
#else
    (void)handle;
#endif
}