#include "xr_overlay.h"

#include <dlfcn.h>
#include <math.h>
#include <stdlib.h>
#include <string.h>

#include <android/log.h>
#include <jni.h>

#include <EGL/egl.h>
#include <GLES3/gl3.h>

#include <SDL.h>

#include "xr_swapchain.h"

#define TAG "QuestShock"
#define LOGE(...) __android_log_print(ANDROID_LOG_ERROR, TAG, __VA_ARGS__)
#define LOGI(...) __android_log_print(ANDROID_LOG_INFO, TAG, __VA_ARGS__)

// Same constant (deliberately not <math.h>'s M_PI, which isn't used
// anywhere else in this codebase and isn't guaranteed by the C standard -
// only POSIX) as xr_session.c's own play-space rebuild math.
#define QS_DEG_TO_RAD (3.14159265358979323846f / 180.0f)

struct XrOverlay {
    XrInstance instance; // for error logging only
    int width, height;
    float distance_m, half_width_m, half_height_m;
    float center_x_m, center_y_m;
    // The *authored* distance_m (config->distance_m at creation time),
    // never touched again afterward by anything, including
    // xr_overlay_set_pose() - kept separate from the live, mutable
    // distance_m above because that one stops meaning "how far in front of
    // the origin" once a rotated set_pose() call (the menu recenter)
    // repurposes it as a raw base-space Z coordinate; reusing *that* as a
    // "how far in front of me" magnitude on a later recenter would drift
    // further with every reopen instead of staying at a stable distance -
    // see xr_overlay_get_default_distance_m().
    float default_distance_m;
    // Rotation about the vertical axis, applied on top of center_x_m/
    // center_y_m/distance_m (see xr_overlay_set_pose()) - 0 for every
    // overlay except the menu launcher immediately after it's freshly
    // reopened (xr_input.c re-centers it in front of wherever the player
    // is currently facing, which needs this to keep the quad actually
    // facing them instead of whatever direction it started out at).
    float yaw_deg;
    bool visible;

    XrSwapchainState swapchain;

    jclass java_class;
    jmethodID native_init_method;
    jmethodID take_pixels_method;
    jmethodID dispatch_touch_method;
    jmethodID update_cursor_method;

    // This overlay's content lives in this texture, uploaded by
    // xr_overlay_upload_source_texture_if_dirty() below. It's created
    // through gl4es (glGenTextures/glBindTexture/glTexImage2D), like any
    // other texture the engine itself creates, so ordinary
    // glTexSubImage2D uploads against it work the normal way.
    GLuint source_texture;
    // A framebuffer with source_texture as its only color attachment, used
    // as the read source for the blit in
    // xr_overlay_render_and_build_layer() below. Created via the real
    // (non-gl4es) GLES entry points below: that blit bypasses gl4es
    // entirely (see xr_overlay_render_and_build_layer() for why), so its
    // source framebuffer has to be a real GL object rather than one gl4es
    // tracks.
    GLuint source_fbo;

    uint32_t pixel_generation;
    uint32_t source_uploaded_generation;
    uint8_t *pixel_cache; // width*height*4 bytes
};

// This file's other GL calls otherwise resolve to gl4es (the only GL
// symbol provider linked into this binary - see android/engine-patches/
// 02-android-opengl-es.patch), fine for anything shared with the engine's
// own gl4es-routed rendering. But the source FBO's blit (see
// xr_overlay_render_and_build_layer()) needs the real driver directly:
// gl4es's fixed-pipeline emulation unconditionally substitutes its own
// shader onto any gl4es-routed draw call, which would silently replace an
// overlay's content with whatever the game itself last rendered. A
// framebuffer blit has no shader stage at all, so going through the real
// driver for it sidesteps the problem entirely - resolved once here and
// shared by every XrOverlay instance, since they're process-wide function
// pointers regardless of how many overlays exist.
typedef void (*PFNQSGETINTEGERV)(GLenum, GLint *);
typedef void (*PFNQSGENFRAMEBUFFERS)(GLsizei, GLuint *);
typedef void (*PFNQSDELETEFRAMEBUFFERS)(GLsizei, const GLuint *);
typedef void (*PFNQSBINDFRAMEBUFFER)(GLenum, GLuint);
typedef void (*PFNQSFRAMEBUFFERTEXTURE2D)(GLenum, GLenum, GLenum, GLuint, GLint);
typedef void (*PFNQSBLITFRAMEBUFFER)(GLint, GLint, GLint, GLint, GLint, GLint, GLint, GLint,
                                      GLbitfield, GLenum);

