#include "pch.h"
#include "oculus_vr.h"
#include "log.h"
#include "rtt.h"
#include "app.h"
#include "canvas.h"
#include <unistd.h>

ovrJava java;
RTT ovr_eyes[3][2];
ovrTextureSwapChain* swap_chain[2];
int swap_chain_count = 0;
int swap_chain_index = 0;
ovrMobile* ovr_context = nullptr;
uint64_t ovr_frame = 0;

struct QuestController : public VRController
{
    ovrDeviceID id = -1;
    ovrTracking tracking;
    ovrInputStateTrackedRemote state;
    void update_state(double predictedDisplayTime, glm::vec3 head_pos)
    {
        vrapi_GetInputTrackingState(ovr_context, id, predictedDisplayTime, &tracking);
        glm::vec3 c_pos = glm::make_vec3((float*)&tracking.HeadPose.Pose.Position) - head_pos;
        auto c_rot_ovr = ovrMatrix4f_CreateFromQuaternion(&tracking.HeadPose.Pose.Orientation);
        glm::mat4 c_rot = glm::make_mat4((float*)&c_rot_ovr);
        m_mat = glm::translate(c_pos) * c_rot;

        // update controllers
        state.Header.ControllerType = ovrControllerType_TrackedRemote;
        vrapi_GetCurrentInputState(ovr_context, id, &state.Header);

        update_analog(kButton::Trigger, {state.IndexTrigger, 0});
        update_analog(kButton::Grip, {state.GripTrigger, 0});
        update_analog(kButton::Pad, glm::make_vec2((float*)&state.Joystick));
        update_digital(kButton::Trigger, state.Buttons & ovrButton_Trigger, {state.IndexTrigger, 0});
        update_digital(kButton::Pad, state.Buttons & ovrButton_Joystick, glm::make_vec2((float*)&state.Joystick));
        update_digital(kButton::A, state.Buttons & ovrButton_A);
    }
    void update_digital(kButton b, bool pressed, glm::vec2 axis = {0, 0})
    {
        if (pressed && !m_buttons[(int)b])
        {
            m_buttons[(int)b] = true;
            App::I.vr_digital(*this, b, kAction::Press, axis);
        }
        if (!pressed && m_buttons[(int)b])
        {
            m_buttons[(int)b] = false;
            App::I.vr_digital(*this, b, kAction::Release, axis);
        }
    }
    void update_analog(kButton b, glm::vec2 force)
    {
        float l = glm::compMax(glm::abs(force));
        const float zero = 0.01f;
        if (l > zero && !m_analog_buttons[(int)b])
        {
            m_analog_buttons[(int)b] = true;
            App::I.vr_analog(*this, b, kAction::Press, force);
        }
        if (l <= zero && m_analog_buttons[(int)b])
        {
            m_analog_buttons[(int)b] = false;
            App::I.vr_analog(*this, b, kAction::Release, force);
        }
    }

    float get_trigger_value() const override {
        if (id == -1)
            return 0.f;
        return state.IndexTrigger;
    }
};

QuestController controllers[2];

void oculus_init(JavaVM* vm, JNIEnv* jni, jobject activity_class)
{
    java.Vm = vm;
    java.Env = jni;
    java.ActivityObject = activity_class;

    LOG("init OVR");
    const ovrInitParms initParms = vrapi_DefaultInitParms( &java );
    int32_t initResult = vrapi_Initialize( &initParms );
    if ( initResult != VRAPI_INITIALIZE_SUCCESS )
    {
        // If intialization failed, vrapi_* function calls will not be available.
        exit( 0 );
    }
}

void oculus_init_vr(EGLDisplay display, EGLContext context, ANativeWindow* surface)
{
    LOG("init swapchain");
    int rtt_w = vrapi_GetSystemPropertyInt(&java, VRAPI_SYS_PROP_SUGGESTED_EYE_TEXTURE_WIDTH);
    int rtt_h = vrapi_GetSystemPropertyInt(&java, VRAPI_SYS_PROP_SUGGESTED_EYE_TEXTURE_HEIGHT);
    LOG("ovr suggested texture size %d %d", rtt_w, rtt_h);
    for (int eye = 0; eye < 2; eye++)
    {
        swap_chain[eye] = vrapi_CreateTextureSwapChain3(VRAPI_TEXTURE_TYPE_2D, GL_RGBA8, rtt_w, rtt_h, 1, 3);
        swap_chain_count = vrapi_GetTextureSwapChainLength(swap_chain[eye]);
        for (int i = 0; i < swap_chain_count; i++)
        {
            auto texid = vrapi_GetTextureSwapChainHandle(swap_chain[eye], i);
            if (!ovr_eyes[i][eye].create(rtt_w, rtt_h, texid, GL_RGBA8, true))
            {
                ovr_eyes[i][eye].bindFramebuffer();
                ovr_eyes[i][eye].clear({1, 0, 1, 1});
                ovr_eyes[i][eye].unbindFramebuffer();
                LOG("FAILED create fb for eye %d", eye);
            }
            else
            {
                LOG("create eye %d", eye);
            }
        }
    }

    LOG("vrapi_DefaultModeParms");
    ovrModeParms parms = vrapi_DefaultModeParms( &java );
    // No need to reset the FLAG_FULLSCREEN window flag when using a View
    //parms.Flags &= ~VRAPI_MODE_FLAG_RESET_WINDOW_FULLSCREEN;

    parms.Flags |= VRAPI_MODE_FLAG_NATIVE_WINDOW;
    parms.Display = (size_t)display;
    parms.WindowSurface = (size_t)surface;
    parms.ShareContext = (size_t)context;
    LOG("enter vr mode");
    ovr_context = vrapi_EnterVrMode(&parms);
    if (!ovr_context)
    {
        LOG("EnterVRMode FAILED");
    }
    LOG("vr mode entered");

    vrapi_SetClockLevels(ovr_context, 2, 2);
    vrapi_SetPerfThread(ovr_context, VRAPI_PERF_THREAD_TYPE_MAIN, gettid());
    vrapi_SetPerfThread(ovr_context, VRAPI_PERF_THREAD_TYPE_RENDERER, 0);
}

