#include "pch.h"
#include "hmd.h"
#include "log.h"
#include <array>

std::map<ViveController::kButton, ViveController::kButtonMask> ViveController::m_mask {
    { ViveController::kButton::Trigger, ViveController::kButtonMask::TriggerBit },
    { ViveController::kButton::Pad, ViveController::kButtonMask::PadBit },
    { ViveController::kButton::Menu, ViveController::kButtonMask::MenuBit },
    { ViveController::kButton::Grip, ViveController::kButtonMask::GripBit },
    { ViveController::kButton::A, ViveController::kButtonMask::ButtonA },
};

bool Vive::Initialize()
{
    vr::EVRInitError error;
    m_hmd = vr::VR_Init(&error, vr::EVRApplicationType::VRApplication_Scene);
    if (error != vr::EVRInitError::VRInitError_None)
        return false;
    
    if (!(m_comp = vr::VRCompositor()))
        return false;
    
    m_settings = vr::VRSettings();

    m_hmd->GetRecommendedRenderTargetSize(&m_eye_width, &m_eye_height);
    App::I->render_task([&]{
        for (int eye = 0; eye < 2; ++eye)
        {
            m_eyes[eye].create(m_eye_width, m_eye_height);
            m_eyes[eye].bindFramebuffer();
            m_eyes[eye].clear({ 1, 0, 0, 1 });
            m_eyes[eye].unbindFramebuffer();
        }
    });

    vr::CVRSettingHelper s(m_settings);
    float timeout = s.GetFloat(vr::k_pch_Power_Section, vr::k_pch_Power_TurnOffScreensTimeout_Float);
    s.SetFloat(vr::k_pch_Power_Section, vr::k_pch_Power_TurnOffScreensTimeout_Float, 0.5);

    // init buttons state
    for (auto& c : m_controllers)
    {
        c.m_buttons.fill(false);
        c.m_analog_buttons.fill(false);
    }

    return true;
}

void Vive::Terminate()
{
    vr::VR_Shutdown();
    for (int eye = 0; eye < 2; ++eye)
        m_eyes[eye].destroy();
}

