#include <android/log.h>
#include <sys/eventfd.h>
#include <thread>

#include <jni.h>
#include <rfb/rfb.h>

#define LOG_TAG "RFBServerNative"

#define CLASS_EVENT_CALLBACK "com/javispedro/vndroid/RFBServer$EventCallback"

static void rfb_err(const char *format, ...) {
    va_list args;
    va_start(args, format);
    __android_log_vprint(ANDROID_LOG_ERROR, LOG_TAG, format, args);
    va_end(args);
}

static void rfb_log(const char *format, ...) {
    va_list args;
    va_start(args, format);
    __android_log_vprint(ANDROID_LOG_DEBUG, LOG_TAG, format, args);
    va_end(args);
}

struct Data
{
    jclass cls;
    jfieldID fid;
    jfieldID fid_cb;

    rfbScreenInfoPtr screen;

    std::thread listenThread;
    int listenEventFd;

    std::thread eventThread;
    int eventEventFd;

    jclass cb_cls;
    jobject cb_obj;
    jmethodID cb_ev_ptr;
    jmethodID cb_ev_key;

    std::atomic<bool> pendingPtrEvent;
    int ptrButtonMask, ptrX, ptrY;

    std::atomic<bool> pendingKeyEvent;
    rfbKeySym key;
    bool keyState;
};

static inline Data * getData(JNIEnv *env, jobject instance)
{
    jclass cls = env->GetObjectClass(instance);
    cls = reinterpret_cast<jclass>(env->NewGlobalRef(cls));

    jfieldID fid = env->GetFieldID(cls, "nativeData", "J");

    Data *data = reinterpret_cast<Data*>(env->GetLongField(instance, fid));

    assert(data);

    return data;
}

static inline Data * getData(rfbClientPtr client)
{
    Data *data = reinterpret_cast<Data*>(client->screen->screenData);
    assert(data);
    return data;
}

static void listen_thread_main(Data *data, JavaVM *vm)
{
    rfbScreenInfoPtr screen = data->screen;

    while (rfbIsActive(screen)) {
        int client_fd = -1;
        fd_set listen_fds;  /* temp file descriptor list for select() */
        FD_ZERO(&listen_fds);
        if(screen->listenSock >= 0)
            FD_SET(screen->listenSock, &listen_fds);
        if(screen->listen6Sock >= 0)
            FD_SET(screen->listen6Sock, &listen_fds);

        assert(data->listenEventFd >= 0);
        FD_SET(data->listenEventFd, &listen_fds);

        int max_fd = std::max(screen->maxFd, data->listenEventFd) + 1;

        if (select(max_fd, &listen_fds, NULL, NULL, NULL) == -1) {
            rfb_err("listen_thread_main: error in select");
            return;
        }

        /* there is something on the listening sockets, handle new connections */
        struct sockaddr_storage peer;
        socklen_t len = sizeof (peer);
        if (screen->listenSock >= 0 && FD_ISSET(screen->listenSock, &listen_fds))
            client_fd = accept(screen->listenSock, (struct sockaddr*)&peer, &len);
        else if (screen->listen6Sock >= 0 && FD_ISSET(screen->listen6Sock, &listen_fds))
            client_fd = accept(screen->listen6Sock, (struct sockaddr*)&peer, &len);
        else if (data->listenEventFd >= 0 && FD_ISSET(data->listenEventFd, &listen_fds))
            break;

        rfbClientPtr cl = NULL;
        if(client_fd >= 0)
            cl = rfbNewClient(screen,client_fd);
        if (cl && !cl->onHold )
            rfbStartOnHoldClient(cl);
    }
}

