/*
 * The JNI layer over the native sink's C interface (wai/extensions/native-sink.md §9): one
 * function per call of `wai.h`, for `wai.sink.WaiNative`. Nothing here decides anything: every
 * input goes to libwai as the caller gave it, every answer comes back as the bytes libwai wrote,
 * and every status other than WAI_OK becomes a `wai.sink.WaiSinkException` carrying that status
 * and `wai_last_error`'s message, read on the same thread before any other call.
 *
 * Strings arrive as UTF-8 byte arrays (the Kotlin side encodes them), so text reaches libwai as
 * real UTF-8 rather than the JVM's modified UTF-8; a NUL inside one is refused here with
 * WAI_ERR_INVALID_ARG, since the C interface would read it as the end of the string.
 */
#include <jni.h>
#include <stdint.h>
#include <stdlib.h>
#include <string.h>

#include "wai.h"

static jclass exc_class;   /* wai.sink.WaiSinkException, a global reference */
static jmethodID exc_init; /* WaiSinkException(int status, byte[] message) */

JNIEXPORT jint JNICALL JNI_OnLoad(JavaVM *vm, void *reserved) {
    (void)reserved;
    JNIEnv *env;
    if ((*vm)->GetEnv(vm, (void **)&env, JNI_VERSION_1_6) != JNI_OK) return JNI_ERR;
    jclass c = (*env)->FindClass(env, "wai/sink/WaiSinkException");
    if (c == NULL) return JNI_ERR;
    exc_class = (jclass)(*env)->NewGlobalRef(env, c);
    (*env)->DeleteLocalRef(env, c);
    if (exc_class == NULL) return JNI_ERR;
    exc_init = (*env)->GetMethodID(env, exc_class, "<init>", "(I[B)V");
    if (exc_init == NULL) return JNI_ERR;
    return JNI_VERSION_1_6;
}

/* Throw WaiSinkException(status, message). A pending exception (an allocation that failed) is
 * left as it is. */
static void throw_status(JNIEnv *env, int status, const char *message) {
    if ((*env)->ExceptionCheck(env)) return;
    size_t n = message ? strlen(message) : 0;
    jbyteArray m = (*env)->NewByteArray(env, (jsize)n);
    if (m == NULL) return;
    if (n) (*env)->SetByteArrayRegion(env, m, 0, (jsize)n, (const jbyte *)message);
    jobject e = (*env)->NewObject(env, exc_class, exc_init, (jint)status, m);
    if (e != NULL) (*env)->Throw(env, (jthrowable)e);
}

/* After a call of the C interface: on any status but WAI_OK, throw it with libwai's message.
 * Returns 1 when the call succeeded. */
static int ok(JNIEnv *env, int status) {
    if (status == WAI_OK) return 1;
    throw_status(env, status, wai_last_error());
    return 0;
}

/* A byte array's contents, copied, or NULL with an exception pending. With `nul`, a NUL byte is
 * appended (a string for the C interface) and a NUL inside is refused. `*len` is the length
 * without the appended NUL. A null array gives NULL with no exception (`*len` 0). */
static uint8_t *take(JNIEnv *env, jbyteArray a, int nul, uintptr_t *len) {
    *len = 0;
    if (a == NULL) return NULL;
    jsize n = (*env)->GetArrayLength(env, a);
    uint8_t *p = (uint8_t *)malloc((size_t)n + 1);
    if (p == NULL) {
        throw_status(env, WAI_ERR_INVALID_ARG, "out of memory copying an argument");
        return NULL;
    }
    if (n) (*env)->GetByteArrayRegion(env, a, 0, n, (jbyte *)p);
    p[n] = 0;
    if (nul && memchr(p, 0, (size_t)n) != NULL) {
        free(p);
        throw_status(env, WAI_ERR_INVALID_ARG, "a string argument holds a NUL byte");
        return NULL;
    }
    *len = (uintptr_t)n;
    return p;
}

