#include "xr_session.h"

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

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

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

#define XR_USE_PLATFORM_ANDROID 1
#define XR_USE_GRAPHICS_API_OPENGL_ES 1
#include <openxr/openxr.h>
#include <openxr/openxr_platform.h>

#include <SDL.h>

#include "xr_input.h"
#include "xr_overlay.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__)

// Fixed pose/size for the single game quad - 2m in front of the local space
// origin, sized to preserve the game's own aspect ratio at a comfortable
// visual angle. No per-eye view/projection work is needed for a quad layer
// at all - the compositor handles reprojecting it per eye itself, which is
// exactly why this milestone doesn't need xrLocateViews or any stereo
// rendering (see the plan's step A notes).
#define QUAD_DISTANCE_METERS 2.0f
#define QUAD_WIDTH_METERS 1.6f

// The menu launcher (just a "Keyboard" button - see MenuOverlay.java) is
// small and sits above where the keyboard defaults to; the keyboard keeps
// its previous size/position. They're independent XrOverlay instances (see
// xr_overlay.h) so both can be shown/hidden/dragged separately - deliberately
// positioned apart by default so they don't start out overlapping.
#define MENU_JAVA_CLASS "de.ladkau.questshock.MenuOverlay"
#define MENU_WIDTH 512
#define MENU_HEIGHT 256
#define MENU_DISTANCE_METERS 1.3f
#define MENU_WIDTH_METERS 0.35f
#define MENU_DEFAULT_CENTER_X_METERS 0.0f
#define MENU_DEFAULT_CENTER_Y_METERS 0.15f

#define KEYBOARD_JAVA_CLASS "de.ladkau.questshock.KeyboardOverlay"
#define KEYBOARD_WIDTH 1024
#define KEYBOARD_HEIGHT 768
#define KEYBOARD_DISTANCE_METERS 1.5f
#define KEYBOARD_WIDTH_METERS 0.8f
#define KEYBOARD_DEFAULT_CENTER_X_METERS 0.0f
#define KEYBOARD_DEFAULT_CENTER_Y_METERS -0.45f

static XrInstance g_instance = XR_NULL_HANDLE;
static XrSystemId g_system_id = XR_NULL_SYSTEM_ID;
static XrSession g_session = XR_NULL_HANDLE;
static XrSpace g_local_space = XR_NULL_HANDLE;
static XrSessionState g_session_state = XR_SESSION_STATE_UNKNOWN;
// Distinct from g_session_state - true from a successful xrBeginSession
// until xrEndSession/loss/exit. Many runtimes only advance the *state*
// (e.g. READY -> SYNCHRONIZED -> VISIBLE -> FOCUSED) in response to the
// app actually calling xrWaitFrame/xrBeginFrame/xrEndFrame continuously
// after xrBeginSession - gating frame submission on already having
// reached SYNCHRONIZED+ (as this used to) is a chicken-and-egg deadlock,
// since the runtime never progresses past READY without the frame loop
// running in the first place. Frame calls just need the session to have
// begun, per the standard OpenXR sample pattern - not any particular
// later state.
static bool g_session_running = false;

static XrSwapchainState g_game_swapchain;
static XrOverlay *g_menu_overlay = NULL;
static XrOverlay *g_keyboard_overlay = NULL;

static XrTime g_predicted_display_time = 0;
static bool g_frame_should_render = false;
static bool g_have_acquired_game_image = false;

static bool xr_check(XrResult result, const char *what) {
    if (XR_SUCCEEDED(result))
        return true;
    char resultString[XR_MAX_RESULT_STRING_SIZE] = {0};
    if (g_instance != XR_NULL_HANDLE)
        xrResultToString(g_instance, result, resultString);
    LOGE("XR: %s failed: %s (%d)", what, resultString[0] ? resultString : "?", (int)result);
    return false;
}

static PFN_xrVoidFunction xr_get_proc(const char *name) {
    PFN_xrVoidFunction fn = NULL;
    // xrGetInstanceProcAddr itself works with a NULL instance for the
    // handful of functions (like xrInitializeLoaderKHR) that must be called
    // before an instance exists.
    if (!xr_check(xrGetInstanceProcAddr(g_instance, name, &fn), name))
        return NULL;
    return fn;
}

