#include "pch.h"
#include "app_core/grid_ui.h"
#include "legacy_canvas_draw_merge_services.h"
#include "legacy_grid_ui_services.h"
#include "legacy_ui_overlay_services.h"
#include "log.h"
#include "node_panel_grid.h"
#include "canvas.h"
#include "app.h"
#include "image.h"
#include "legacy_ui_gl_dispatch.h"
#include "renderer_gl/opengl_capabilities.h"
#include "util.h"

namespace {

void apply_grid_capability(std::uint32_t state, bool enabled)
{
    pp::legacy::ui_gl::set_capability(state, enabled, "Grid");
}

bool query_grid_capability(std::uint32_t state)
{
    return pp::legacy::ui_gl::query_capability(state, "Grid");
}

void set_grid_active_texture_unit(std::uint32_t unit_index)
{
    pp::legacy::ui_gl::activate_texture_unit(unit_index, "Grid");
}

void apply_grid_viewport(std::int32_t x, std::int32_t y, std::int32_t width, std::int32_t height)
{
    pp::legacy::ui_gl::apply_viewport(x, y, width, height, "Grid");
}

pp::renderer::gl::OpenGlViewportRect query_grid_viewport()
{
    return pp::legacy::ui_gl::query_viewport_rect("Grid");
}

void clear_grid_depth_buffer()
{
    pp::legacy::ui_gl::clear_depth_buffer("Grid");
}

void apply_grid_color_mask(std::uint8_t r, std::uint8_t g, std::uint8_t b, std::uint8_t a)
{
    pp::legacy::ui_gl::set_color_write_mask(r, g, b, a, "Grid");
}

}

Node* NodePanelGrid::clone_instantiate() const
{
    return new NodePanelGrid();
}

void NodePanelGrid::clone_finalize(Node* dest) const
{
    NodePanelGrid* n = static_cast<NodePanelGrid*>(dest);
    n->init_controls();
}

void NodePanelGrid::init()
{
    init_template_file("data/dialogs/panel-grid.xml", "tpl-panel-grid");
    init_controls();
}