/* A required string argument: NULL (with an exception) when absent or refused. */
static char *text(JNIEnv *env, jbyteArray a, const char *what) {
    uintptr_t n;
    if (a == NULL) {
        throw_status(env, WAI_ERR_INVALID_ARG, what);
        return NULL;
    }
    return (char *)take(env, a, 1, &n);
}

/* A required byte argument. A zero-length array is given as a valid pointer and length 0. */
static uint8_t *bytes(JNIEnv *env, jbyteArray a, uintptr_t *len, const char *what) {
    if (a == NULL) {
        *len = 0;
        throw_status(env, WAI_ERR_INVALID_ARG, what);
        return NULL;
    }
    return take(env, a, 0, len);
}

/* The output buffer as a new byte array, the buffer freed either way. */
static jbyteArray give(JNIEnv *env, wai_buffer_t b) {
    jbyteArray r = (*env)->NewByteArray(env, (jsize)b.len);
    if (r != NULL && b.len) (*env)->SetByteArrayRegion(env, r, 0, (jsize)b.len, (const jbyte *)b.data);
    wai_buffer_free(b);
    return r;
}

static const wai_sink_t *sink_of(jlong h) { return (const wai_sink_t *)(intptr_t)h; }

#define EMPTY_BUFFER {NULL, 0, 0}

JNIEXPORT jint JNICALL Java_wai_sink_WaiNative_abiVersion(JNIEnv *env, jclass cls) {
    (void)env; (void)cls;
    return (jint)wai_sink_abi_version();
}

JNIEXPORT jlong JNICALL Java_wai_sink_WaiNative_nowUs(JNIEnv *env, jclass cls) {
    (void)env; (void)cls;
    return (jlong)wai_sink_now_us();
}

JNIEXPORT jlong JNICALL Java_wai_sink_WaiNative_open(JNIEnv *env, jclass cls, jbyteArray config, jbyteArray seed, jlong start_us) {
    (void)cls;
    char *c = text(env, config, "config_json is null");
    if (c == NULL) return 0;
    uintptr_t sn = 0;
    uint8_t *s = take(env, seed, 0, &sn);
    if (seed != NULL && s == NULL) { free(c); return 0; }
    wai_sink_t *out = NULL;
    int st = wai_sink_open(c, s, sn, (uint64_t)start_us, &out);
    free(c);
    free(s);
    if (!ok(env, st)) return 0;
    return (jlong)(intptr_t)out;
}

JNIEXPORT void JNICALL Java_wai_sink_WaiNative_free(JNIEnv *env, jclass cls, jlong h) {
    (void)env; (void)cls;
    wai_sink_free((wai_sink_t *)(intptr_t)h);
}

JNIEXPORT jbyteArray JNICALL Java_wai_sink_WaiNative_info(JNIEnv *env, jclass cls, jlong h) {
    (void)cls;
    wai_buffer_t out = EMPTY_BUFFER;
    if (!ok(env, wai_sink_info(sink_of(h), &out))) return NULL;
    return give(env, out);
}

JNIEXPORT jbyteArray JNICALL Java_wai_sink_WaiNative_receiptName(JNIEnv *env, jclass cls, jbyteArray playlist) {
    (void)cls;
    uintptr_t n;
    uint8_t *p = bytes(env, playlist, &n, "playlist is null");
    if (p == NULL) return NULL;
    wai_buffer_t out = EMPTY_BUFFER;
    int st = wai_sink_receipt_name(p, n, &out);
    free(p);
    if (!ok(env, st)) return NULL;
    return give(env, out);
}

JNIEXPORT jbyteArray JNICALL Java_wai_sink_WaiNative_playlist(JNIEnv *env, jclass cls, jlong h, jbyteArray url, jbyteArray body, jbyteArray receipt) {
    (void)cls;
    char *u = text(env, url, "url is null");
    if (u == NULL) return NULL;
    uintptr_t n, rn = 0;
    uint8_t *b = bytes(env, body, &n, "bytes is null");
    if (b == NULL) { free(u); return NULL; }
    uint8_t *r = take(env, receipt, 0, &rn);
    if (receipt != NULL && r == NULL) { free(u); free(b); return NULL; }
    wai_buffer_t out = EMPTY_BUFFER;
    int st = wai_sink_playlist(sink_of(h), u, b, n, r, rn, &out);
    free(u); free(b); free(r);
    if (!ok(env, st)) return NULL;
    return give(env, out);
}