void oculus_draw(float dt)
{
    const double predictedDisplayTime = vrapi_GetPredictedDisplayTime(ovr_context, ovr_frame);
    const ovrTracking2 tracking = vrapi_GetPredictedTracking2(ovr_context, predictedDisplayTime);
    auto pose_ovr = ovrMatrix4f_CreateFromQuaternion(&tracking.HeadPose.Pose.Orientation);
    auto pose_ovr_tp = ovrMatrix4f_Transpose(&pose_ovr);
    glm::mat4 pose = glm::make_mat4(reinterpret_cast<float*>(&pose_ovr_tp));
    glm::vec3 head_pos = glm::make_vec3((float*)&tracking.HeadPose.Pose.Position);

    if (controllers[0].id == -1)
    {
        // init controller
        ovrInputCapabilityHeader capsHeader;
        for (int i = 0; ; i++)
        {
            if ( vrapi_EnumerateInputDevices( ovr_context, i, &capsHeader ) >= 0 )
            {
                if ( capsHeader.Type == ovrControllerType_TrackedRemote )
                {
                    ovrInputTrackedRemoteCapabilities remoteCaps;
                    remoteCaps.Header = capsHeader;
                    if ( vrapi_GetInputDeviceCapabilities( ovr_context, &remoteCaps.Header ) >= 0 )
                    {
                        // remote is connected
                        if (remoteCaps.ControllerCapabilities & ovrControllerCaps_RightHand)
                        {
                            // right controller found
                            controllers[0].id = capsHeader.DeviceID;
                            LOG("found controller id %d", (int)controllers[0].id);
                            break;
                        }
                    }
                }
            }
            else
            {
                break;
            }
        }
    }

    if (controllers[0].id != -1)
    {
        controllers[0].update_state(predictedDisplayTime, head_pos);
        App::I.vr_controllers[0] = controllers[0];
    }

    App::I.vr_update(dt);

    // update hmd
    auto layer = vrapi_DefaultLayerProjection2();
    //ovrVector4f red = {1, 1, 0, 1};
    //auto layer = vrapi_DefaultLayerSolidColorProjection2(&red);
    layer.HeadPose = tracking.HeadPose;
    for (int eye = 0; eye < 2; eye++)
    {
        auto& rtt = ovr_eyes[swap_chain_index][eye];
        rtt.bindFramebuffer();
        rtt.clear({1, 0, 1, 1});
        glViewport(0, 0, rtt.getWidth(), rtt.getHeight());
        auto proj_ovr = ovrMatrix4f_Transpose(&tracking.Eye[eye].ProjectionMatrix);
        glm::mat4 proj = glm::make_mat4(reinterpret_cast<const float*>(&proj_ovr));
        auto view_ovr = ovrMatrix4f_Transpose(&tracking.Eye[eye].ViewMatrix);
        glm::mat4 view = glm::make_mat4(reinterpret_cast<const float*>(&view_ovr));
        App::I.vr_draw(proj, view, pose);
        rtt.unbindFramebuffer();
        layer.Textures[eye].ColorSwapChain = swap_chain[eye];
        layer.Textures[eye].SwapChainIndex = swap_chain_index;
        layer.Textures[eye].TexCoordsFromTanAngles =
                ovrMatrix4f_TanAngleMatrixFromProjection(&tracking.Eye[eye].ProjectionMatrix);
    }
    layer.Header.Flags |= VRAPI_FRAME_LAYER_FLAG_CHROMATIC_ABERRATION_CORRECTION;

    const ovrLayerHeader2 * layers[] = { &layer.Header };

    ovrSubmitFrameDescription2 frameDesc = { 0 };
    frameDesc.Flags = 0;
    frameDesc.SwapInterval = 1;
    frameDesc.FrameIndex = ovr_frame;
    frameDesc.DisplayTime = predictedDisplayTime;
    frameDesc.LayerCount = 1;
    frameDesc.Layers = layers;

    vrapi_SubmitFrame2(ovr_context, &frameDesc);

    ovr_frame++;
    swap_chain_index = (swap_chain_index + 1) % swap_chain_count;
}