void NodePanelGrid::init_controls()
{
    m_groud_opacity = find<NodeSliderH>("grid-ground-opacity");
    m_groud_value = find<NodeSliderH>("grid-ground-value");
    m_groud_resolution = find<NodeComboBox>("grid-ground-resolution");
    m_groud_offset = find<NodeSliderH>("grid-ground-offset");

    m_hm_preview = find<NodeImageTexture>("grid-heightmap-preview");
    m_hm_load = find<NodeButton>("grid-heightmap-load");
    m_hm_clear = find<NodeButton>("grid-heightmap-clear");
    m_hm_reload = find<NodeButton>("grid-heightmap-reload");
    m_hm_wireframe = find<NodeSliderH>("grid-heightmap-wireframe");
    m_hm_thickness = find<NodeSliderH>("grid-heightmap-thickness");
    m_hm_height = find<NodeSliderH>("grid-heightmap-height");
    m_hm_lyaw = find<NodeSliderH>("grid-heightmap-lyaw");
    m_hm_lpitch = find<NodeSliderH>("grid-heightmap-lpitch");
    m_hm_ambient = find<NodeSliderH>("grid-heightmap-ambient");
    m_hm_radius = find<NodeSliderH>("grid-heightmap-radius");
    m_hm_shading = find<NodeComboBox>("grid-heightmap-shading");
    m_hm_texres = find<NodeComboBox>("grid-heightmap-texres");
    m_hm_samples = find<NodeComboBox>("grid-heightmap-samples");
    m_render = find<NodeButton>("grid-render");
    m_commit = find<NodeButton>("grid-commit");

    m_hm_preview->SetHeight(0);
    m_hm_plane.create(1, 1, 100);

    m_hm_preview_nav = m_hm_preview->add_child<NodeHeightmapOverlay>();
    m_hm_preview_nav->SetWidthP(100);
    m_hm_preview_nav->SetHeightP(100);

    //m_hm_height->on_value_changed = update_hm;
    m_groud_resolution->on_select = [this](Node* target, int index) {
        if (m_hm_image.data())
            m_hm_plane.create(1, 1, m_hm_image, get_resolution(), get_height());
        else
            m_hm_plane.create(1, 1, 100 * get_resolution());
        m_rt_dirty = true;
        LOG("resolution index %d", index);
    };

    m_hm_height->on_value_final = [this](Node* target, float v) {
        if (m_hm_image.data())
            m_hm_plane.create(1, 1, m_hm_image, get_resolution(), get_height());
        m_rt_dirty = true;
        LOG("height value %f", v);
    };

    m_hm_shading->on_select = [this](Node*, int index) {
        m_shade_mode = (ShadeMode)index;
    };

    m_hm_load->on_click = [this](Node*) {
        const auto plan = pp::app::plan_grid_heightmap_pick();
        const auto status = pp::panopainter::execute_legacy_grid_ui_plan(*this, plan);
        if (!status.ok())
            LOG("Grid heightmap pick action failed: %s", status.message);
    };

    m_hm_clear->on_click = [this](Node*)
    {
        const auto plan = pp::app::plan_grid_heightmap_clear(static_cast<bool>(m_hm_image.data()));
        const auto status = pp::panopainter::execute_legacy_grid_ui_plan(*this, plan);
        if (!status.ok())
            LOG("Grid heightmap clear action failed: %s", status.message);
    };

    m_hm_reload->on_click = [this](Node*)
    {
        load_heightmap_file(m_file_path, false);
    };

    m_render->on_click = [this](Node*)
    {
        const auto renderer_features = ShaderManager::render_device_features();
        const auto plan = pp::app::plan_grid_lightmap_render(
            static_cast<bool>(m_hm_image.data()),
            renderer_features.float32_render_targets,
            renderer_features.float16_render_targets,
            get_texres(),
            get_samples());
        if (!plan)
            return;
        const auto status = pp::panopainter::execute_legacy_grid_ui_plan(*this, plan.value());
        if (!status.ok())
            LOG("Grid lightmap render action failed: %s", status.message);
    };
    m_commit->on_click = [this](Node*)
    {
        const auto plan = pp::app::plan_grid_heightmap_commit(Canvas::I != nullptr);
        const auto status = pp::panopainter::execute_legacy_grid_ui_plan(*this, plan);
        if (!status.ok())
            LOG("Grid heightmap commit action failed: %s", status.message);
    };
    m_hm_texres->on_select = [this](Node*, int index) {
        int texres = get_texres();
        if (texres == m_texture.size().x)
            return;

#if defined(__GLES__)
        m_texture.create(texres, texres);
        m_texture.create_mipmaps();
#else
        // get the texture data and resize it
        Image img = m_texture.get_image();
        Image resized = img.resize(texres, texres);
        m_texture.create(resized);
        m_texture.create_mipmaps();
#endif // 
    };
    int rexres = get_texres();
    m_texture.create(rexres, rexres);
    m_sampler_mipmap.create();
    m_sampler_mipmap.set_filter(
        pp::renderer::gl::linear_mipmap_linear_texture_filter(),
        pp::renderer::gl::linear_texture_filter());
    m_sampler_linear.create();
    m_plane.create<1>(1, 1);

    TextureManager::load("data/sun.png");
}

int NodePanelGrid::get_samples() const
{
    return atoi(m_hm_samples->m_items[m_hm_samples->m_current_index].c_str());
}

int NodePanelGrid::get_texres() const
{
    return atoi(m_hm_texres->m_items[m_hm_texres->m_current_index].c_str());
}

float NodePanelGrid::get_radius() const
{
    return m_hm_radius->get_value() * 0.5f + 0.01f;
}

float NodePanelGrid::get_ambient() const
{
    return glm::pow(m_hm_ambient->get_value(), 3);
}