JNIEXPORT void JNICALL Java_wai_sink_WaiNative_playlistParent(JNIEnv *env, jclass cls, jlong h, jbyteArray url, jbyteArray name, jbyteArray body) {
    (void)cls;
    char *u = text(env, url, "url is null");
    if (u == NULL) return;
    char *nm = text(env, name, "name is null");
    if (nm == NULL) { free(u); return; }
    uintptr_t n = 0;
    uint8_t *b = take(env, body, 0, &n);
    if (body != NULL && b == NULL) { free(u); free(nm); return; }
    int st = wai_sink_playlist_parent(sink_of(h), u, nm, b, n);
    free(u); free(nm); free(b);
    ok(env, st);
}

JNIEXPORT jbyteArray JNICALL Java_wai_sink_WaiNative_objectVerdict(JNIEnv *env, jclass cls, jlong h, jbyteArray meta, jbyteArray body) {
    (void)cls;
    char *m = text(env, meta, "meta_json is null");
    if (m == NULL) return NULL;
    uintptr_t n;
    uint8_t *b = bytes(env, body, &n, "bytes is null");
    if (b == NULL) { free(m); return NULL; }
    wai_buffer_t out = EMPTY_BUFFER;
    int st = wai_sink_object(sink_of(h), m, b, n, &out);
    free(m); free(b);
    if (!ok(env, st)) return NULL;
    return give(env, out);
}

JNIEXPORT void JNICALL Java_wai_sink_WaiNative_event(JNIEnv *env, jclass cls, jlong h, jbyteArray name, jlong t_us) {
    (void)cls;
    char *nm = text(env, name, "name is null");
    if (nm == NULL) return;
    int st = wai_sink_event(sink_of(h), nm, (uint64_t)t_us);
    free(nm);
    ok(env, st);
}

JNIEXPORT jbyteArray JNICALL Java_wai_sink_WaiNative_observe(JNIEnv *env, jclass cls, jlong h, jbyteArray obs) {
    (void)cls;
    char *o = text(env, obs, "obs_json is null");
    if (o == NULL) return NULL;
    wai_buffer_t out = EMPTY_BUFFER;
    int st = wai_sink_observe(sink_of(h), o, &out);
    free(o);
    if (!ok(env, st)) return NULL;
    return give(env, out);
}

JNIEXPORT void JNICALL Java_wai_sink_WaiNative_energy(JNIEnv *env, jclass cls, jlong h, jbyteArray entry) {
    (void)cls;
    char *e = text(env, entry, "entry_json is null");
    if (e == NULL) return;
    int st = wai_sink_energy(sink_of(h), e);
    free(e);
    ok(env, st);
}

JNIEXPORT void JNICALL Java_wai_sink_WaiNative_attach(JNIEnv *env, jclass cls, jlong h, jbyteArray name, jbyteArray body) {
    (void)cls;
    char *nm = text(env, name, "name is null");
    if (nm == NULL) return;
    uintptr_t n;
    uint8_t *b = bytes(env, body, &n, "bytes is null");
    if (b == NULL) { free(nm); return; }
    int st = wai_sink_attach(sink_of(h), nm, b, n);
    free(nm); free(b);
    ok(env, st);
}

JNIEXPORT jbyteArray JNICALL Java_wai_sink_WaiNative_takeCmcd(JNIEnv *env, jclass cls, jlong h) {
    (void)cls;
    wai_buffer_t out = EMPTY_BUFFER;
    if (!ok(env, wai_sink_take_cmcd(sink_of(h), &out))) return NULL;
    return give(env, out);
}