static PFNQSGETINTEGERV real_glGetIntegerv;
static PFNQSGENFRAMEBUFFERS real_glGenFramebuffers;
static PFNQSDELETEFRAMEBUFFERS real_glDeleteFramebuffers;
static PFNQSBINDFRAMEBUFFER real_glBindFramebuffer;
static PFNQSFRAMEBUFFERTEXTURE2D real_glFramebufferTexture2D;
// glBlitFramebuffer is a GLES 3.0 addition; resolved via eglGetProcAddress
// rather than dlsym, since Android's driver dispatch doesn't guarantee ES3+
// symbols are dlsym-able by name from libGLESv2.so, unlike the GLES2-core
// functions above.
static PFNQSBLITFRAMEBUFFER real_glBlitFramebuffer;
static bool g_real_gles_loaded = false;

static bool xr_overlay_load_real_gles(void) {
    if (g_real_gles_loaded)
        return true;
    void *lib = dlopen("libGLESv2.so", RTLD_NOW | RTLD_LOCAL);
    if (lib == NULL) {
        LOGE("XR: overlay dlopen(libGLESv2.so) failed: %s", dlerror());
        return false;
    }
    real_glGetIntegerv = (PFNQSGETINTEGERV)dlsym(lib, "glGetIntegerv");
    real_glGenFramebuffers = (PFNQSGENFRAMEBUFFERS)dlsym(lib, "glGenFramebuffers");
    real_glDeleteFramebuffers = (PFNQSDELETEFRAMEBUFFERS)dlsym(lib, "glDeleteFramebuffers");
    real_glBindFramebuffer = (PFNQSBINDFRAMEBUFFER)dlsym(lib, "glBindFramebuffer");
    real_glFramebufferTexture2D = (PFNQSFRAMEBUFFERTEXTURE2D)dlsym(lib, "glFramebufferTexture2D");
    real_glBlitFramebuffer = (PFNQSBLITFRAMEBUFFER)eglGetProcAddress("glBlitFramebuffer");
    g_real_gles_loaded = real_glGetIntegerv && real_glGenFramebuffers &&
                          real_glDeleteFramebuffers && real_glBindFramebuffer &&
                          real_glFramebufferTexture2D && real_glBlitFramebuffer;
    return g_real_gles_loaded;
}

// FindClass() from this thread (SDL's native thread, attached to the JVM
// via AttachCurrentThread rather than spawned from Java) resolves against
// the bootstrap classloader, which only knows framework classes - it can't
// see app classes like de.ladkau.questshock.KeyboardOverlay at all.
// Routing through the activity's own classloader is the standard,
// documented workaround.
static jclass xr_overlay_find_class(JNIEnv *env, jobject activity, const char *name) {
    jclass activityClass = (*env)->GetObjectClass(env, activity);
    jmethodID getClassLoader =
        (*env)->GetMethodID(env, activityClass, "getClassLoader", "()Ljava/lang/ClassLoader;");
    jobject classLoader = (*env)->CallObjectMethod(env, activity, getClassLoader);
    jclass classLoaderClass = (*env)->FindClass(env, "java/lang/ClassLoader");
    jmethodID loadClass = (*env)->GetMethodID(env, classLoaderClass, "loadClass",
                                               "(Ljava/lang/String;)Ljava/lang/Class;");
    jstring className = (*env)->NewStringUTF(env, name);
    jclass result = (jclass)(*env)->CallObjectMethod(env, classLoader, loadClass, className);
    (*env)->DeleteLocalRef(env, className);
    return result;
}