float NodePanelGrid::get_thickness() const
{
    return glm::lerp(m_line_range[0], m_line_range[1], m_hm_thickness->get_value());
}

float NodePanelGrid::get_resolution() const
{
    return atof(m_groud_resolution->m_items[m_groud_resolution->m_current_index].c_str());
}

float NodePanelGrid::get_height() const
{
    return -glm::pow(m_hm_height->get_value() - 0.5f, 3.f) * 10.f;
}

float NodePanelGrid::get_offset() const
{
    return glm::pow(m_groud_offset->get_value() - 0.5f, 3);
}

bool NodePanelGrid::load_heightmap_file(const std::string& path, bool raise_ground_opacity)
{
    const auto plan = raise_ground_opacity
        ? pp::app::plan_grid_heightmap_load(path)
        : pp::app::plan_grid_heightmap_reload(path);
    if (!plan)
        return false;

    const auto status = pp::panopainter::execute_legacy_grid_ui_plan(*this, plan.value());
    if (!status.ok())
        LOG("Grid heightmap load action failed: %s", status.message);
    return status.ok();
}

void NodePanelGrid::draw_heightmap(const glm::mat4& proj, const glm::mat4& camera, bool commit) const
{
    assert(App::I->is_render_thread());
    if (m_groud_opacity->get_value() > 0.f)
    {
        const bool depth = query_grid_capability(pp::renderer::gl::depth_test_state());
        apply_grid_capability(pp::renderer::gl::depth_test_state(), true);
        clear_grid_depth_buffer();

        auto nav = m_hm_image.m_data ? -(m_hm_preview_nav->m_value - 0.5f) : glm::vec2(0);
        auto mvp = proj * camera
            * glm::scale(glm::vec3(100))
            * glm::translate(glm::vec3(0, get_offset(), 0))
            * glm::translate(glm::vec3(nav.x, get_offset(), nav.y))
            ;

        auto light_yaw = m_hm_lyaw->get_value() * glm::pi<float>() * 2.f;
        auto light_pitch = m_hm_lpitch->get_value() * 5;
        auto light_pos = glm::vec3(sinf(light_yaw) + nav.x, light_pitch + get_offset(), cosf(light_yaw) + nav.y);
        auto light_dir = glm::normalize(light_pos);

        apply_grid_capability(pp::renderer::gl::blend_state(), false);
        // DRAW SOLID
        if (m_hm_image.m_data)
        {

            if (m_shade_mode == ShadeMode::Solid)
            {
                apply_grid_capability(pp::renderer::gl::blend_state(), false);
                pp::panopainter::configure_legacy_grid_lambert_shader(
                    mvp,
                    light_dir,
                    get_ambient());
                m_hm_plane.draw_fill();
            }
            else if (m_shade_mode == ShadeMode::Flat)
            {
                pp::panopainter::configure_legacy_ui_color_shader(
                    mvp,
                    glm::vec4(
                        glm::vec3(1.f - m_groud_value->get_value()),
                        m_groud_opacity->get_value()));
                m_hm_plane.draw_fill();
            }
            else if(m_shade_mode == ShadeMode::Textured)
            {
                pp::panopainter::configure_legacy_grid_lambert_lightmap_shader(
                    mvp,
                    light_dir,
                    get_ambient(),
                    0);
                m_sampler_mipmap.bind(0);
                set_grid_active_texture_unit(0U);
                m_texture.bind();
                m_hm_plane.draw_fill();
                m_texture.unbind();
            }
        }

        // DRAW GRIDS
        auto wire_alpha = m_hm_image.m_data ? m_hm_wireframe->get_value() : 1.f;
        if (wire_alpha > 0.f)
        {
            apply_grid_capability(pp::renderer::gl::blend_state(), true);
            pp::panopainter::configure_legacy_ui_color_shader(
                mvp,
                glm::vec4(
                    glm::vec3(m_groud_value->get_value()),
                    m_groud_opacity->get_value() *
                    wire_alpha));
            if (m_hm_image.m_data && m_shade_mode == ShadeMode::Transparent)
            {
                apply_grid_color_mask(
                    pp::renderer::gl::color_write_disabled(),
                    pp::renderer::gl::color_write_disabled(),
                    pp::renderer::gl::color_write_disabled(),
                    pp::renderer::gl::color_write_disabled());
                m_hm_plane.draw_fill();
                apply_grid_color_mask(
                    pp::renderer::gl::color_write_enabled(),
                    pp::renderer::gl::color_write_enabled(),
                    pp::renderer::gl::color_write_enabled(),
                    pp::renderer::gl::color_write_enabled());
            }
            m_hm_plane.draw_stroke();
        }

        // DRAW SUN SPHERE
        if (!commit)
        {
            auto c = proj * camera * glm::translate(light_pos) * glm::vec4(0, 0, 0, 1);
            if (c.w > 0)
            {
                auto p2d = xy(c) / c.z;
                const auto viewport = query_grid_viewport();
                auto aspect_ratio = static_cast<float>(viewport.height) / static_cast<float>(viewport.width);
                apply_grid_capability(pp::renderer::gl::blend_state(), true);
                apply_grid_capability(pp::renderer::gl::depth_test_state(), false);
                pp::panopainter::setup_legacy_canvas_draw_merge_texture_shader({
                    .mvp = glm::ortho(-1.f, 1.f, -1.f, 1.f) *
                        //glm::scale(glm::vec3(100)) *
                        glm::translate(glm::vec3(p2d, 0)) * glm::scale(glm::vec3(.1f * aspect_ratio, .1f, 1.f)),
                    .texture_slot = 0});
                set_grid_active_texture_unit(0U);
                m_sampler_linear.bind(0);
                constexpr auto sun_tex = const_hash("data/sun.png");
                TextureManager::get(sun_tex).bind();
                m_plane.draw_fill();
            }
        }
        depth ? apply_grid_capability(pp::renderer::gl::depth_test_state(), true) : apply_grid_capability(pp::renderer::gl::depth_test_state(), false);
    }
}