JNIEXPORT jbyteArray JNICALL Java_wai_sink_WaiNative_end(JNIEnv *env, jclass cls, jlong h, jbyteArray ended, jlong t_us) {
    (void)cls;
    char *e = text(env, ended, "ended is null");
    if (e == NULL) return NULL;
    wai_buffer_t out = EMPTY_BUFFER;
    int st = wai_sink_end(sink_of(h), e, (uint64_t)t_us, &out);
    free(e);
    if (!ok(env, st)) return NULL;
    return give(env, out);
}

JNIEXPORT jbyteArray JNICALL Java_wai_sink_WaiNative_verifyBundle(JNIEnv *env, jclass cls, jbyteArray bundle, jbyteArray owner, jbyteArray channels) {
    (void)cls;
    uintptr_t n, on;
    uint8_t *b = bytes(env, bundle, &n, "bundle is null");
    if (b == NULL) return NULL;
    uint8_t *o = bytes(env, owner, &on, "owner is null");
    if (o == NULL) { free(b); return NULL; }
    if (on != 32) {
        free(b); free(o);
        throw_status(env, WAI_ERR_INVALID_ARG, "an owner key is 32 bytes");
        return NULL;
    }
    char *ch = text(env, channels, "channels_json is null");
    if (ch == NULL) { free(b); free(o); return NULL; }
    wai_buffer_t out = EMPTY_BUFFER;
    int st = wai_audit_verify_bundle(b, n, o, ch, &out);
    free(b); free(o); free(ch);
    if (!ok(env, st)) return NULL;
    return give(env, out);
}

/* ---------------------------------------------------------------- the integer decode
 *
 * The decoder of native-sink §13 (`wai_int_*`): a decoder handle, a parameter set handle, and a
 * decode answering two byte arrays, the output's bytes and the answer's JSON. A refused decode is
 * an answer, not an error (WAI_OK with no bytes); only a status other than WAI_OK throws. */

static const wai_int_decoder_t *decoder_of(jlong h) { return (const wai_int_decoder_t *)(intptr_t)h; }
static const wai_int_set_t *set_of(jlong h) { return (const wai_int_set_t *)(intptr_t)h; }

/* { bytes, answer } as a byte[][] of two, both buffers freed either way. */
static jobjectArray give_pair(JNIEnv *env, wai_buffer_t bytes, wai_buffer_t answer) {
    jbyteArray b = give(env, bytes);
    jbyteArray a = give(env, answer);
    if (b == NULL || a == NULL) return NULL;
    jclass ba = (*env)->FindClass(env, "[B");
    if (ba == NULL) return NULL;
    jobjectArray r = (*env)->NewObjectArray(env, 2, ba, NULL);
    (*env)->DeleteLocalRef(env, ba);
    if (r == NULL) return NULL;
    (*env)->SetObjectArrayElement(env, r, 0, b);
    (*env)->SetObjectArrayElement(env, r, 1, a);
    return r;
}

JNIEXPORT jlong JNICALL Java_wai_sink_WaiNative_intOpen(JNIEnv *env, jclass cls, jbyteArray config) {
    (void)cls;
    uintptr_t n = 0;
    char *c = (char *)take(env, config, 1, &n);
    if (config != NULL && c == NULL) return 0;
    wai_int_decoder_t *out = NULL;
    int st = wai_int_decoder_open(c, &out);
    free(c);
    if (!ok(env, st)) return 0;
    return (jlong)(intptr_t)out;
}

JNIEXPORT void JNICALL Java_wai_sink_WaiNative_intFree(JNIEnv *env, jclass cls, jlong h) {
    (void)env; (void)cls;
    wai_int_decoder_free((wai_int_decoder_t *)(intptr_t)h);
}

JNIEXPORT jbyteArray JNICALL Java_wai_sink_WaiNative_intInfo(JNIEnv *env, jclass cls, jlong h) {
    (void)cls;
    wai_buffer_t out = EMPTY_BUFFER;
    if (!ok(env, wai_int_decoder_info(decoder_of(h), &out))) return NULL;
    return give(env, out);
}

JNIEXPORT jlong JNICALL Java_wai_sink_WaiNative_intSetNew(JNIEnv *env, jclass cls) {
    (void)cls;
    wai_int_set_t *out = NULL;
    if (!ok(env, wai_int_set_new(&out))) return 0;
    return (jlong)(intptr_t)out;
}

