#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;
ovrDeviceID controller = -1;

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()
{
    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 (controller == -1)
    {
        // init controller
        ovrInputCapabilityHeader capsHeader;
        for (int i = 0; ; i++)
        {
            LOG("")
            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
                            controller = capsHeader.DeviceID;
                            LOG("found controller id %d", (int)controller);
                            break;
                        }
                    }
                }
            }
            else
            {
                break;
            }
        }
    }

    if (controller != -1)
    {
        ovrTracking trackingState;
        vrapi_GetInputTrackingState(ovr_context, controller, predictedDisplayTime, &trackingState);
        glm::vec3 c_pos = glm::make_vec3((float*)&trackingState.HeadPose.Pose.Position) - head_pos;
        auto c_rot_ovr = ovrMatrix4f_CreateFromQuaternion(&trackingState.HeadPose.Pose.Orientation);
        glm::mat4 c_rot = glm::make_mat4((float*)&c_rot_ovr);
        glm::mat4 c_mat = glm::translate(c_pos);

        App::I.vr_controller = c_mat;
        App::I.vr_controller_pos = c_pos;

        // update controllers
        ovrInputStateTrackedRemote remote_state;
        remote_state.Header.ControllerType = ovrControllerType_TrackedRemote;
        vrapi_GetCurrentInputState(ovr_context, controller, &remote_state.Header);

        static cbuffer<glm::vec3> controller_points(10);
        static bool down = false;
        static glm::vec3 last_point;

        auto cur = glm::normalize(c_pos) * 800.f;

        if (remote_state.IndexTrigger > 0.f)
        {
            if (!down)
            {
                controller_points.clear();
                last_point = cur;
                Canvas::I->stroke_start(cur, remote_state.IndexTrigger);
                down = true;
            }
            else
            {
                controller_points.add(cur);
                auto p = controller_points.average();
                if (glm::distance(last_point, p) > 0.1f)
                {
                    Canvas::I->stroke_update(p, remote_state.IndexTrigger);
                    App::I.update(0);
                    last_point = p;
                }
            }
        }
        else if (down)
        {
            down = false;
            Canvas::I->stroke_end();
        }
    }

    // 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;
}