#include "pch.h"

#include <cstdint>

#include "app.h"
#include "legacy_canvas_draw_merge_services.h"
#include "legacy_canvas_stroke_composite_services.h"
#include "legacy_canvas_stroke_erase_services.h"
#include "legacy_canvas_stroke_preview_services.h"
#include "legacy_ui_gl_dispatch.h"
#include "node_panel_grid.h"
#include "util.h"
#include "shape.h"
#include "renderer_gl/opengl_capabilities.h"

namespace {

void set_active_texture_unit(std::uint32_t unit_index)
{
    pp::legacy::ui_gl::activate_texture_unit(unit_index, "OpenGL VR");
}

void unbind_texture_2d()
{
    pp::legacy::ui_gl::unbind_texture_2d("OpenGL VR");
}

void apply_vr_ui_viewport(pp::renderer::gl::OpenGlViewportRect viewport)
{
    const auto status = pp::renderer::gl::apply_opengl_viewport(
        viewport,
        pp::renderer::gl::OpenGlViewportDispatch {
            .viewport = pp::legacy::ui_gl::set_opengl_viewport,
        });
    if (!status.ok())
        LOG("OpenGL VR UI viewport failed: %s", status.message);
}

void apply_vr_ui_scissor_test(bool enabled)
{
    const auto status = pp::renderer::gl::apply_opengl_scissor_test(
        enabled,
        pp::renderer::gl::OpenGlScissorTestDispatch {
            .enable = pp::legacy::ui_gl::enable_opengl_state,
            .disable = pp::legacy::ui_gl::disable_opengl_state,
        });
    if (!status.ok())
        LOG("OpenGL VR UI scissor test failed: %s", status.message);
}

void apply_vr_render_capability(std::uint32_t state, bool enabled)
{
    const auto status = pp::renderer::gl::apply_opengl_capability(
        state,
        enabled,
        pp::renderer::gl::OpenGlCapabilityDispatch {
            .enable = pp::legacy::ui_gl::enable_opengl_state,
            .disable = pp::legacy::ui_gl::disable_opengl_state,
        });
    if (!status.ok())
        LOG("OpenGL VR render state failed: %s", status.message);
}

bool query_vr_render_capability(std::uint32_t state)
{
    const auto result = pp::renderer::gl::query_opengl_capability_state(
        state,
        pp::renderer::gl::OpenGlCapabilityStateQueryDispatch {
            .is_enabled = pp::legacy::ui_gl::is_opengl_state_enabled,
        });
    if (!result.ok()) {
        LOG("OpenGL VR render state query failed: %s", result.status().message);
        return false;
    }
    return result.value();
}

void clear_vr_depth_buffer()
{
    const auto status = pp::renderer::gl::clear_opengl_buffers(
        pp::renderer::gl::framebuffer_depth_buffer_mask(),
        pp::renderer::gl::OpenGlBufferClearDispatch {
            .clear = pp::legacy::ui_gl::clear_opengl_buffer,
        });
    if (!status.ok())
        LOG("OpenGL VR depth clear failed: %s", status.message);
}

}

bool trigger_down = false;
cbuffer<glm::vec3> controller_points(10);
glm::vec3 controller_last_point;
const VRController* down_controller = nullptr;
glm::vec2 controller_cursor;
bool ui_inside = false;
bool ui_capture = false;
Sphere controller_ray;

bool App::vr_start()
{
    return start_platform_vr_mode();
}

void App::vr_stop()
{
    stop_platform_vr_mode();
}

void App::vr_draw_ui()
{
    uirtt.bindFramebuffer();
    uirtt.clear();
    apply_vr_ui_viewport(pp::renderer::gl::OpenGlViewportRect {
        .width = static_cast<std::int32_t>(uirtt.getWidth()),
        .height = static_cast<std::int32_t>(uirtt.getHeight()),
    });
    apply_vr_ui_scissor_test(true);
    auto observer = std::bind(&App::update_ui_observer, this, std::placeholders::_1);
    for (int i = 1; i < layout[main_id]->m_children.size(); i++)
        layout[main_id]->m_children[i]->watch(observer);
    //msgbox->watch(observer);
    apply_vr_ui_scissor_test(false);
    uirtt.unbindFramebuffer();
}