static bool xr_overlay_init_jni(XrOverlay *overlay, const char *java_class_name) {
    JNIEnv *env = (JNIEnv *)SDL_AndroidGetJNIEnv();
    jobject activity = (jobject)SDL_AndroidGetActivity();
    if (env == NULL || activity == NULL) {
        LOGE("XR: overlay - no JNIEnv/Activity from SDL");
        return false;
    }

    jclass localClass = xr_overlay_find_class(env, activity, java_class_name);
    if (localClass == NULL) {
        LOGE("XR: overlay - could not find class %s", java_class_name);
        return false;
    }
    overlay->java_class = (jclass)(*env)->NewGlobalRef(env, localClass);

    overlay->native_init_method = (*env)->GetStaticMethodID(env, overlay->java_class, "nativeInit",
                                                              "(Landroid/app/Activity;)V");
    overlay->take_pixels_method =
        (*env)->GetStaticMethodID(env, overlay->java_class, "nativeTakePixelsIfDirty", "()[B");
    overlay->dispatch_touch_method =
        (*env)->GetStaticMethodID(env, overlay->java_class, "nativeDispatchTouch", "(FFZ)V");
    overlay->update_cursor_method =
        (*env)->GetStaticMethodID(env, overlay->java_class, "nativeUpdateCursor", "(FFZ)V");
    if (!overlay->native_init_method || !overlay->take_pixels_method ||
        !overlay->dispatch_touch_method || !overlay->update_cursor_method) {
        LOGE("XR: overlay - could not resolve %s JNI methods", java_class_name);
        return false;
    }

    (*env)->CallStaticVoidMethod(env, overlay->java_class, overlay->native_init_method, activity);
    // CallStaticVoidMethod doesn't surface Java exceptions on its own - if
    // the constructor throws (it runs its View measure/layout/draw calls
    // on this native render thread rather than the Android UI thread, a
    // plausible crash vector), the exception would otherwise be left
    // silently pending and corrupt whatever JNI call runs next.
    if ((*env)->ExceptionCheck(env)) {
        LOGE("XR: overlay %s.nativeInit() threw a pending Java exception:", java_class_name);
        (*env)->ExceptionDescribe(env);
        (*env)->ExceptionClear(env);
        return false;
    }
    return true;
}

static bool xr_overlay_create_source_texture(XrOverlay *overlay) {
    glGenTextures(1, &overlay->source_texture);
    glBindTexture(GL_TEXTURE_2D, overlay->source_texture);
    glTexParameteri(GL_TEXTURE_2D, GL_TEXTURE_MIN_FILTER, GL_LINEAR);
    glTexParameteri(GL_TEXTURE_2D, GL_TEXTURE_MAG_FILTER, GL_LINEAR);
    glTexParameteri(GL_TEXTURE_2D, GL_TEXTURE_WRAP_S, GL_CLAMP_TO_EDGE);
    glTexParameteri(GL_TEXTURE_2D, GL_TEXTURE_WRAP_T, GL_CLAMP_TO_EDGE);
    glTexImage2D(GL_TEXTURE_2D, 0, GL_RGBA, overlay->width, overlay->height, 0, GL_RGBA,
                 GL_UNSIGNED_BYTE, NULL);

    GLint prevFbo = 0;
    real_glGetIntegerv(GL_FRAMEBUFFER_BINDING, &prevFbo);
    real_glGenFramebuffers(1, &overlay->source_fbo);
    real_glBindFramebuffer(GL_FRAMEBUFFER, overlay->source_fbo);
    real_glFramebufferTexture2D(GL_FRAMEBUFFER, GL_COLOR_ATTACHMENT0, GL_TEXTURE_2D,
                                 overlay->source_texture, 0);
    real_glBindFramebuffer(GL_FRAMEBUFFER, (GLuint)prevFbo);

    return overlay->source_texture != 0;
}

