#include "pch.h"
#include "log.h"
#include "node_panel_grid.h"
#include "canvas.h"
#include "app.h"
#include "image.h"
#include "util.h"

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*) {
        App::I->pick_image([this](std::string path) {
            Image img;
            if (img.load_file(path))
            {
                m_file_path = path;
                m_hm_image = img.resize(128, 128);
                m_hm_preview->tex.create(m_hm_image);
                m_hm_preview->tex.create_mipmaps();
                auto sz = m_hm_preview->tex.size();
                m_hm_preview->SetAspectRatio(sz.x / sz.y);
                m_hm_plane.create(1, 1, m_hm_image, get_resolution(), get_height());
                m_hm_preview->SetHeight(100);
                if (m_groud_opacity->get_value() == 0.f)
                    m_groud_opacity->set_value(1.f);
                m_rt_dirty = true;
            }
        });
    };

    m_hm_clear->on_click = [this](Node*)
    {
        m_hm_plane.create(1, 1, 100 * get_resolution());
        m_hm_image.destroy();
        m_hm_preview->tex.destroy();
        m_hm_preview->SetHeight(0);
    };

    m_hm_reload->on_click = [this](Node*)
    {
        Image img;
        if (img.load_file(m_file_path))
        {
            m_hm_image = img.resize(128, 128);
            m_hm_preview->tex.create(m_hm_image);
            m_hm_preview->tex.create_mipmaps();
            auto sz = m_hm_preview->tex.size();
            m_hm_preview->SetAspectRatio(sz.x / sz.y);
            m_hm_plane.create(1, 1, m_hm_image, get_resolution(), get_height());
            m_hm_preview->SetHeight(100);
            m_rt_dirty = true;
        }
    };

    m_render->on_click = [this](Node*)
    {
        if (ShaderManager::ext_float32 || ShaderManager::ext_float16)
        {
            std::thread([this] {
                bake_uvs();
                m_hm_shading->set_index(3);
                m_shade_mode = ShadeMode::Textured;
            }).detach();
        }
        else
        {
            App::I->message_box("Rendering failed", 
                "Your hardware does not support lightmap rendering.");
        }
    };
    m_commit->on_click = [this](Node*)
    {
        Canvas::I->draw_objects([this](const glm::mat4& camera, const glm::mat4& proj, int i) {
            draw_heightmap(proj, camera, true);
        });
        m_groud_opacity->set_value(0);
    };
    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;
        img.create(m_texture.size().x, m_texture.size().y);
        App::I->render_task([&]
        {
            m_texture.bind();
            glGetTexImage(GL_TEXTURE_2D, 0, GL_RGBA, GL_UNSIGNED_BYTE, img.m_data.get());
            m_texture.unbind();
        });
        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(GL_LINEAR_MIPMAP_LINEAR, GL_LINEAR);
    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);
}

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)
    {
        bool depth = glIsEnabled(GL_DEPTH_TEST);
        glEnable(GL_DEPTH_TEST);
        glClear(GL_DEPTH_BUFFER_BIT);

        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);

        glDisable(GL_BLEND);
        // DRAW SOLID
        if (m_hm_image.m_data)
        {

            if (m_shade_mode == ShadeMode::Solid)
            {
                glDisable(GL_BLEND);
                ShaderManager::use(kShader::Lambert);
                ShaderManager::u_mat4(kShaderUniform::MVP, mvp);
                ShaderManager::u_vec3(kShaderUniform::LightDir, light_dir);
                ShaderManager::u_float(kShaderUniform::Ambient, get_ambient());
                m_hm_plane.draw_fill();
            }
            else if (m_shade_mode == ShadeMode::Flat)
            {
                ShaderManager::use(kShader::Color);
                ShaderManager::u_vec4(kShaderUniform::Col, glm::vec4(
                    glm::vec3(1.f - m_groud_value->get_value()),
                    m_groud_opacity->get_value()
                ));
                ShaderManager::u_mat4(kShaderUniform::MVP, mvp);
                m_hm_plane.draw_fill();
            }
            else if(m_shade_mode == ShadeMode::Textured)
            {
                ShaderManager::use(kShader::LambertLightmap);
                ShaderManager::u_mat4(kShaderUniform::MVP, mvp);
                ShaderManager::u_vec3(kShaderUniform::LightDir, light_dir);
                ShaderManager::u_float(kShaderUniform::Ambient, get_ambient());
                ShaderManager::u_int(kShaderUniform::Tex, 0);
                m_sampler_mipmap.bind(0);
                glActiveTexture(GL_TEXTURE0);
                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)
        {
            glEnable(GL_BLEND);
            ShaderManager::use(kShader::Color);
            ShaderManager::u_vec4(kShaderUniform::Col, glm::vec4(
                glm::vec3(m_groud_value->get_value()),
                m_groud_opacity->get_value() *
                wire_alpha
            ));
            ShaderManager::u_mat4(kShaderUniform::MVP, mvp);
            if (m_hm_image.m_data && m_shade_mode == ShadeMode::Transparent)
            {
                glColorMask(GL_FALSE, GL_FALSE, GL_FALSE, GL_FALSE);
                m_hm_plane.draw_fill();
                glColorMask(GL_TRUE, GL_TRUE, GL_TRUE, GL_TRUE);
            }
            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;
                GLint vp[4];
                glGetIntegerv(GL_VIEWPORT, vp);
                auto aspect_ratio = (float)vp[3] / (float)vp[2];
                glEnable(GL_BLEND);
                glDisable(GL_DEPTH_TEST);
                ShaderManager::use(kShader::Texture);
                ShaderManager::u_int(kShaderUniform::Tex, 0);
                ShaderManager::u_mat4(kShaderUniform::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)));
                glActiveTexture(GL_TEXTURE0);
                m_sampler_linear.bind(0);
                constexpr auto sun_tex = const_hash("data/sun.png");
                TextureManager::get(sun_tex).bind();
                m_plane.draw_fill();
            }
        }
        depth ? glEnable(GL_DEPTH_TEST) : glDisable(GL_DEPTH_TEST);
    }
}