void App::vr_update(float dt)
{
    if (!vr_controllers_enabled)
        return;

    const auto vr_session = vr_session_snapshot();
    canvas->m_canvas->m_cam_fov = 60;
    float tan_fov = glm::tan(glm::radians(canvas->m_canvas->m_cam_fov / 2.f));
    glm::vec3 aspect = { (float)uirtt.getWidth() / (float)uirtt.getHeight(), 1.f, 1.f };
    glm::mat4 m = (
        glm::scale(glm::vec3(100)) *
        glm::transpose(canvas->m_canvas->m_cam_rot) *
        glm::translate(glm::vec3(0, 0, -1)) *
        glm::scale(aspect * tan_fov)
    );
    glm::mat4 mm = glm::inverse(m);
    auto o = glm::vec3(m * glm::vec4(0, 0, 0, 1));
    auto r = glm::normalize(glm::vec3(glm::vec4(1, 0, 0, 0) * mm));
    auto n = glm::normalize(glm::vec3(glm::vec4(0, 0, 1, 0) * mm));
    auto u = glm::normalize(glm::vec3(glm::vec4(0, 1, 0, 0) * mm));
    auto co = vr_session.vr_controllers[0].get_pos();
    auto cd = glm::mat3(vr_session.vr_controllers[0].m_mat) * glm::mat3(glm::eulerAngleX(glm::radians(-30.f))) * glm::vec3(0, 0, -1);
    ui_inside = false;
    glm::vec3 hit;
    float t;
    if (ray_intersect(co, cd, o, n, r, hit, t))
    {
        glm::mat4 plane_camera = m;
        glm::vec4 plane_local = (mm * glm::vec4(hit, 1));
        if (glm::all(glm::lessThanEqual(glm::abs(xy(plane_local)), glm::vec2(1.0f))))
        {
            controller_cursor.x = -(-plane_local.x * 0.5f - 0.5f) * width;
            controller_cursor.y =  (-plane_local.y * 0.5f + 0.5f) * height;
            if (!down_controller && ui_visible)
            {
                ui_task_async([this,p=controller_cursor] {
                    mouse_move(p.x, p.y, 1.f, kEventSource::Mouse, false);
                });
                ui_inside = true;
            }
        }
    }

    if (down_controller)
    {
        glm::vec3 head_position = vr_session.vr_head[3];
        glm::vec3 c_pos = glm::normalize(down_controller->get_pos() - head_position) * 800.f;
        controller_points.add(c_pos);
        auto p = controller_points.average();
        if (glm::distance(p, controller_last_point) > 1)
        {
            render_task_async([p, f=down_controller->get_trigger_value()]{
                Canvas::I->stroke_update(p, f);
                Canvas::I->stroke_draw();
            });
            controller_last_point = p;
            redraw = true;
        }
    }
}

void App::vr_analog(const VRController& c, VRController::kButton b, VRController::kAction a, glm::vec2 force)
{
    if (!vr_controllers_enabled)
        return;
    if (b == VRController::kButton::Trigger)
    {
        if (!down_controller && (ui_inside || ui_capture) && ui_visible)
        {
            if (a == VRController::kAction::Press)
            {
                ui_task_async([this, p=controller_cursor] {
                    mouse_down(0, p.x, p.y, 1.f, kEventSource::Mouse, false);
                });
                ui_capture = true;
            }
            else
            {
                ui_task_async([this, p=controller_cursor] {
                    mouse_up(0, p.x, p.y, kEventSource::Mouse, false);
                });
                ui_capture = false;
            }
        }
        else
        {
            if (a == VRController::kAction::Press)
            {
                glm::vec3 head_position = vr_session_snapshot().vr_head[3];
                glm::vec3 c_pos = glm::normalize(c.get_pos() - head_position) * 800.f;
                render_task_async([=] {
                    Canvas::I->stroke_start(c_pos, force.x);
                });
                controller_last_point = c_pos;
                controller_points.clear();
                down_controller = &c;
            }
            if (a == VRController::kAction::Release)
            {
                down_controller = nullptr;
                render_task_async([this] {
                    Canvas::I->stroke_end();
                });
            }
        }
    }
}