XrOverlay *xr_overlay_create(XrInstance instance, XrSession session, int64_t swapchain_format,
                              const XrOverlayConfig *config) {
    if (!xr_overlay_load_real_gles()) {
        LOGE("XR: overlay couldn't resolve real GLES functions via dlsym/eglGetProcAddress");
        return NULL;
    }

    XrOverlay *overlay = (XrOverlay *)calloc(1, sizeof(XrOverlay));
    if (overlay == NULL)
        return NULL;

    overlay->instance = instance;
    overlay->width = config->width;
    overlay->height = config->height;
    overlay->distance_m = config->distance_m;
    overlay->default_distance_m = config->distance_m;
    overlay->half_width_m = config->width_m * 0.5f;
    overlay->half_height_m = config->width_m * 0.5f * (float)config->height / (float)config->width;
    overlay->center_x_m = config->default_center_x_m;
    overlay->center_y_m = config->default_center_y_m;

    if (!xr_swapchain_create(instance, session, swapchain_format, config->width, config->height,
                              &overlay->swapchain)) {
        LOGE("XR: overlay %s swapchain setup failed", config->java_class_name);
        free(overlay);
        return NULL;
    }

    if (!xr_overlay_create_source_texture(overlay)) {
        xr_swapchain_destroy(&overlay->swapchain);
        free(overlay);
        return NULL;
    }

    if (!xr_overlay_init_jni(overlay, config->java_class_name)) {
        xr_swapchain_destroy(&overlay->swapchain);
        free(overlay);
        return NULL;
    }

    overlay->pixel_cache = (uint8_t *)malloc((size_t)config->width * config->height * 4);
    if (overlay->pixel_cache == NULL) {
        xr_swapchain_destroy(&overlay->swapchain);
        free(overlay);
        return NULL;
    }
    overlay->source_uploaded_generation = (uint32_t)-1;

    LOGI("XR: overlay %s ready (%dx%d)", config->java_class_name, config->width, config->height);
    return overlay;
}

void xr_overlay_destroy(XrOverlay *overlay) {
    if (overlay == NULL)
        return;

    if (overlay->source_texture != 0)
        glDeleteTextures(1, &overlay->source_texture);
    if (overlay->source_fbo != 0 && real_glDeleteFramebuffers != NULL)
        real_glDeleteFramebuffers(1, &overlay->source_fbo);

    if (overlay->java_class != NULL) {
        JNIEnv *env = (JNIEnv *)SDL_AndroidGetJNIEnv();
        if (env != NULL)
            (*env)->DeleteGlobalRef(env, overlay->java_class);
    }

    xr_swapchain_destroy(&overlay->swapchain);
    free(overlay->pixel_cache);
    free(overlay);
}

// Every function below tolerates overlay == NULL (a no-op, or the
// obvious "not visible"/zeroed-extent default) - xr_overlay_create() can
// fail (e.g. its backing Java class couldn't be resolved), and unlike the
// single static-global menu quad this replaced, an overlay instance is
// otherwise a heap pointer callers (xr_session.c, questshock_native.c) can
// hold onto and use across frames without re-checking for failure every
// time.
void xr_overlay_toggle_visible(XrOverlay *overlay) {
    if (overlay != NULL)
        overlay->visible = !overlay->visible;
}

void xr_overlay_set_visible(XrOverlay *overlay, bool visible) {
    if (overlay != NULL)
        overlay->visible = visible;
}

bool xr_overlay_is_visible(const XrOverlay *overlay) { return overlay != NULL && overlay->visible; }

void xr_overlay_get_quad_extent(const XrOverlay *overlay, float *center_x_m, float *center_y_m,
                                 float *distance_m, float *half_width_m, float *half_height_m,
                                 float *yaw_deg) {
    if (overlay == NULL) {
        *center_x_m = *center_y_m = *distance_m = *half_width_m = *half_height_m = *yaw_deg = 0.0f;
        return;
    }
    *center_x_m = overlay->center_x_m;
    *center_y_m = overlay->center_y_m;
    *distance_m = overlay->distance_m;
    *half_width_m = overlay->half_width_m;
    *half_height_m = overlay->half_height_m;
    *yaw_deg = overlay->yaw_deg;
}

void xr_overlay_set_position(XrOverlay *overlay, float center_x_m, float center_y_m) {
    if (overlay == NULL)
        return;
    overlay->center_x_m = center_x_m;
    overlay->center_y_m = center_y_m;
}

void xr_overlay_set_pose(XrOverlay *overlay, float center_x_m, float center_y_m, float distance_m,
                          float yaw_deg) {
    if (overlay == NULL)
        return;
    overlay->center_x_m = center_x_m;
    overlay->center_y_m = center_y_m;
    overlay->distance_m = distance_m;
    overlay->yaw_deg = yaw_deg;
}

float xr_overlay_get_default_distance_m(const XrOverlay *overlay) {
    return overlay == NULL ? 0.0f : overlay->default_distance_m;
}