kEventResult NodePanelGrid::handle_event(Event* e)
{
    switch (e->m_type)
    {
    case kEventType::MouseUpL:
        if (!m_mouse_inside)
        {
            pp::panopainter::release_legacy_mouse_capture(*this);
            if (m_parent)
            {
                pp::panopainter::detach_legacy_node_from_parent(*this);
            }
            if (on_popup_close)
            {
                on_popup_close(this);
            }
        }
        break;
    default:
        return kEventResult::Available;
        break;
    }
    return kEventResult::Available;
}

void NodePanelGrid::bake_uvs()
{
    if (!m_hm_image.m_data)
        return;

    RTT fb;
    const auto renderer_features = ShaderManager::render_device_features();
    if (renderer_features.float32_render_targets)
    {
        fb.create(m_texture.size().x, m_texture.size().y, -1, pp::renderer::gl::rgba32f_internal_format());
    }
    else if (renderer_features.float16_render_targets)
    {
        fb.create(m_texture.size().x, m_texture.size().y, -1, pp::renderer::gl::rgba16f_internal_format());
    }

    std::unique_ptr<float[]> data_nor;
    std::unique_ptr<float[]> data_pos;
    App::I->render_task([&]{
        fb.bindFramebuffer();
        fb.clear({ 1, 0, 0, 1 });
        apply_grid_capability(pp::renderer::gl::blend_state(), false);
        apply_grid_capability(pp::renderer::gl::depth_test_state(), false);
        apply_grid_viewport(0, 0, fb.getWidth(), fb.getHeight());
        pp::panopainter::configure_legacy_grid_bake_uv_shader(glm::mat4(1));

        // bake normal
        ShaderManager::u_int(kShaderUniform::Mode, 0);
        m_hm_plane.draw_fill();
        data_nor.reset(fb.readTextureDataFloat());
        //stbi_write_png("bake-nor.png", fb.getWidth(), fb.getHeight(), 4, fb.readTextureData(), 0);

        // bake position
        ShaderManager::u_int(kShaderUniform::Mode, 1);
        m_hm_plane.draw_fill();
        data_pos.reset(fb.readTextureDataFloat());
        //stbi_write_png("bake-pos.png", fb.getWidth(), fb.getHeight(), 4, fb.readTextureData(), 0);

        fb.unbindFramebuffer();
        fb.destroy();
    });

    auto pb = App::I->show_progress("Lightmap Rendering");

    if (m_rt_dirty)
    {
        nanort::BVHBuildOptions<float> build_options;  // Use default option
        build_options.cache_bbox = false;
        nanort::TriangleMesh<float> triangle_mesh(reinterpret_cast<float*>(m_hm_plane.vertices.data()), m_hm_plane.idx.data(), sizeof(vertex_t));
        nanort::TriangleSAHPred<float> triangle_pred(reinterpret_cast<float*>(m_hm_plane.vertices.data()), m_hm_plane.idx.data(), sizeof(vertex_t));
        bool ret = m_rt_accel.Build(static_cast<unsigned int>(m_hm_plane.idx.size() / 3), triangle_mesh, triangle_pred, build_options);
        if (!ret)
            return;
        m_rt_dirty = false;
    }
    
    auto light_yaw = m_hm_lyaw->get_value() * glm::pi<float>() * 2.f;
    auto light_pitch = m_hm_lpitch->get_value() * 5;
    auto light_pos = glm::vec3(sinf(light_yaw), light_pitch + get_offset(), cosf(light_yaw));
    auto light_dir = glm::normalize(light_pos);

    std::atomic_int pb_value(0);

    auto data_out = std::make_unique<uint8_t[]>(fb.getWidth() * fb.getHeight() * 4);
    const auto samples = get_samples();
    const auto radius = get_radius();
    {
        std::jthread worker([&]
        {
            BT_SetTerminate();
            float* d_pos = data_pos.get();
            float* d_nor = data_nor.get();
            auto* d_out = reinterpret_cast<glm::u8vec4*>(data_out.get());
            parallel_for(fb.getHeight(), [&](size_t y)
            {
                for (int x = 0; x < fb.getWidth(); x++)
                {
                    int i = (int)y * fb.getHeight() + x;
                    auto nor = glm::make_vec3(&d_nor[i * 4]);
                    auto pos = glm::make_vec3(&d_pos[i * 4]);
                    auto& out = d_out[i];

                    if (glm::dot(nor, light_dir) <= 0.f)
                    {
                        out = { 50, 50, 50, 255 };
                        continue;
                    }

                    int hit = 0;
                    for (int s = 0; s < samples; s++)
                    {
                        auto dir = glm::normalize(light_dir + glm::sphericalRand(radius));

                        nanort::Ray<float> ray;
                        ray.org[0] = pos.x;// + nor.x * 0.005;
                        ray.org[1] = pos.y;// + nor.y * 0.005;
                        ray.org[2] = pos.z;// + nor.z * 0.005;
                        ray.dir[0] = dir.x;
                        ray.dir[1] = dir.y;
                        ray.dir[2] = dir.z;

                        float kFar = 2000.0;
                        ray.min_t = 0.005f;
                        ray.max_t = kFar;

                        nanort::TriangleIntersector<> triangle_intersector(reinterpret_cast<float*>(m_hm_plane.vertices.data()), m_hm_plane.idx.data(), sizeof(vertex_t));
                        nanort::TriangleIntersection<> isect;
                        hit += m_rt_accel.Traverse(ray, triangle_intersector, &isect);
                    }
                    out = glm::lerp(glm::vec4(255, 255, 255, 255), glm::vec4(50, 50, 50, 255), hit / (float)samples);
                }
                pb_value++;
            });
        });
        while (pb_value < fb.getHeight())
        {
            pb->set_progress((float)pb_value / (float)fb.getHeight());
            std::this_thread::sleep_for(std::chrono::milliseconds(100));
        }
    }
    pp::panopainter::close_legacy_dialog_node(*pb);
    //stbi_write_jpg("bake-out.jpg", fb.getWidth(), fb.getHeight(), 4, data_out.get(), 75);
    m_texture.update(data_out.get());
    m_texture.create_mipmaps();
    App::I->async_redraw();
}