void Vive::Update()
{
    if (!(m_hmd && m_comp))
        return;

    vr::TrackedDevicePose_t poses[vr::k_unMaxTrackedDeviceCount];
    m_comp->WaitGetPoses(poses, vr::k_unMaxTrackedDeviceCount, NULL, 0);
    auto& current_pose = poses[vr::k_unTrackedDeviceIndex_Hmd];

    for (int eye = 0; eye < 2; ++eye)
    {

        // Get view and projection matrices
        vr::HmdMatrix44_t p = m_hmd->GetProjectionMatrix((vr::EVREye)eye, 0.2f, 1000.f);
        float data_eye_proj[16] = {
            p.m[0][0], p.m[1][0], p.m[2][0], p.m[3][0],
            p.m[0][1], p.m[1][1], p.m[2][1], p.m[3][1],
            p.m[0][2], p.m[1][2], p.m[2][2], p.m[3][2],
            p.m[0][3], p.m[1][3], p.m[2][3], p.m[3][3],
        };

        vr::HmdMatrix34_t e = m_hmd->GetEyeToHeadTransform((vr::EVREye)eye);
        float data_eye_pose[16] = {
            e.m[0][0], e.m[1][0], e.m[2][0], 0,
            e.m[0][1], e.m[1][1], e.m[2][1], 0,
            e.m[0][2], e.m[1][2], e.m[2][2], 0,
            e.m[0][3], e.m[1][3], e.m[2][3], 1,
        };

        vr::HmdMatrix34_t h = current_pose.mDeviceToAbsoluteTracking;
        float data_hmd_pose[16] = {
            h.m[0][0], h.m[1][0], h.m[2][0], 0,
            h.m[0][1], h.m[1][1], h.m[2][1], 0,
            h.m[0][2], h.m[1][2], h.m[2][2], 0,
            //        0,         0,         0, 1, // rotation only
            h.m[0][3], h.m[1][3], h.m[2][3], 1,
        };

        if (current_pose.bPoseIsValid && !m_position_valid)
        {
            m_initial_position = { h.m[0][3], h.m[1][3], h.m[2][3] };
            m_position_valid = true;
        }

        auto mat_proj = glm::make_mat4(data_eye_proj);
        auto mat_pose = glm::inverse(glm::make_mat4(data_hmd_pose));
        auto mat_eye = glm::inverse(glm::make_mat4(data_eye_pose));

        m_proj[eye] = mat_proj;
        m_view[eye] = mat_eye * mat_pose/* * glm::translate(m_initial_position)*/;
        m_pose = glm::make_mat4(data_hmd_pose);

        // invalidate controller state
        for (auto& c : m_controllers)
            c.m_valid = false;
        int controller_index = 0;

        for (int id = 0; id < vr::k_unMaxTrackedDeviceCount; id++)
        {
            // device not tracking
            if (!poses[id].bPoseIsValid)
                continue;

            switch (m_hmd->GetTrackedDeviceClass(id))
            {
                case vr::TrackedDeviceClass_HMD:
                    m_active = (m_hmd->GetTrackedDeviceActivityLevel(id) == vr::k_EDeviceActivityLevel_UserInteraction);
                    break;
                case vr::TrackedDeviceClass_Controller:
                {
                    m_controllers[controller_index].m_valid = true;
                    m_hmd->GetControllerState(id, &m_controllers[controller_index].m_state, sizeof(vr::VRControllerState_t));
                    m_controllers[controller_index].m_mat = /*glm::translate(-m_initial_position) * */Pose2Mat(poses[id].mDeviceToAbsoluteTracking);

                    auto pressed_mask = m_controllers[controller_index].m_state.ulButtonPressed;
                    for (uint8_t bi = 0; bi < (uint8_t)ViveController::kButton::COUNT; bi++)
                    {
                        auto b = (ViveController::kButton)bi;
                        glm::vec2 force = { m_controllers[controller_index].m_state.rAxis[(int)b].x,
                            m_controllers[controller_index].m_state.rAxis[(int)b].y };

                        if (glm::compMax(glm::abs(force)) > 0.01f && !m_controllers[controller_index].m_analog_buttons[bi])
                        {
                            m_controllers[controller_index].m_analog_buttons[bi] = true;
                            if (on_analog_button)
                                on_analog_button(m_controllers[controller_index], b, ViveController::kAction::Press, force);
                        }
                        if (glm::compMax(glm::abs(force)) <= 0.01f && m_controllers[controller_index].m_analog_buttons[bi])
                        {
                            m_controllers[controller_index].m_analog_buttons[bi] = false;
                            if (on_analog_button)
                                on_analog_button(m_controllers[controller_index], b, ViveController::kAction::Release, force);
                        }

                        bool b_pressed = pressed_mask & (uint64_t)m_controllers[controller_index].m_mask[b];
                        if (b_pressed == true && !m_controllers[controller_index].m_buttons[bi])
                        {
                            m_controllers[controller_index].m_buttons[bi] = true;
                            if (on_button)
                                on_button(m_controllers[controller_index], b, ViveController::kAction::Press, force);
                        }
                        if (b_pressed == false && m_controllers[controller_index].m_buttons[bi])
                        {
                            m_controllers[controller_index].m_buttons[bi] = false;
                            if (on_button)
                                on_button(m_controllers[controller_index], b, ViveController::kAction::Release, force);
                        }
                    }
                    controller_index++;
                }
                default:
                    break;
            }
        }
    }
}

void Vive::Draw()
{
    if (!Valid())
        return;

    for (int eye = 0; eye < 2; ++eye)
    {
        m_eyes[eye].bindFramebuffer();
        m_eyes[eye].clear();
        glViewport(0, 0, m_eyes[eye].getWidth(), m_eyes[eye].getHeight());

        if (on_draw)
            on_draw(m_proj[eye], m_view[eye], m_pose);

        m_eyes[eye].unbindFramebuffer();
    }

    vr::Texture_t eyeTexture0 = { (void*)m_eyes[0].getTextureID(), vr::ETextureType::TextureType_OpenGL, vr::ColorSpace_Linear };
    if (auto err = m_comp->Submit(vr::EVREye::Eye_Left, &eyeTexture0) != vr::EVRCompositorError::VRCompositorError_None)
    {
        LOG("Compositor error %d", err);
        comp_attempts++;
        return;
    }
    vr::Texture_t eyeTexture1 = { (void*)m_eyes[1].getTextureID(), vr::ETextureType::TextureType_OpenGL, vr::ColorSpace_Linear };
    if (auto err = m_comp->Submit(vr::EVREye::Eye_Right, &eyeTexture1) != vr::EVRCompositorError::VRCompositorError_None)
    {
        LOG("Compositor error %d", err);
        comp_attempts++;
        return;
    }
    comp_attempts = 0;

    //glFlush();
}

bool Vive::Valid()
{
    return m_hmd && m_comp && comp_attempts < 10;
}

glm::mat4 Vive::Pose2Mat(vr::HmdMatrix34_t const& m)
{
    float data[16] = {
        m.m[0][0], m.m[1][0], m.m[2][0], 0,
        m.m[0][1], m.m[1][1], m.m[2][1], 0,
        m.m[0][2], m.m[1][2], m.m[2][2], 0,
        m.m[0][3], m.m[1][3], m.m[2][3], 1,
    };
    return glm::make_mat4(data);
}