void xr_overlay_touch(XrOverlay *overlay, float u, float v, bool down) {
    if (overlay == NULL)
        return;
    JNIEnv *env = (JNIEnv *)SDL_AndroidGetJNIEnv();
    if (env == NULL)
        return;
    // Passed via CallStaticVoidMethodA/jvalue rather than the variadic
    // Call*Method form - deliberately sidesteps relying on JNI
    // implementations correctly un-doing C's float-to-double default
    // argument promotion for varargs float parameters (a well-known JNI
    // footgun; Android's ART handles it correctly, but there's no reason
    // to depend on that when the jvalue form is unambiguous either way).
    jvalue args[3];
    args[0].f = u;
    args[1].f = v;
    args[2].z = (jboolean)down;
    (*env)->CallStaticVoidMethodA(env, overlay->java_class, overlay->dispatch_touch_method, args);
}

void xr_overlay_update_cursor(XrOverlay *overlay, float u, float v, bool visible) {
    if (overlay == NULL)
        return;
    JNIEnv *env = (JNIEnv *)SDL_AndroidGetJNIEnv();
    if (env == NULL)
        return;
    // Same jvalue-form rationale as xr_overlay_touch() above.
    jvalue args[3];
    args[0].f = u;
    args[1].f = v;
    args[2].z = (jboolean)visible;
    (*env)->CallStaticVoidMethodA(env, overlay->java_class, overlay->update_cursor_method, args);
}

// Pulls the Java view's latest pixels (if it redrew since the last check)
// into pixel_cache and bumps pixel_generation - called once per rendered
// frame, before deciding whether source_texture needs a fresh upload.
static void xr_overlay_refresh_pixel_cache(XrOverlay *overlay) {
    JNIEnv *env = (JNIEnv *)SDL_AndroidGetJNIEnv();
    if (env == NULL)
        return;
    jbyteArray pixels = (jbyteArray)(*env)->CallStaticObjectMethod(
        env, overlay->java_class, overlay->take_pixels_method);
    // CallStaticObjectMethod also doesn't surface Java exceptions on its
    // own - if nativeTakePixelsIfDirty() itself throws, it would otherwise
    // silently look identical to "nothing changed yet" (a null return).
    if ((*env)->ExceptionCheck(env)) {
        LOGE("XR: overlay nativeTakePixelsIfDirty() threw a pending Java exception:");
        (*env)->ExceptionDescribe(env);
        (*env)->ExceptionClear(env);
        return;
    }
    if (pixels == NULL)
        return;
    (*env)->GetByteArrayRegion(env, pixels, 0, overlay->width * overlay->height * 4,
                                (jbyte *)overlay->pixel_cache);
    (*env)->DeleteLocalRef(env, pixels);
    overlay->pixel_generation++;
}

// Uploads pixel_cache into source_texture via plain (gl4es-routed)
// glBindTexture/glTexSubImage2D, matching how the texture was created.
// Skipped when nothing changed since the last upload. Texture-unit-0 state
// is saved and restored around the upload, since this call is gl4es-routed
// and the engine's own next-frame rendering assumes nothing touched its
// texture bindings since it last drew.
static void xr_overlay_upload_source_texture_if_dirty(XrOverlay *overlay) {
    if (overlay->source_uploaded_generation == overlay->pixel_generation)
        return;

    GLint prevActiveTexture = GL_TEXTURE0;
    glGetIntegerv(GL_ACTIVE_TEXTURE, &prevActiveTexture);
    glActiveTexture(GL_TEXTURE0);
    GLint prevTexture2D = 0;
    glGetIntegerv(GL_TEXTURE_BINDING_2D, &prevTexture2D);

    glBindTexture(GL_TEXTURE_2D, overlay->source_texture);
    glTexSubImage2D(GL_TEXTURE_2D, 0, 0, 0, overlay->width, overlay->height, GL_RGBA,
                     GL_UNSIGNED_BYTE, overlay->pixel_cache);
    overlay->source_uploaded_generation = overlay->pixel_generation;

    glBindTexture(GL_TEXTURE_2D, (GLuint)prevTexture2D);
    if (prevActiveTexture != GL_TEXTURE0)
        glActiveTexture((GLenum)prevActiveTexture);
}