void App::vr_digital(const VRController& c, VRController::kButton b, VRController::kAction a, glm::vec2 axis)
{
    if (!vr_controllers_enabled)
        return;
    if ((b == VRController::kButton::A || b == VRController::kButton::Menu || b == VRController::kButton::Pad)
        && a == VRController::kAction::Press)
    {
        if (glm::length(axis) < 0.5f)
        {
            if (!ui_visible)
                Canvas::I->m_cam_rot = vr_rot;
            ui_task_async([this] {
                toggle_ui();
            });
        }
    }
}

void App::vr_draw(const glm::mat4& proj, const glm::mat4& camera, const glm::mat4& pose)
{
    //glm::mat4 ortho_proj = glm::ortho(0.f, box.z, 0.f, box.w, -1000.f, 1000.f);
    //glm::mat4 proj = glm::perspective(glm::radians(m_canvas->m_cam_fov), box.z / box.w, 0.01f, 1000.f);
    //glm::mat4 camera = glm::eulerAngleXY(m_canvas->m_cam_rot.y, m_canvas->m_cam_rot.x) *
    //    glm::translate(m_canvas->m_cam_pos);

    //         auto plane_mvp = proj * camera * transform *
    //             glm::scale(glm::vec3(sz, 1));

    //glm::extractEulerAngleXYZ(camera, vr_rot.y, vr_rot.x, vr_rot.z);
    glm::vec3 origin = glm::vec3(0, 0, -1) * glm::transpose(glm::mat3(pose));
    vr_rot = glm::lookAt({ 0, 0, 0 }, origin, { 0, 1, 0 });

    const bool blend = query_vr_render_capability(pp::renderer::gl::blend_state());
    const bool depth = query_vr_render_capability(pp::renderer::gl::depth_test_state());

    apply_vr_render_capability(pp::renderer::gl::blend_state(), false);
    apply_vr_render_capability(pp::renderer::gl::depth_test_state(), false);
    clear_vr_depth_buffer();

    for (int plane_index = 0; plane_index < 6; plane_index++)
    {
        auto plane_mvp = proj * camera * 
            glm::scale(glm::vec3(canvas->m_canvas->m_layers.size() * 20)) * 
            canvas->m_canvas->m_plane_transform[plane_index] *
            glm::translate(glm::vec3(0, 0, -1));

        pp::panopainter::setup_legacy_canvas_draw_merge_checkerboard_shader(
            pp::panopainter::LegacyCanvasDrawMergeCheckerboardUniforms {
                .mvp = plane_mvp,
                .colorize = false,
            });
        m_face_plane.draw_fill();
    }

    apply_vr_render_capability(pp::renderer::gl::blend_state(), true);
    apply_vr_render_capability(pp::renderer::gl::depth_test_state(), true);
    clear_vr_depth_buffer();

    for (size_t i = 0; i < canvas->m_canvas->m_layers.size(); i++)
    {
        auto layer_index = i;
        for (int plane_index = 0; plane_index < 6; plane_index++)
        {
            if (!(canvas->m_canvas->m_show_tmp && canvas->m_canvas->m_current_layer_idx == layer_index) &&
                (!canvas->m_canvas->m_layers[layer_index]->m_visible ||
                canvas->m_canvas->m_layers[layer_index]->m_opacity == .0f ||
                !canvas->m_canvas->m_layers[layer_index]->face(plane_index)))
                continue;

            int z = (int)(canvas->m_canvas->m_layers.size() - i);
            auto plane_mvp_z = proj * camera * 
                glm::scale(glm::vec3(100.f)) * 
                //glm::eulerAngleYXZ(yaw, pitch, roll) *
                canvas->m_canvas->m_plane_transform[plane_index] *
                glm::translate(glm::vec3(0, 0, -1));
            
            if (canvas->m_canvas->m_current_mode == kCanvasMode::Erase && canvas->m_canvas->m_show_tmp && canvas->m_canvas->m_current_layer_idx == layer_index)
            {
                sampler.bind(0);
                sampler.bind(1);
                sampler.bind(2);

                //ShaderManager::u_int(kShaderUniform::Lock, m_canvas->m_layers[layer_index]->m_alpha_locked);
                pp::panopainter::setup_legacy_stroke_erase_shader(
                    pp::panopainter::LegacyStrokeEraseUniforms {
                        .mvp = plane_mvp_z,
                        .texture_slot = 0,
                        .stroke_texture_slot = 1,
                        .mask_texture_slot = 2,
                        .alpha = canvas->m_canvas->m_layers[layer_index]->m_opacity,
                        .mask_enabled = canvas->m_canvas->m_smask_active,
                    });
                set_active_texture_unit(0);
                canvas->m_canvas->m_layers[layer_index]->rtt(plane_index).bindTexture();
                set_active_texture_unit(1);
                canvas->m_canvas->m_tmp[plane_index].bindTexture();
                set_active_texture_unit(2);
                canvas->m_canvas->m_smask.rtt(plane_index).bindTexture();
                m_face_plane.draw_fill();
                canvas->m_canvas->m_smask.rtt(plane_index).unbindTexture();
                set_active_texture_unit(1);
                canvas->m_canvas->m_tmp[plane_index].unbindTexture();
                set_active_texture_unit(0);
                canvas->m_canvas->m_layers[layer_index]->rtt(plane_index).unbindTexture();
            }
            else if (canvas->m_canvas->m_show_tmp && canvas->m_canvas->m_current_layer_idx == layer_index)
            {
                const auto& b = canvas->m_canvas->m_current_stroke->m_brush;
                sampler.bind(0);
                sampler.bind(1);
                sampler.bind(2);
                sampler.bind(3);
                sampler_stencil.bind(4);

                glm::vec2 patt_scale = glm::vec2(b->m_pattern_scale);
                if (b->m_pattern_flipx) patt_scale.x *= -1.f;
                if (b->m_pattern_flipy) patt_scale.y *= -1.f;

                pp::panopainter::setup_legacy_stroke_composite_shader(
                    pp::panopainter::LegacyStrokeCompositeUniforms {
                        .resolution = canvas->m_canvas->m_size,
                        .pattern = {
                            .scale = patt_scale,
                            .invert = static_cast<float>(b->m_pattern_invert),
                            .brightness = b->m_pattern_brightness,
                            .contrast = b->m_pattern_contrast,
                            .depth = b->m_pattern_depth,
                            .blend_mode = b->m_pattern_blend_mode,
                            .offset = Canvas::I->m_pattern_offset,
                        },
                        .mvp = plane_mvp_z,
                        .layer_alpha = canvas->m_canvas->m_layers[layer_index]->m_opacity,
                        .alpha_lock = canvas->m_canvas->m_layers[layer_index]->m_alpha_locked,
                        .mask_enabled = canvas->m_canvas->m_smask_active,
                        .use_fragcoord = false,
                        .blend_mode = b->m_blend_mode,
                        .use_dual = b->m_dual_enabled,
                        .dual_blend_mode = b->m_dual_blend_mode,
                        .dual_alpha = b->m_dual_opacity,
                        .use_pattern = b->m_pattern_enabled && !b->m_pattern_eachsample,
                    });

                set_active_texture_unit(0);
                canvas->m_canvas->m_layers[layer_index]->rtt(plane_index).bindTexture();
                set_active_texture_unit(1);
                canvas->m_canvas->m_tmp[plane_index].bindTexture();
                set_active_texture_unit(2);
                canvas->m_canvas->m_smask.rtt(plane_index).bindTexture();
                set_active_texture_unit(3);
                if (b->m_dual_enabled)
                    canvas->m_canvas->m_tmp_dual[plane_index].bindTexture();
                set_active_texture_unit(4);
                b->m_pattern_texture ?
                    b->m_pattern_texture->bind() :
                    unbind_texture_2d();
                m_face_plane.draw_fill();
                set_active_texture_unit(3);
                if (b->m_dual_enabled)
                    canvas->m_canvas->m_tmp_dual[plane_index].unbindTexture();
                set_active_texture_unit(2);
                canvas->m_canvas->m_smask.rtt(plane_index).unbindTexture();
                set_active_texture_unit(1);
                canvas->m_canvas->m_tmp[plane_index].unbindTexture();
                set_active_texture_unit(0);
                canvas->m_canvas->m_layers[layer_index]->rtt(plane_index).unbindTexture();
            }
            else
            {
                sampler.bind(0);
                pp::panopainter::setup_legacy_canvas_draw_merge_texture_alpha_shader(
                    pp::panopainter::LegacyCanvasDrawMergeTextureAlphaUniforms {
                        .mvp = plane_mvp_z,
                        .texture_slot = 0,
                        .alpha = canvas->m_canvas->m_layers[layer_index]->m_opacity,
                        .highlight = canvas->m_canvas->m_layers[layer_index]->m_hightlight,
                    });
                
                set_active_texture_unit(0);
                canvas->m_canvas->m_layers[layer_index]->rtt(plane_index).bindTexture();
                m_face_plane.draw_fill();
                canvas->m_canvas->m_layers[layer_index]->rtt(plane_index).unbindTexture();
            }
        }
    }

    App::I->grid->draw_heightmap(proj, camera, false);

    float tan_fov = glm::tan(glm::radians(canvas->m_canvas->m_cam_fov / 2.f));
    glm::vec3 aspect = { (float)uirtt.getWidth() / (float)uirtt.getHeight(), 1.f, 1.f };

    // draw the frame
    if (ui_visible)
    {
        auto mvp = proj * camera *
            glm::scale(glm::vec3(100)) *
            glm::transpose(canvas->m_canvas->m_cam_rot) *
            glm::translate(glm::vec3(0, 0, -1)) *
            glm::scale(aspect * tan_fov);
        pp::panopainter::setup_legacy_vr_color_shader(
            pp::panopainter::LegacyVrColorUniforms {
                .color = { 0, 0, 0, 1 },
                .mvp = mvp,
            });
        m_face_plane.draw_stroke();
    }

    apply_vr_render_capability(pp::renderer::gl::depth_test_state(), false);
    // draw the brush
    /*
    auto mode = dynamic_cast<CanvasModePen*>(canvas->m_canvas->modes[(int)canvas->m_canvas->m_current_mode][0]);
    if (ui_visible && ui_inside && mode)
    {
        auto pos = mode->m_resizing ? mode->m_size_pos_start : mode->m_cur_pos;
        if (keys[(int)kKey::KeyAlt] && !mode->m_resizing)
            pos.x = pos.x - canvas->m_canvas->m_current_brush->m_tip_size * 500;
        auto cur = (glm::vec2(pos.x / width, 1.f - pos.y / height) - 0.5f) * 2.f;

        ShaderManager::use(kShader::StrokePreview);
        ShaderManager::u_int(kShaderUniform::Tex, 0);
        ShaderManager::u_float(kShaderUniform::Alpha, canvas->m_canvas->m_current_brush->m_tip_flow);
        ShaderManager::u_int(kShaderUniform::DrawOutline, false);
        auto tip_color = glm::vec4(glm::vec3(canvas->m_canvas->m_current_brush->m_tip_color), 1);
        ShaderManager::u_vec4(kShaderUniform::Col, tip_color);
        ShaderManager::u_mat4(kShaderUniform::MVP,
            proj * camera *
            glm::scale(glm::vec3(100)) *
            glm::transpose(canvas->m_canvas->m_cam_rot) *
            glm::translate(glm::vec3(cur * glm::vec2(aspect * tan_fov), -1)) *
            glm::scale(glm::vec3(canvas->m_canvas->m_current_brush->m_tip_size / height)) *
            glm::eulerAngleZ(canvas->m_canvas->m_current_brush->m_tip_angle  * (float)(M_PI * 2.0))
        );
        apply_vr_render_capability(pp::renderer::gl::blend_state(), true);
        set_active_texture_unit(0);
        auto& tex = *canvas->m_canvas->m_current_brush->m_tip_texture;
        tex.bind();
        sampler_linear.bind(0);
        m_face_plane.draw_fill();
        tex.unbind();
    }
     */
    
    // draw the 2D UI
    if (ui_visible)
    {
        auto mvp = proj * camera *
            glm::scale(glm::vec3(100)) *
            glm::transpose(canvas->m_canvas->m_cam_rot) *
            glm::translate(glm::vec3(0, 0, -1)) *
            glm::scale(aspect * tan_fov);
        sampler_linear.bind(0);
        pp::panopainter::setup_legacy_canvas_draw_merge_texture_shader(
            pp::panopainter::LegacyCanvasDrawMergeTextureUniforms {
                .mvp = mvp,
                .texture_slot = 0,
            });
        set_active_texture_unit(0);
        uirtt.bindTexture();
        m_face_plane.draw_fill();
        uirtt.unbindTexture();
    }

    // draw the cursor
    if (ui_visible /*&& ui_inside*/)
    {
        auto cur = (glm::vec2(cursor.x / width * zoom, 1.f - cursor.y / height * zoom) - 0.5f) * 2.f;
        auto mvp = proj * camera *
            glm::scale(glm::vec3(100)) *
            glm::transpose(canvas->m_canvas->m_cam_rot) *
            glm::translate(glm::vec3(cur * glm::vec2(aspect * tan_fov), -1));
        pp::panopainter::setup_legacy_vr_color_shader(
            pp::panopainter::LegacyVrColorUniforms {
                .color = { 0, 0, 0, 1 },
                .mvp = mvp * glm::scale(glm::vec3(.01f)),
            });
        m_face_plane.draw_fill();
        pp::panopainter::setup_legacy_vr_color_shader(
            pp::panopainter::LegacyVrColorUniforms {
                .color = { 1, 1, 1, 1 },
                .mvp = mvp * glm::scale(glm::vec3(.005f)),
            });
        m_face_plane.draw_fill();
    }


    // draw the motion controller sphere
    if (vr_controllers_enabled && ui_visible && ui_inside)
    {
        const auto vr_session = vr_session_snapshot();
        auto mvp = proj * camera * vr_session.vr_controllers[0].m_mat * glm::eulerAngleX(glm::radians(-30.f));
        pp::panopainter::setup_legacy_vr_color_shader(
            pp::panopainter::LegacyVrColorUniforms {
                .color = { 1, 0, 1, 1 },
                .mvp = mvp * glm::scale(glm::vec3(.0125f, .0125f, .07f)),
            });
        sphere.draw_fill();
    }


    // draw the motion controller brush
    if (vr_controllers_enabled && (!ui_visible || !ui_inside))
    {
        const auto vr_session = vr_session_snapshot();
        glm::vec3 cpos = vr_session.vr_controllers[0].get_pos() - xyz(pose[3]);
        auto pos = glm::translate(glm::normalize(cpos) * 100.f);
        auto tip_color = glm::vec4(glm::vec3(canvas->m_canvas->m_current_brush->m_tip_color), 1);
        pp::panopainter::setup_legacy_vr_stroke_preview_shader(
            pp::panopainter::LegacyVrStrokePreviewUniforms {
                .texture_slot = 0,
                .alpha = canvas->m_canvas->m_current_brush->m_tip_flow,
                .draw_outline = false,
                .color = tip_color,
                .mvp = proj * camera * pos *
                    glm::inverse(glm::lookAt({ 0, 0, 0 }, cpos, { 0, 1, 0 })) *
                    glm::scale(glm::vec3(canvas->m_canvas->m_current_brush->m_tip_size * 100.f / height)) *
                    glm::eulerAngleZ(canvas->m_canvas->m_current_brush->m_tip_angle * (float)(M_PI * 2.0)),
            });
        apply_vr_render_capability(pp::renderer::gl::blend_state(), true);
        set_active_texture_unit(0);
        auto& tex = *canvas->m_canvas->m_current_brush->m_tip_texture;
        tex.bind();
        sampler_linear.bind(0);
        m_face_plane.draw_fill();
        tex.unbind();
    }

    glm::mat4 ortho_proj(1);

/*
    for (auto& mode : *canvas->m_canvas->m_mode)
        mode->on_Draw(ortho_proj, proj, camera);
    
    glDisable(pp::renderer::gl::depth_test_state());
    if (canvas->m_canvas->m_smask_active)
    {
        canvas->m_canvas->modes[(int)kCanvasMode::MaskFree][0]->on_Draw(ortho_proj, proj, camera);
        canvas->m_canvas->modes[(int)kCanvasMode::MaskLine][0]->on_Draw(ortho_proj, proj, camera);
    }
    
    // keep drawing the grids
    if (canvas->m_canvas->m_current_mode != kCanvasMode::Grid)
        for (auto& mode : Canvas::modes[(int)kCanvasMode::Grid])
            mode->on_Draw(ortho_proj, proj, camera);
*/

    apply_vr_render_capability(pp::renderer::gl::blend_state(), blend);
    apply_vr_render_capability(pp::renderer::gl::depth_test_state(), depth);
    sampler.unbind();
}