JNIEXPORT void JNICALL Java_wai_sink_WaiNative_intSetAdd(JNIEnv *env, jclass cls, jlong set, jbyteArray name, jbyteArray body) {
    (void)cls;
    char *nm = text(env, name, "name is null");
    if (nm == NULL) return;
    uintptr_t n;
    uint8_t *b = bytes(env, body, &n, "bytes is null");
    if (b == NULL) { free(nm); return; }
    int st = wai_int_set_add(set_of(set), nm, b, n);
    free(nm); free(b);
    ok(env, st);
}

JNIEXPORT void JNICALL Java_wai_sink_WaiNative_intSetListing(JNIEnv *env, jclass cls, jlong set, jbyteArray listing) {
    (void)cls;
    uintptr_t n;
    uint8_t *b = bytes(env, listing, &n, "listing is null");
    if (b == NULL) return;
    int st = wai_int_set_listing(set_of(set), b, n);
    free(b);
    ok(env, st);
}

JNIEXPORT void JNICALL Java_wai_sink_WaiNative_intSetFree(JNIEnv *env, jclass cls, jlong set) {
    (void)env; (void)cls;
    wai_int_set_free((wai_int_set_t *)(intptr_t)set);
}

JNIEXPORT jobjectArray JNICALL Java_wai_sink_WaiNative_intDecode(JNIEnv *env, jclass cls, jlong h, jbyteArray capability, jbyteArray payload, jlong set, jbyteArray options) {
    (void)cls;
    char *c = text(env, capability, "capability is null");
    if (c == NULL) return NULL;
    uintptr_t n, on = 0;
    uint8_t *p = bytes(env, payload, &n, "payload is null");
    if (p == NULL) { free(c); return NULL; }
    char *o = (char *)take(env, options, 1, &on);
    if (options != NULL && o == NULL) { free(c); free(p); return NULL; }
    wai_buffer_t out_bytes = EMPTY_BUFFER, out_answer = EMPTY_BUFFER;
    int st = wai_int_decode(decoder_of(h), c, p, n, set_of(set), o, &out_bytes, &out_answer);
    free(c); free(p); free(o);
    if (!ok(env, st)) { wai_buffer_free(out_bytes); wai_buffer_free(out_answer); return NULL; }
    return give_pair(env, out_bytes, out_answer);
}

JNIEXPORT jobjectArray JNICALL Java_wai_sink_WaiNative_intDecodeEnvelope(JNIEnv *env, jclass cls, jlong h, jbyteArray envelope, jlong set) {
    (void)cls;
    uintptr_t n;
    uint8_t *e = bytes(env, envelope, &n, "envelope is null");
    if (e == NULL) return NULL;
    wai_buffer_t out_bytes = EMPTY_BUFFER, out_answer = EMPTY_BUFFER;
    int st = wai_int_decode_envelope(decoder_of(h), e, n, set_of(set), &out_bytes, &out_answer);
    free(e);
    if (!ok(env, st)) { wai_buffer_free(out_bytes); wai_buffer_free(out_answer); return NULL; }
    return give_pair(env, out_bytes, out_answer);
}

/* A stream of one integer frame form (native-sink §13.6), one frame an object: open, reset,
 * free, and the next object decoded against it. */

static const wai_int_frames_t *frames_of(jlong h) { return (const wai_int_frames_t *)(intptr_t)h; }

JNIEXPORT jlong JNICALL Java_wai_sink_WaiNative_intFramesOpen(JNIEnv *env, jclass cls, jbyteArray capability, jlong set) {
    (void)cls;
    char *c = text(env, capability, "capability is null");
    if (c == NULL) return 0;
    wai_int_frames_t *out = NULL;
    int st = wai_int_frames_open(c, set_of(set), &out);
    free(c);
    if (!ok(env, st)) return 0;
    return (jlong)(intptr_t)out;
}