///////////////////////////////////////////////////////////////////////////////

Node* NodeHeightmapOverlay::clone_instantiate() const
{
    return new NodeHeightmapOverlay();
}

void NodeHeightmapOverlay::clone_finalize(Node* dest) const
{
    auto n = (NodeHeightmapOverlay*)dest;
    n->m_picker = (NodeBorder*)n->m_children[0].get();
}

void NodeHeightmapOverlay::init()
{
    m_picker = new NodeBorder;
    m_picker->SetSize({ 10, 10 });
    m_picker->SetPositioning(YGPositionTypeAbsolute);
    m_picker->SetPosition(0, 0);
    m_picker->m_border_color = glm::vec4(0, 0, 0, 1);
    m_picker->m_thinkness = 1;
    m_picker->m_color = glm::vec4(1, 1, 1, 0.25f);
    add_child(m_picker);
}

void NodeHeightmapOverlay::set_value(float x, float y)
{
    auto sz = m_size;
    auto pos = glm::clamp(glm::vec2(x, y) * sz, { 0, 0 }, sz);
    //m_picker->SetPosition(pos - m_picker->GetSize() * .5f);
    m_value = pos / glm::max({ 1,1 }, sz); // avoid div0
    if (on_value_changed)
        on_value_changed(this, m_value);
}