// The Android OpenXR loader needs the JavaVM/context before anything else -
// separate from XR_KHR_android_create_instance below, which only affects
// xrCreateInstance itself. Skipping this is a common cause of
// xrCreateInstance silently failing to find a runtime on Android.
static bool xr_initialize_android_loader(JavaVM *vm, jobject activity) {
    PFN_xrInitializeLoaderKHR initializeLoader =
        (PFN_xrInitializeLoaderKHR)xr_get_proc("xrInitializeLoaderKHR");
    if (initializeLoader == NULL) {
        LOGE("XR: xrInitializeLoaderKHR not available - no Android OpenXR loader present?");
        return false;
    }

    XrLoaderInitInfoAndroidKHR loaderInitInfo = {XR_TYPE_LOADER_INIT_INFO_ANDROID_KHR};
    loaderInitInfo.applicationVM = vm;
    loaderInitInfo.applicationContext = activity;
    return xr_check(initializeLoader((const XrLoaderInitInfoBaseHeaderKHR *)&loaderInitInfo),
                     "xrInitializeLoaderKHR");
}

// eglGetCurrentDisplay()/eglGetCurrentContext() hand us the display/context,
// but not the EGLConfig they were created with - eglQueryContext's
// EGL_CONFIG_ID plus eglChooseConfig is the standard way to recover it (the
// same trick Khronos' own hello_xr sample uses).
static EGLConfig xr_get_current_egl_config(EGLDisplay display, EGLContext context) {
    EGLint configId = 0;
    eglQueryContext(display, context, EGL_CONFIG_ID, &configId);

    EGLint attribs[] = {EGL_CONFIG_ID, configId, EGL_NONE};
    EGLConfig config = NULL;
    EGLint numConfigs = 0;
    eglChooseConfig(display, attribs, &config, 1, &numConfigs);
    return config;
}