JNIEXPORT void JNICALL Java_wai_sink_WaiNative_intFramesReset(JNIEnv *env, jclass cls, jlong s) {
    (void)cls;
    ok(env, wai_int_frames_reset(frames_of(s)));
}

JNIEXPORT void JNICALL Java_wai_sink_WaiNative_intFramesFree(JNIEnv *env, jclass cls, jlong s) {
    (void)env; (void)cls;
    wai_int_frames_free((wai_int_frames_t *)(intptr_t)s);
}

JNIEXPORT jobjectArray JNICALL Java_wai_sink_WaiNative_intDecodeFrame(JNIEnv *env, jclass cls, jlong h, jlong s, jbyteArray envelope) {
    (void)cls;
    uintptr_t n;
    uint8_t *e = bytes(env, envelope, &n, "envelope is null");
    if (e == NULL) return NULL;
    wai_buffer_t out_bytes = EMPTY_BUFFER, out_answer = EMPTY_BUFFER;
    int st = wai_int_decode_frame(decoder_of(h), frames_of(s), e, n, &out_bytes, &out_answer);
    free(e);
    if (!ok(env, st)) { wai_buffer_free(out_bytes); wai_buffer_free(out_answer); return NULL; }
    return give_pair(env, out_bytes, out_answer);
}

/* The decoder under its operator's model policy (native-sink §14): the policy held, the clock it
 * reads, what it holds, and the run report signed for a period with the device key's seed. A
 * decode the policy refuses is an answer, not an error; a policy, clock or request that does not
 * read or cannot be met throws with libwai's message, which starts with the policy's code. */

JNIEXPORT void JNICALL Java_wai_sink_WaiNative_intPolicy(JNIEnv *env, jclass cls, jlong h, jbyteArray config) {
    (void)cls;
    char *c = text(env, config, "config is null");
    if (c == NULL) return;
    int st = wai_int_decoder_policy(decoder_of(h), c);
    free(c);
    ok(env, st);
}

JNIEXPORT void JNICALL Java_wai_sink_WaiNative_intClock(JNIEnv *env, jclass cls, jlong h, jbyteArray at) {
    (void)cls;
    uintptr_t n = 0;
    char *t = (char *)take(env, at, 1, &n);
    if (at != NULL && t == NULL) return;
    int st = wai_int_decoder_clock(decoder_of(h), t);
    free(t);
    ok(env, st);
}

JNIEXPORT jbyteArray JNICALL Java_wai_sink_WaiNative_intModel(JNIEnv *env, jclass cls, jlong h) {
    (void)cls;
    wai_buffer_t out = EMPTY_BUFFER;
    if (!ok(env, wai_int_decoder_model(decoder_of(h), &out))) return NULL;
    return give(env, out);
}

JNIEXPORT jbyteArray JNICALL Java_wai_sink_WaiNative_intRunReport(JNIEnv *env, jclass cls, jlong h, jbyteArray request, jbyteArray seed) {
    (void)cls;
    char *r = text(env, request, "request is null");
    if (r == NULL) return NULL;
    uintptr_t n;
    uint8_t *s = bytes(env, seed, &n, "seed is null");
    if (s == NULL) { free(r); return NULL; }
    wai_buffer_t out = EMPTY_BUFFER;
    int st = wai_int_run_report(decoder_of(h), r, s, n, &out);
    memset(s, 0, (size_t)n);
    free(r); free(s);
    if (!ok(env, st)) { wai_buffer_free(out); return NULL; }
    return give(env, out);
}

/* A model-governance bundle verified as `wai model verify` verifies it: the outcome's JSON. */
JNIEXPORT jbyteArray JNICALL Java_wai_sink_WaiNative_modelVerify(JNIEnv *env, jclass cls, jbyteArray bundle) {
    (void)cls;
    uintptr_t n;
    uint8_t *b = bytes(env, bundle, &n, "bundle is null");
    if (b == NULL) return NULL;
    wai_buffer_t out = EMPTY_BUFFER;
    int st = wai_model_verify(b, n, &out);
    free(b);
    if (!ok(env, st)) { wai_buffer_free(out); return NULL; }
    return give(env, out);
}