kEventResult NodeHeightmapOverlay::handle_event(Event* e)
{
    NodeBorder::handle_event(e);
    switch (e->m_type)
    {
    case kEventType::MouseDownL:
    {
        m_old_value = m_value;
        dragging = true;
        mouse_capture();
        auto sz = m_size;
        auto pos = glm::clamp(((MouseEvent*)e)->m_pos - m_pos, { 0, 0 }, sz);
        m_picker->SetPosition(pos - m_picker->GetSize() * .5f);
        m_value = pos / glm::max({ 1,1 }, sz); // avoid div0
        if (on_value_changed)
            on_value_changed(this, m_value);
    }
    break;
    case kEventType::MouseUpL:
        pp::panopainter::release_legacy_mouse_capture(*this);
        dragging = false;
        break;
    case kEventType::MouseMove:
        if (dragging)
        {
            auto sz = m_size;
            auto pos = glm::clamp(((MouseEvent*)e)->m_pos - m_pos, { 0, 0 }, sz);
            m_picker->SetPosition(pos - m_picker->GetSize() * .5f);
            m_value = pos / glm::max({ 1,1 }, sz); // avoid div0
            if (on_value_changed)
                on_value_changed(this, m_value);
        }
        break;
    case kEventType::MouseCancel:
        pp::panopainter::release_legacy_mouse_capture(*this);
        dragging = false;
        m_value = m_old_value;
        set_value(m_value.x, m_value.y);
        if (on_value_changed)
            on_value_changed(this, m_value);
        break;
    default:
        return kEventResult::Available;
        break;
    }
    return kEventResult::Consumed;
}

void NodeHeightmapOverlay::draw()
{
    auto sz = m_size;
    auto pos = glm::clamp(m_value * sz, { 0, 0 }, sz);
    m_picker->SetPosition(pos - m_picker->GetSize() * .5f);
}

void NodeHeightmapOverlay::handle_resize(glm::vec2 old_size, glm::vec2 new_size, float zoom)
{
    NodeBorder::handle_resize(old_size, new_size, zoom);
    set_value(m_value.x, m_value.y);
}