bool xr_overlay_render_and_build_layer(XrOverlay *overlay, XrCompositionLayerQuad *out_quad) {
    if (overlay == NULL || !overlay->visible)
        return false;
    if (!xr_swapchain_acquire(overlay->instance, &overlay->swapchain))
        return false;

    xr_overlay_refresh_pixel_cache(overlay);
    xr_overlay_upload_source_texture_if_dirty(overlay);

    // Copies source_texture (via source_fbo) directly into the swapchain
    // image xr_swapchain_acquire() just bound - sized to exactly
    // width x height, so the destination is always the whole image - using
    // the real GLES3 hardware blit (glBlitFramebuffer) rather than a
    // shader-based full-screen-quad draw. A blit has no vertex/fragment
    // shading stage, no shader program, and no texture units involved at
    // all, so gl4es's fixed-pipeline-emulation layer - which
    // unconditionally substitutes a customized shader (reproducing the
    // game's own last-bound texture/fixed-function state) onto any
    // gl4es-routed draw call - has nothing to intercept here. It also
    // never touches gl4es's own tracked program/vertex-array/
    // texture-binding shadow state, so the only piece of state that needs
    // save/restore is the READ framebuffer binding; binding only
    // GL_READ_FRAMEBUFFER, rather than the combined GL_FRAMEBUFFER target,
    // leaves the DRAW side (the swapchain image itself) untouched
    // throughout.
    //
    // Y is flipped between source and destination: source_texture's texel
    // row v=0 holds the Java view's Bitmap row 0 (the top of the rendered
    // content, standard top-row-first raster order), and reading that same
    // texture via an attached FBO, framebuffer y=0 accesses that identical
    // row. The destination swapchain framebuffer follows the opposite
    // convention - its own y=0 is its bottom edge - so swapping dstY0/dstY1
    // keeps the content right-side up (glBlitFramebuffer supports inverted
    // src/dst rects for exactly this).
    GLint prevReadFbo = 0;
    real_glGetIntegerv(GL_READ_FRAMEBUFFER_BINDING, &prevReadFbo);

    real_glBindFramebuffer(GL_READ_FRAMEBUFFER, overlay->source_fbo);
    real_glBlitFramebuffer(0, 0, overlay->width, overlay->height, 0, overlay->height,
                            overlay->width, 0, GL_COLOR_BUFFER_BIT, GL_LINEAR);
    real_glBindFramebuffer(GL_READ_FRAMEBUFFER, (GLuint)prevReadFbo);

    // Without this, nothing here waits for the blit to actually finish on
    // the GPU before xr_swapchain_release() hands the image back to the
    // compositor - glBlitFramebuffer only queues the copy. A cross-process
    // consumer reading the image before the driver has finished writing it
    // is a textbook source of the diagonal-banded/torn corruption seen on
    // device (see xr_swapchain_acquire()'s counterpart in OpenGL.cc, which
    // had the same gap).
    glFinish();

    xr_swapchain_release(overlay->instance, &overlay->swapchain);

    memset(out_quad, 0, sizeof(*out_quad));
    out_quad->type = XR_TYPE_COMPOSITION_LAYER_QUAD;
    out_quad->subImage.swapchain = overlay->swapchain.swapchain;
    out_quad->subImage.imageRect.extent.width = overlay->swapchain.width;
    out_quad->subImage.imageRect.extent.height = overlay->swapchain.height;
    // Rotation about the vertical axis only (see the struct field's own
    // comment) - same R_y(-yaw_deg) quaternion construction as
    // xr_session.c's own play-space rebuild, for the same reason: this
    // quad's pose is expressed as "yaw_deg turned from its own local -Z",
    // so a point/ray in that rotated frame maps back into the (unrotated)
    // base space by rotating by -yaw_deg.
    float half = -(overlay->yaw_deg * QS_DEG_TO_RAD) * 0.5f;
    out_quad->pose.orientation.x = 0.0f;
    out_quad->pose.orientation.y = sinf(half);
    out_quad->pose.orientation.z = 0.0f;
    out_quad->pose.orientation.w = cosf(half);
    out_quad->pose.position.x = overlay->center_x_m;
    out_quad->pose.position.y = overlay->center_y_m;
    out_quad->pose.position.z = -overlay->distance_m;
    out_quad->size.width = overlay->half_width_m * 2.0f;
    out_quad->size.height = overlay->half_height_m * 2.0f;
    return true;
}