kEventResult NodePanelGrid::handle_event(Event* e)
{
    switch (e->m_type)
    {
    case kEventType::MouseUpL:
        if (!m_mouse_inside)
        {
            mouse_release();
            m_parent->remove_child(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;
    if (ShaderManager::ext_float32)
    {
        fb.create(m_texture.size().x, m_texture.size().y, -1, GL_RGBA32F);
    }
    else if (ShaderManager::ext_float16)
    {
        fb.create(m_texture.size().x, m_texture.size().y, -1, GL_RGBA16F);
    }

    std::unique_ptr<float[]> data_nor;
    std::unique_ptr<float[]> data_pos;
    App::I->render_task([&]{
        fb.bindFramebuffer();
        fb.clear({ 1, 0, 0, 1 });
        glDisable(GL_BLEND);
        glDisable(GL_DEPTH_TEST);
        glViewport(0, 0, fb.getWidth(), fb.getHeight());
        ShaderManager::use(kShader::BakeUV);
        ShaderManager::u_mat4(kShaderUniform::MVP, 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::thread worker([&]
    {
        BT_SetTerminate();
        __block float* d_pos = data_pos.get();
        __block float* d_nor = data_nor.get();
        __block glm::u8vec4* d_out = reinterpret_cast<glm::u8vec4*>(data_out.get());
#if _WIN32
        concurrency::parallel_for(int(0), fb.getHeight(), [&](int y)
#elif __IOS__ || __OSX__
        dispatch_queue_t queue = dispatch_get_global_queue(DISPATCH_QUEUE_PRIORITY_DEFAULT, 0);
        dispatch_apply((size_t)fb.getHeight(), queue, ^ (size_t y)
#else
        parallel_for(fb.getHeight(), [&](int y)
#endif
        {
            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));
    }
    worker.join();
    pb->destroy();
    //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:
        mouse_release();
        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:
        mouse_release();
        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);
}