static bool xr_create_instance_and_session(int game_width, int game_height) {
    JNIEnv *env = (JNIEnv *)SDL_AndroidGetJNIEnv();
    jobject activity = (jobject)SDL_AndroidGetActivity();
    if (env == NULL || activity == NULL) {
        LOGE("XR: no JNIEnv/Activity from SDL - can't init OpenXR");
        return false;
    }
    JavaVM *vm = NULL;
    (*env)->GetJavaVM(env, &vm);

    if (!xr_initialize_android_loader(vm, activity))
        return false;

    const char *extensions[] = {
        XR_KHR_ANDROID_CREATE_INSTANCE_EXTENSION_NAME,
        XR_KHR_OPENGL_ES_ENABLE_EXTENSION_NAME,
    };

    XrInstanceCreateInfoAndroidKHR androidInfo = {XR_TYPE_INSTANCE_CREATE_INFO_ANDROID_KHR};
    androidInfo.applicationVM = vm;
    androidInfo.applicationActivity = activity;

    XrInstanceCreateInfo createInfo = {XR_TYPE_INSTANCE_CREATE_INFO};
    createInfo.next = &androidInfo;
    createInfo.enabledExtensionCount = sizeof(extensions) / sizeof(extensions[0]);
    createInfo.enabledExtensionNames = extensions;
    strncpy(createInfo.applicationInfo.applicationName, "QuestShock",
            XR_MAX_APPLICATION_NAME_SIZE - 1);
    strncpy(createInfo.applicationInfo.engineName, "Shockolate", XR_MAX_ENGINE_NAME_SIZE - 1);
    createInfo.applicationInfo.applicationVersion = 1;
    createInfo.applicationInfo.engineVersion = 1;
    createInfo.applicationInfo.apiVersion = XR_CURRENT_API_VERSION;

    if (!xr_check(xrCreateInstance(&createInfo, &g_instance), "xrCreateInstance"))
        return false;

    XrSystemGetInfo systemInfo = {XR_TYPE_SYSTEM_GET_INFO};
    systemInfo.formFactor = XR_FORM_FACTOR_HEAD_MOUNTED_DISPLAY;
    if (!xr_check(xrGetSystem(g_instance, &systemInfo, &g_system_id), "xrGetSystem"))
        return false;

    // Required by spec before creating an OpenGL ES-backed session, even
    // though we don't gate on the returned min/max API version here.
    PFN_xrGetOpenGLESGraphicsRequirementsKHR getGLESRequirements =
        (PFN_xrGetOpenGLESGraphicsRequirementsKHR)xr_get_proc(
            "xrGetOpenGLESGraphicsRequirementsKHR");
    if (getGLESRequirements == NULL)
        return false;
    XrGraphicsRequirementsOpenGLESKHR glesRequirements = {
        XR_TYPE_GRAPHICS_REQUIREMENTS_OPENGL_ES_KHR};
    if (!xr_check(getGLESRequirements(g_instance, g_system_id, &glesRequirements),
                  "xrGetOpenGLESGraphicsRequirementsKHR"))
        return false;

    EGLDisplay display = eglGetCurrentDisplay();
    EGLContext context = eglGetCurrentContext();
    if (display == EGL_NO_DISPLAY || context == EGL_NO_CONTEXT) {
        LOGE("XR: no current EGL display/context - call xr_init() after init_opengl()");
        return false;
    }

    XrGraphicsBindingOpenGLESAndroidKHR graphicsBinding = {
        XR_TYPE_GRAPHICS_BINDING_OPENGL_ES_ANDROID_KHR};
    graphicsBinding.display = display;
    graphicsBinding.config = xr_get_current_egl_config(display, context);
    graphicsBinding.context = context;

    XrSessionCreateInfo sessionInfo = {XR_TYPE_SESSION_CREATE_INFO};
    sessionInfo.next = &graphicsBinding;
    sessionInfo.systemId = g_system_id;
    if (!xr_check(xrCreateSession(g_instance, &sessionInfo, &g_session), "xrCreateSession"))
        return false;

    XrReferenceSpaceCreateInfo spaceInfo = {XR_TYPE_REFERENCE_SPACE_CREATE_INFO};
    spaceInfo.referenceSpaceType = XR_REFERENCE_SPACE_TYPE_LOCAL;
    spaceInfo.poseInReferenceSpace.orientation.w = 1.0f;
    if (!xr_check(xrCreateReferenceSpace(g_session, &spaceInfo, &g_local_space),
                  "xrCreateReferenceSpace"))
        return false;

    // Non-fatal if it fails (e.g. no controllers bound yet) - rendering
    // keeps working either way, just without the laser pointer.
    xr_input_init(g_instance, g_session);

    // Prefer a plain linear 8-bit format over GL_SRGB8_ALPHA8: the source
    // SDL surface pixels are already sRGB-encoded (as ordinary 8-bit image
    // data conventionally is) and get uploaded/sampled as plain linear
    // GL_RGBA with no decode step anywhere in this path, matching the
    // desktop SDL_RenderCopy path this replaces - an sRGB swapchain format
    // would auto-gamma-encode on write and double-encode already-encoded
    // data. Every swapchain below (game quad, menu/keyboard overlays)
    // shares this one choice - the compositor's supported format set
    // doesn't depend on swapchain size.
    uint32_t formatCount = 0;
    xrEnumerateSwapchainFormats(g_session, 0, &formatCount, NULL);
    int64_t *formats = (int64_t *)malloc(sizeof(int64_t) * formatCount);
    xrEnumerateSwapchainFormats(g_session, formatCount, &formatCount, formats);
    int64_t chosenFormat = formats[0];
    for (uint32_t i = 0; i < formatCount; i++) {
        LOGI("XR: swapchain format[%u] = 0x%llx", i, (unsigned long long)formats[i]);
        if (formats[i] == GL_RGBA8) {
            chosenFormat = GL_RGBA8;
            break;
        }
    }
    free(formats);

    if (!xr_swapchain_create(g_instance, g_session, chosenFormat, game_width, game_height,
                              &g_game_swapchain)) {
        LOGE("XR: game swapchain setup failed");
        return false;
    }

    // Both non-fatal - the game keeps rendering with no such panel if
    // either fails (e.g. the Java class couldn't be resolved via JNI).
    XrOverlayConfig menuConfig = {
        .java_class_name = MENU_JAVA_CLASS,
        .width = MENU_WIDTH,
        .height = MENU_HEIGHT,
        .distance_m = MENU_DISTANCE_METERS,
        .width_m = MENU_WIDTH_METERS,
        .default_center_x_m = MENU_DEFAULT_CENTER_X_METERS,
        .default_center_y_m = MENU_DEFAULT_CENTER_Y_METERS,
    };
    g_menu_overlay = xr_overlay_create(g_instance, g_session, chosenFormat, &menuConfig);

    XrOverlayConfig keyboardConfig = {
        .java_class_name = KEYBOARD_JAVA_CLASS,
        .width = KEYBOARD_WIDTH,
        .height = KEYBOARD_HEIGHT,
        .distance_m = KEYBOARD_DISTANCE_METERS,
        .width_m = KEYBOARD_WIDTH_METERS,
        .default_center_x_m = KEYBOARD_DEFAULT_CENTER_X_METERS,
        .default_center_y_m = KEYBOARD_DEFAULT_CENTER_Y_METERS,
    };
    g_keyboard_overlay = xr_overlay_create(g_instance, g_session, chosenFormat, &keyboardConfig);

    LOGI("XR: instance/session ready (game %dx%d, menu %dx%d, keyboard %dx%d)", game_width,
         game_height, MENU_WIDTH, MENU_HEIGHT, KEYBOARD_WIDTH, KEYBOARD_HEIGHT);
    return true;
}