static void event_thread_main(Data *data, JavaVM *vm)
{
    rfbScreenInfoPtr screen = data->screen;

    JNIEnv *env = NULL;
    vm->AttachCurrentThread(&env, 0);
    assert(env);

    eventfd_t value;
    while (rfbIsActive(screen)) {
        if (eventfd_read(data->eventEventFd, &value) != 0)
            break;

        if (data->pendingPtrEvent) {
            env->CallVoidMethod(data->cb_obj, data->cb_ev_ptr, (jint)data->ptrButtonMask, data->ptrX, data->ptrY);
            data->pendingPtrEvent = false;
        }

        if (data->pendingKeyEvent) {
            env->CallVoidMethod(data->cb_obj, data->cb_ev_key, (jint)data->key, (jboolean)data->keyState);
            data->pendingKeyEvent = false;
        }
    }

    vm->DetachCurrentThread();
}

static void ptr_event(int buttonMask, int x, int y, rfbClientPtr cl)
{
    Data *data = getData(cl);

    data->ptrButtonMask = buttonMask;
    data->ptrX = x;
    data->ptrY = y;
    data->pendingPtrEvent = true;

    eventfd_write(data->eventEventFd, 1);

    rfbDefaultPtrAddEvent(buttonMask, x, y, cl);
}

static void key_event(rfbBool down, rfbKeySym keySym, rfbClientPtr cl)
{
    Data *data = getData(cl);

    // TODO This misses keys...

    data->key = keySym;
    data->keyState = down;
    data->pendingKeyEvent = true;

    eventfd_write(data->eventEventFd, 1);
}

extern "C"
JNIEXPORT jboolean JNICALL
Java_com_javispedro_vndroid_RFBServer_allocate(JNIEnv *env, jobject instance)
{
    jclass cls = env->GetObjectClass(instance);
    cls = reinterpret_cast<jclass>(env->NewGlobalRef(cls));

    jfieldID fid = env->GetFieldID(cls, "nativeData", "J");

    assert(fid);
    assert(env->GetLongField(instance, fid) == 0);

    rfb_log("starting server");

    rfbErr = rfb_err;
    rfbLog = rfb_log;

    Data *data = new Data();
    assert(data);

    data->cls = cls;
    data->fid = fid;
    data->fid_cb = env->GetFieldID(cls, "callback", "L" CLASS_EVENT_CALLBACK ";");
    assert(data->fid_cb);

    jclass cb_cls = env->FindClass(CLASS_EVENT_CALLBACK);
    assert(cb_cls);
    data->cb_cls = cb_cls = reinterpret_cast<jclass>(env->NewGlobalRef(cb_cls));
    data->cb_obj = 0;
    data->cb_ev_ptr = env->GetMethodID(cb_cls, "onPointerEvent", "(III)V");
    assert(data->cb_ev_ptr);
    data->cb_ev_key = env->GetMethodID(cb_cls, "onKeyEvent", "(IZ)V");
    assert(data->cb_ev_key);

    data->screen = 0;
    data->listenEventFd = -1;
    data->eventEventFd = -1;

    env->SetLongField(instance, fid, reinterpret_cast<jlong>(data));

    return JNI_TRUE;
}

extern "C"
JNIEXPORT void JNICALL
Java_com_javispedro_vndroid_RFBServer_deallocate(JNIEnv *env, jobject instance)
{
    Data *data = getData(env, instance);

    assert(data->screen == 0);

    env->SetLongField(instance, data->fid, 0);
    env->DeleteGlobalRef(data->cb_cls);
    if (data->cb_obj) env->DeleteGlobalRef(data->cb_obj);
    env->DeleteGlobalRef(data->cls);

    data->cb_cls = 0;
    data->cb_obj = 0;
    data->cls = 0;

    delete data;
}


