#include "pch.h"
#include "log.h"
#include "node_panel_grid.h"
#include "canvas.h"
#include "app.h"
#include "image.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("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<NodeSliderH>("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_height = find<NodeSliderH>("grid-heightmap-height");
    m_hm_lyaw = find<NodeSliderH>("grid-heightmap-lyaw");
    m_hm_lpitch = find<NodeSliderH>("grid-heightmap-lpitch");
    m_hm_shading = find<NodeComboBox>("grid-heightmap-shading");
    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_height->on_value_changed = update_hm;
    m_groud_resolution->on_value_final = [this](Node* target, float v) {
        if (m_hm_image.data())
            m_hm_plane.create(1, 1, m_hm_image, v * 5.f, get_height());
        else
            m_hm_plane.create(1, 1, 100 * v * 5.f);
        m_rt_dirty = true;
        LOG("resolution value %f", v);
    };

    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, m_groud_resolution->get_value() * 5.f, 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;
            async_start();
            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, 
                    m_groud_resolution->get_value() * 5.f, 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;
            }
            async_update();
            async_end();
        });
    };

    m_hm_clear->on_click = [this](Node*)
    {
        m_hm_plane.create(1, 1, 100 * m_groud_resolution->get_value() * 5.f);
        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, 
                m_groud_resolution->get_value() * 5.f, get_height());
            m_hm_preview->SetHeight(100);
            m_rt_dirty = true;
        }
    };

    m_render->on_click = [this](Node*)
    {
        gl_state gl;
        gl.save();
        bake_uvs();
        m_hm_shading->set_index(3);
        m_shade_mode = ShadeMode::Textured;
        gl.restore();
    };
    m_commit->on_click = [this](Node*)
    {
        gl_state gl;
        gl.save();
        Canvas::I->draw_objects([this](const glm::mat4& camera, const glm::mat4& proj, int i) {
            draw_heightmap(proj, camera);
        });
        m_groud_opacity->set_value(0);
        gl.restore();
    };
    m_texture.create(1024, 1024);
    m_sampler_linear.create();
}

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) const
{
    if (m_groud_opacity->get_value() > 0.f)
    {
        glEnable(GL_DEPTH_TEST);
        glClear(GL_DEPTH_BUFFER_BIT);

        auto mvp = proj * camera
            * glm::translate(glm::vec3(0, get_offset(), 0));

        // DRAW SOLID
        if (m_hm_image.m_data)
        {
            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, cosf(light_yaw));
            auto light_dir = glm::normalize(light_pos);

            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);
                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_int(kShaderUniform::Tex, 0);
                m_sampler_linear.bind(0);
                glActiveTexture(GL_TEXTURE0);
                m_texture.bind();
                m_hm_plane.draw_fill();
            }
        }

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

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

    RTT fb;
    fb.create(m_texture.size().x, m_texture.size().y, -1, GL_RGBA32F);
    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();
    std::unique_ptr<float[]> data_nor(fb.readTextureDataFloat());
    //stbi_write_jpg("bake-nor.jpg", fb.getWidth(), fb.getHeight(), 4, fb.readTextureData(), 75);
    
    // bake position
    ShaderManager::u_int(kShaderUniform::Mode, 1);
    m_hm_plane.draw_fill();
    std::unique_ptr<float[]> data_pos(fb.readTextureDataFloat());
    //stbi_write_jpg("bake-pos.jpg", fb.getWidth(), fb.getHeight(), 4, fb.readTextureData(), 75);

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

    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, cosf(light_yaw));
    auto light_dir = glm::normalize(light_pos);

    auto data_out = std::make_unique<uint8_t[]>(fb.getWidth() * fb.getHeight() * 4);
    for (int y = 0; y < fb.getHeight(); y++)
    {
        for (int x = 0; x < fb.getWidth(); x++)
        {
            int i = y * fb.getHeight() + x;
            auto nor = glm::make_vec4(&data_nor[i * 4]);
            auto pos = glm::make_vec4(&data_pos[i * 4]);
            auto& out = *reinterpret_cast<glm::i8vec4*>(&data_out[i * 4]);

            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] = light_dir.x;
            ray.dir[1] = light_dir.y;
            ray.dir[2] = light_dir.z;

            float kFar = 1.0e+30f;
            ray.min_t = 0.001f;
            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;
            bool hit = m_rt_accel.Traverse(ray, triangle_intersector, &isect);
            if (hit) 
            {
                out = { 50, 50, 50, 255 };
            }
            else
            {
                out = { 255, 255, 255, 255 };
            }
        }
    }
    //stbi_write_jpg("bake-out.jpg", fb.getWidth(), fb.getHeight(), 4, data_out.get(), 75);
    m_texture.update(data_out.get());
}