bool xr_init(int game_width, int game_height) {
    if (!xr_create_instance_and_session(game_width, game_height)) {
        LOGE("XR: bring-up failed - falling back to the normal window present path");
        xr_shutdown();
        return false;
    }
    return true;
}

XrOverlay *xr_session_get_menu_overlay(void) { return g_menu_overlay; }
XrOverlay *xr_session_get_keyboard_overlay(void) { return g_keyboard_overlay; }

void xr_poll_events(void) {
    if (g_instance == XR_NULL_HANDLE)
        return;

    while (true) {
        XrEventDataBuffer event = {XR_TYPE_EVENT_DATA_BUFFER};
        XrResult result = xrPollEvent(g_instance, &event);
        if (result == XR_EVENT_UNAVAILABLE)
            break;
        if (!xr_check(result, "xrPollEvent"))
            break;

        if (event.type == XR_TYPE_EVENT_DATA_SESSION_STATE_CHANGED) {
            XrEventDataSessionStateChanged *stateEvent = (XrEventDataSessionStateChanged *)&event;
            g_session_state = stateEvent->state;
            LOGI("XR: session state -> %d", (int)g_session_state);

            if (g_session_state == XR_SESSION_STATE_READY) {
                XrSessionBeginInfo beginInfo = {XR_TYPE_SESSION_BEGIN_INFO};
                beginInfo.primaryViewConfigurationType = XR_VIEW_CONFIGURATION_TYPE_PRIMARY_STEREO;
                g_session_running =
                    xr_check(xrBeginSession(g_session, &beginInfo), "xrBeginSession");
            } else if (g_session_state == XR_SESSION_STATE_STOPPING) {
                xr_check(xrEndSession(g_session), "xrEndSession");
                g_session_running = false;
            } else if (g_session_state == XR_SESSION_STATE_EXITING ||
                       g_session_state == XR_SESSION_STATE_LOSS_PENDING) {
                g_session_running = false;
            }
        }
    }
}

bool xr_is_session_running(void) { return g_session_running; }

bool xr_frame_begin(void) {
    g_have_acquired_game_image = false;
    if (!xr_is_session_running())
        return false;

    XrFrameWaitInfo waitInfo = {XR_TYPE_FRAME_WAIT_INFO};
    XrFrameState frameState = {XR_TYPE_FRAME_STATE};
    if (!xr_check(xrWaitFrame(g_session, &waitInfo, &frameState), "xrWaitFrame"))
        return false;
    g_predicted_display_time = frameState.predictedDisplayTime;
    g_frame_should_render = frameState.shouldRender;

    XrFrameBeginInfo beginInfo = {XR_TYPE_FRAME_BEGIN_INFO};
    if (!xr_check(xrBeginFrame(g_session, &beginInfo), "xrBeginFrame"))
        return false;

    if (!g_frame_should_render)
        return false;

    g_have_acquired_game_image = xr_swapchain_acquire(g_instance, &g_game_swapchain);
    return g_have_acquired_game_image;
}