extern "C"
JNIEXPORT jboolean JNICALL
Java_com_javispedro_vndroid_RFBServer_init(JNIEnv *env, jobject instance)
{
    Data *data = getData(env, instance);

    if (data->screen) {
        rfb_err("server already initialized");
        return JNI_FALSE;
    }

    JavaVM *vm = NULL;
    env->GetJavaVM(&vm);

    rfb_log("allocating rfb screen");
    data->screen = rfbGetScreen(NULL, NULL, 0, 0, 0, 0, 0);
    rfbInitServer(data->screen);

    data->screen->screenData = data;
    data->screen->ptrAddEvent = ptr_event;
    data->screen->kbdAddEvent = key_event;

    rfb_log("starting listen thread");
    data->screen->backgroundLoop = TRUE;
    data->listenEventFd = eventfd(0, EFD_CLOEXEC);
    if (data->listenEventFd == -1) {
        rfbLogPerror("eventfd");
        return FALSE;
    }
    data->listenThread = std::thread(listen_thread_main, data, vm);

    rfb_log("starting event thread");
    data->eventThread = std::thread(event_thread_main, data, vm);
    data->eventEventFd = eventfd(0, EFD_CLOEXEC);
    if (data->eventEventFd == -1) {
        rfbLogPerror("eventfd");
        return FALSE;
    }

    return JNI_TRUE;
}

extern "C"
JNIEXPORT void JNICALL
Java_com_javispedro_vndroid_RFBServer_shutdown(JNIEnv *env, jobject instance)
{
    Data *data = getData(env, instance);

    if (!data->screen) {
        return; // Nothing to do
    }

    rfb_log("shutting down rfb screen");
    rfbShutdownServer(data->screen, TRUE);

    rfb_log("waiting for listen thread");
    if (data->listenThread.joinable()) {
        assert(data->listenEventFd != -1);
        eventfd_write(data->listenEventFd, 1);
        data->listenThread.join();
    }
    close(data->listenEventFd);
    data->listenEventFd = -1;

    rfb_log("waiting for event thread");
    if (data->eventThread.joinable()) {
        assert(data->eventEventFd != -1);
        eventfd_write(data->eventEventFd, 1);
        data->eventThread.join();
    }
    close(data->eventEventFd);
    data->eventEventFd = -1;

    rfb_log("freeing up resources");
    rfbScreenCleanup(data->screen);
    data->screen = 0;
}

extern "C"
JNIEXPORT void JNICALL
Java_com_javispedro_vndroid_RFBServer_set_1event_1callback(JNIEnv *env, jobject instance,
                                                           jobject callback) {
    Data *data = getData(env, instance);

    if (data->cb_obj) {
        env->DeleteGlobalRef(data->cb_obj);
    }

    if (callback) {
        data->cb_obj = env->NewGlobalRef(callback);
    } else {
        data->cb_obj = 0;
    }
}

extern "C"
JNIEXPORT jboolean JNICALL
Java_com_javispedro_vndroid_RFBServer_put_1image(JNIEnv *env, jobject instance, jint width,
                                                 jint height, jobject buffer, jint pixel_stride,
                                                 jint row_stride) {
    Data *data = getData(env, instance);
    if (!data->screen) {
        rfb_err("no screen");
        return JNI_FALSE;
    }

    const int bitsPerSample = 8;
    const int samplesPerPixel = 3;
    const int bytesPerPixel = pixel_stride;
    const int bitsPerPixel = 8 * bytesPerPixel;
    const int bytesPerRow = row_stride;

    bool size_change = data->screen->width != width || data->screen->height != height;
    bool format_change = data->screen->bitsPerPixel != bitsPerPixel or data->screen->paddedWidthInBytes != bytesPerRow;

    char * fb = reinterpret_cast<char *>(env->GetDirectBufferAddress(buffer));

    if (size_change || format_change) {
        rfb_log("size or format changed");

        rfbNewFramebuffer(data->screen, fb, width, height, bitsPerSample, samplesPerPixel,
                          bytesPerPixel, bytesPerRow);
    } else {
        data->screen->frameBuffer = fb;
        rfbMarkRectAsModified(data->screen, 0, 0, width, height);
    }

    return JNI_TRUE;
}