void xr_frame_end(void) {
    if (g_session == XR_NULL_HANDLE)
        return;

    // Draw the laser-pointer reticle into the still-bound game swapchain
    // framebuffer before releasing it, so it composites on top of whatever
    // this frame's game content already drew there. Syncing actions
    // happens even when there's nothing to draw (no acquired image this
    // frame), so edge detection (trigger/menu-button clicks) doesn't miss a
    // frame. Only ever draws while neither overlay is visible (see
    // xr_input.c) - the menu/keyboard render into their own independent
    // swapchain images below.
    if (xr_is_session_running())
        xr_input_sync_and_draw(g_local_space, g_predicted_display_time,
                                g_have_acquired_game_image);

    if (g_have_acquired_game_image)
        xr_swapchain_release(g_instance, &g_game_swapchain);

    if (!xr_is_session_running())
        return;

    // Up to 3 layers: the game quad (if a frame was actually rendered),
    // and, while visible, the menu launcher and keyboard overlays - each
    // gets its own acquire/render/release cycle against its own swapchain
    // (see xr_overlay_render_and_build_layer()) any time before xrEndFrame,
    // unlike the game quad there's no per-frame engine rendering to wrap
    // around here, just each overlay's own blit.
    XrCompositionLayerQuad gameQuad = {XR_TYPE_COMPOSITION_LAYER_QUAD};
    gameQuad.space = g_local_space;
    gameQuad.eyeVisibility = XR_EYE_VISIBILITY_BOTH;
    gameQuad.subImage.swapchain = g_game_swapchain.swapchain;
    gameQuad.subImage.imageRect.extent.width = g_game_swapchain.width;
    gameQuad.subImage.imageRect.extent.height = g_game_swapchain.height;
    gameQuad.pose.orientation.w = 1.0f;
    gameQuad.pose.position.z = -QUAD_DISTANCE_METERS;
    gameQuad.size.width = QUAD_WIDTH_METERS;
    gameQuad.size.height =
        QUAD_WIDTH_METERS * (float)g_game_swapchain.height / (float)g_game_swapchain.width;

    const XrCompositionLayerBaseHeader *layers[3];
    uint32_t layerCount = 0;
    if (g_have_acquired_game_image)
        layers[layerCount++] = (XrCompositionLayerBaseHeader *)&gameQuad;

    XrCompositionLayerQuad menuQuad;
    if (g_frame_should_render && xr_overlay_render_and_build_layer(g_menu_overlay, &menuQuad)) {
        menuQuad.space = g_local_space;
        menuQuad.eyeVisibility = XR_EYE_VISIBILITY_BOTH;
        layers[layerCount++] = (XrCompositionLayerBaseHeader *)&menuQuad;
    }

    XrCompositionLayerQuad keyboardQuad;
    if (g_frame_should_render &&
        xr_overlay_render_and_build_layer(g_keyboard_overlay, &keyboardQuad)) {
        keyboardQuad.space = g_local_space;
        keyboardQuad.eyeVisibility = XR_EYE_VISIBILITY_BOTH;
        // Lets the keyboard sit on screen translucently while playing
        // rather than fully hiding whatever's behind it - KeyboardOverlay
        // writes real sub-255 alpha into the swapchain texture (see
        // OverlayPanel.compositeAndPublish()), already in Android's
        // default premultiplied format, which is what this flag's absence
        // of UNPREMULTIPLIED_ALPHA_BIT assumes. The menu launcher above
        // deliberately doesn't get this flag - it stays fully opaque.
        keyboardQuad.layerFlags = XR_COMPOSITION_LAYER_BLEND_TEXTURE_SOURCE_ALPHA_BIT;
        layers[layerCount++] = (XrCompositionLayerBaseHeader *)&keyboardQuad;
    }

    XrFrameEndInfo endInfo = {XR_TYPE_FRAME_END_INFO};
    endInfo.displayTime = g_predicted_display_time;
    endInfo.environmentBlendMode = XR_ENVIRONMENT_BLEND_MODE_OPAQUE;
    endInfo.layerCount = layerCount;
    endInfo.layers = layerCount > 0 ? layers : NULL;
    xr_check(xrEndFrame(g_session, &endInfo), "xrEndFrame");
}

void xr_shutdown(void) {
    xr_swapchain_destroy(&g_game_swapchain);

    if (g_local_space != XR_NULL_HANDLE)
        xrDestroySpace(g_local_space);
    g_local_space = XR_NULL_HANDLE;

    // Before the session/instance they were created from.
    xr_input_shutdown();
    xr_overlay_destroy(g_menu_overlay);
    g_menu_overlay = NULL;
    xr_overlay_destroy(g_keyboard_overlay);
    g_keyboard_overlay = NULL;

    if (g_session != XR_NULL_HANDLE)
        xrDestroySession(g_session);
    g_session = XR_NULL_HANDLE;

    if (g_instance != XR_NULL_HANDLE)
        xrDestroyInstance(g_instance);
    g_instance = XR_NULL_HANDLE;

    g_session_state = XR_SESSION_STATE_UNKNOWN;
    g_session_running = false;
}

void xr_get_game_quad_extent(float *distance_m, float *half_width_m, float *half_height_m) {
    *distance_m = QUAD_DISTANCE_METERS;
    *half_width_m = QUAD_WIDTH_METERS * 0.5f;
    *half_height_m =
        QUAD_WIDTH_METERS * 0.5f * (float)g_game_swapchain.height / (float)g_game_swapchain.width;
}
