#include "pch.h"
#include "canvas_layer.h"
#include "app.h"
#include "legacy_canvas_draw_merge_services.h"
#include "legacy_ui_gl_dispatch.h"
#include "log.h"
#include "renderer_gl/opengl_capabilities.h"
#include "rtt.h"
#include "util.h"

#include <array>

uint32_t Layer::s_count = 0;

namespace {

void apply_layer_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, "Layer");
}

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

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

void bind_layer_texture_cube(std::uint32_t texture_id)
{
    pp::legacy::ui_gl::bind_texture_cube(texture_id, "Layer");
}

void copy_layer_framebuffer_to_texture(
    std::uint32_t texture_target,
    std::int32_t destination_x,
    std::int32_t destination_y,
    std::int32_t source_x,
    std::int32_t source_y,
    std::int32_t width,
    std::int32_t height)
{
    copy_framebuffer_to_texture_target(
        texture_target,
        destination_x,
        destination_y,
        source_x,
        source_y,
        width,
        height);
}

void clear_layer_color_buffer(const glm::vec4& color)
{
    pp::legacy::ui_gl::clear_color_buffer({ color.r, color.g, color.b, color.a }, "Layer");
}

std::array<float, 4> query_layer_clear_color()
{
    return pp::legacy::ui_gl::query_clear_color("Layer");
}

void restore_layer_clear_color(std::array<float, 4> color)
{
    pp::legacy::ui_gl::set_clear_color(color, "Layer");
}

}

RTT& Layer::rtt(int i, int frame /*= -1*/)
{
    if (frame == -1)
        frame = m_frame_index;
    assert(m_frames[frame].m_rtt[i].valid());
    return m_frames[frame].m_rtt[i];
}

glm::vec4& Layer::box(int i, int frame /*= -1*/)
{
    if (frame == -1)
        frame = m_frame_index;
    return m_frames[frame].m_dirty_box[i];
}

bool& Layer::face(int i, int frame /*= -1*/)
{
    if (frame == -1)
        frame = m_frame_index;
    return m_frames[frame].m_dirty_face[i];
}

LayerFrame& Layer::frame(int frame /*= -1*/)
{
    if (frame == -1)
        frame = m_frame_index;
    return m_frames[frame];
}

TextureCube Layer::gen_cube()
{
    TextureCube ret;
    ret.create(w);
    for (int i = 0; i < 6; i++)
    {
        App::I->render_task([&]
        {
            ret.bind();
            rtt(i).bindFramebuffer();
            copy_layer_framebuffer_to_texture(pp::renderer::gl::cube_face_texture_target(i), 0, 0, 0, 0, w, w);
            rtt(i).unbindFramebuffer();
        });
    }
    return ret;
}

Texture2D Layer::gen_equirect(glm::ivec2 size /*= { 0, 0 }*/)
{
    Texture2D ret;
    
    App::I->render_task([&]
    {
        gl_state gl;
        gl.save();

        TextureCube cube;
        RTT latlong;
        
        if (size.x == 0 || size.y == 0)
            size = { w, h };

        cube = gen_cube();
        latlong.create(size.x * 4, size.y * 2);
        ret.create(size.x * 4, size.y * 2);

        apply_layer_capability(pp::renderer::gl::blend_state(), false);

        latlong.bindFramebuffer();

        //latlong.clear({ 0, 1, 1, 1 });

        apply_layer_viewport(0, 0, latlong.getWidth(), latlong.getHeight());

        set_layer_active_texture_unit(0U);
        bind_layer_texture_cube(cube.m_cubetex_id);

        pp::panopainter::setup_legacy_canvas_draw_merge_equirect_shader(
            {.mvp = glm::ortho(-.5f, .5f, -.5f, .5f, -1.f, 1.f), .texture_slot = 0});
        Canvas::I->m_sampler.bind(0);
        Canvas::I->m_plane.draw_fill();

        ret.bind();
        copy_framebuffer_to_texture_2d(0, 0, 0, 0, latlong.getWidth(), latlong.getHeight());

        latlong.unbindFramebuffer();

        latlong.destroy();
        cube.destroy();

        gl.restore();
    });

    return ret;
}

PBO Layer::gen_equirect_pbo(glm::ivec2 size /*= { 0, 0 }*/)
{
    if (size.x == 0 || size.y == 0)
        size = { w, h };

    PBO pbo;
    TextureCube cube = gen_cube();
    std::this_thread::yield();
    RTT latlong;
    latlong.create(size.x, size.y);
    std::this_thread::yield();

    App::I->render_task([&]
    {
        apply_layer_capability(pp::renderer::gl::blend_state(), false);

        latlong.bindFramebuffer();
        
        apply_layer_viewport(0, 0, latlong.getWidth(), latlong.getHeight());
        set_layer_active_texture_unit(0U);
        bind_layer_texture_cube(cube.m_cubetex_id);

        pp::panopainter::setup_legacy_canvas_draw_merge_equirect_shader(
            {.mvp = glm::ortho(-.5f, .5f, -.5f, .5f, -1.f, 1.f), .texture_slot = 0});
        Canvas::I->m_sampler.bind(0);
        Canvas::I->m_plane.draw_fill();

        latlong.unbindFramebuffer();
    });
    std::this_thread::yield();
    pbo.create(latlong);
    std::this_thread::yield();
    latlong.destroy();
    cube.destroy();
    return pbo;
}

void Layer::destroy()
{
    for (int i = 0; i < 6; i++)
        rtt(i).destroy();
}

void Layer::optimize(int frame /*= -1*/)
{
    if (frame == -1)
        frame = m_frame_index;
    int saved_bytes = 0;
    for (int i = 0; i < 6; i++)
    {
        if (!face(i, frame))
            continue;

        auto data = std::unique_ptr<glm::u8vec4[]>(reinterpret_cast<glm::u8vec4*>(rtt(i, frame).readTextureData()));
        glm::ivec2 bbmin(w, h);
        glm::ivec2 bbmax(0);
        for (int y = 0; y < h; y++)
        {
            for (int x = 0; x < w; x++)
            {
                if (data[x + y * w].a > 0)
                {
                    bbmin = glm::min(bbmin, { x, y });
                    bbmax = glm::max(bbmax, { x + 1, y + 1 });
                }
            }
        }
        glm::vec2 bbsz = bbmax - bbmin;
        glm::vec2 old_size = zw(box(i, frame)) - xy(box(i, frame));
        glm::vec2 diff;
        if (bbsz.x <= 0 || bbmax.y <= 0)
        {
            face(i, frame) = false;
            box(i, frame) = glm::vec4(0);
            diff = old_size;
        }
        else
        {
            box(i, frame) = { bbmin, bbmax };

            diff = old_size - bbsz;
        }
        saved_bytes += (int)(diff.x * diff.y * 4);
    }
    LOG("optimized %d bytes", saved_bytes);
}

void Layer::restore(const LayerFrame::Snapshot& snap, int frame /*= -1*/)
{
    if (frame == -1)
        frame = m_frame_index;
    auto loaded = m_frames[frame].gpu_load();
    m_frames[frame].restore(snap);
    if (!loaded) m_frames[frame].gpu_unload();
}

LayerFrame::Snapshot Layer::snapshot(int frame /*= -1*/, std::array<glm::vec4, 6>* dirty_box /*= nullptr*/, std::array<bool, 6>* dirty_face /*= nullptr*/)
{
    if (frame == -1)
        frame = m_frame_index;
    auto loaded = m_frames[frame].gpu_load();
    auto snap = m_frames[frame].snapshot(dirty_box, dirty_face);
    if (!loaded) m_frames[frame].gpu_unload();
    return snap;
}

void Layer::clear(const glm::vec4& c, int frame /*= -1*/)
{
    if (frame == -1)
        frame = m_frame_index;
    m_frames[frame].clear(c);
}

bool Layer::create(int width, int height, std::string name)
{
    m_name = name;
    w = width;
    h = height;
    m_frame_index = 0;
    m_frames.clear();
    m_frames.emplace_back();
    m_frames.back().create(width, height);
    frames_gpu_update();
    return true;
}

bool Layer::add_frame(int index /*= -1*/)
{
    auto pos = index == -1 ? m_frames.end() : m_frames.begin() + index;
    auto it = m_frames.emplace(pos);
    it->create(w, h);
    frames_gpu_update();
    return true;
}

void Layer::remove_frame(int frame)
{
    m_frames.erase(m_frames.begin() + frame);
    m_frame_index = glm::clamp(m_frame_index, 0, (int)m_frames.size() - 1);
    frames_gpu_update();
}

void Layer::duplicate_frame(int frame)
{
    if (frame == -1)
        frame = m_frame_index;
    m_frames.insert(m_frames.begin() + frame + 1, m_frames[frame].clone());
    frames_gpu_update();
}

void Layer::frames_gpu_update()
{
    App::I->render_task_async([=]
    {
        for (int j = 0; j < m_frames.size(); j++)
        {
            int dist = glm::abs(j - m_frame_index);
            int onion = App::I->animation ? App::I->animation->get_onion_size() : 1;
            if (dist <= onion)
                m_frames[j].gpu_load();
            else
                m_frames[j].gpu_unload();
        }
    });
}

int Layer::move_frame_offset(int frame, int offset) noexcept
{
    int new_pos = glm::clamp(frame + offset, 0, (int)m_frames.size() - 1);
    auto from = m_frames.begin() + frame;
    auto to = m_frames.begin() + new_pos;

    if (new_pos < frame)
        std::rotate(to, from, from + 1);
    if (new_pos > frame)
        std::rotate(from, from + 1, to + 1);

    frames_gpu_update();
    return new_pos;
}

int Layer::total_duration() const noexcept
{
    int duration = 0;
    for (auto& f : m_frames)
        duration += f.m_duration;
    return duration;
}

void Layer::goto_frame(int frame) noexcept
{
    int i = 0;
    for (i = 0; i < m_frames.size() && frame >= 0; i++)
        frame -= m_frames[i].m_duration;
    m_frame_index = i - 1;
    m_frames[m_frame_index].gpu_load();
    frames_gpu_update();
}

void Layer::resize(int width, int height)
{
    w = width;
    h = height;
    for (auto& frame : m_frames)
        frame.resize(width, height);
}

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

void LayerFrame::Snapshot::create(int w, int h)
{
    width = w;
    height = h;
    for (int i = 0; i < 6; i++)
    {
        m_dirty_face[i] = false;
        m_dirty_box[i] = glm::vec4(width, height, 0, 0);
        image[i] = std::make_unique<uint8_t[]>(w * h * 4);
        std::fill_n(image[i].get(), w * h * 4, 0);
    }
}

void LayerFrame::Snapshot::clear()
{
    for (int i = 0; i < 6; i++)
    {
        m_dirty_face[i] = false;
        m_dirty_box[i] = glm::vec4(width, height, 0, 0);
        std::fill_n(image[i].get(), width * height * 4, 0);
    }
}

void LayerFrame::Snapshot::optimize()
{
    for (int i = 0; i < 6; i++)
    {
        if (!m_dirty_face[i] || !image[i])
            continue;
        auto data = reinterpret_cast<glm::u8vec4*>(image[i].get());
        glm::ivec2 bbmin(width, height);
        glm::ivec2 bbmax(0);
        for (int y = 0; y < height; y++)
        {
            for (int x = 0; x < width; x++)
            {
                if (data[x + y * width].a > 0)
                {
                    bbmin = glm::min(bbmin, { x, y });
                    bbmax = glm::max(bbmax, { x + 1, y + 1 });
                }
            }
        }
        //glm::vec2 bbsz = bbmax - bbmin;
    }
}

int LayerFrame::Snapshot::memsize() const
{
    int ret = 0;
    for (int i = 0; i < 6; i++)
    {
        if (m_dirty_face[i])
        {
            glm::vec2 sz = zw(m_dirty_box[i]) - xy(m_dirty_box[i]);
            ret += sz.x * sz.y * 4;
        }
    }
    return ret;
}

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

/*
LayerFrame::LayerFrame(LayerFrame&& other)
{
    LOG("LayerFrame move-ctor");
    for (int i = 0; i < 6; i++)
    {
        m_rtt[i] = std::move(other.m_rtt[i]);
        m_dirty_box[i] = other.m_dirty_box[i];
        m_dirty_face[i] = other.m_dirty_face[i];
    }
    m_duration = other.m_duration;
    w = other.w;
    h = other.h;
}

LayerFrame& LayerFrame::operator=(LayerFrame&& other)
{
    LOG("LayerFrame move-assignment");
    for (int i = 0; i < 6; i++)
    {
        m_rtt[i] = std::move(other.m_rtt[i]);
        m_dirty_box[i] = other.m_dirty_box[i];
        m_dirty_face[i] = other.m_dirty_face[i];
    }
    m_duration = other.m_duration;
    w = other.w;
    h = other.h;
    return *this;
}
*/

bool LayerFrame::create(int width, int height, int duration /*= 1*/)
{
    bool success = true;
    //App::I->render_task([&]
    //{
    //    for (int i = 0; i < 6; i++)
    //    {
    //        if (!m_rtt[i].create(width, height))
    //        {
    //            success = false;
    //            return;
    //        }
    //        m_rtt[i].bindFramebuffer();
    //        m_rtt[i].clear();
    //        m_rtt[i].unbindFramebuffer();
    //        m_dirty_box[i] = glm::vec4(width, height, 0, 0); // reset bounding box
    //        m_dirty_face[i] = false;
    //    }
    //});
    for (int i = 0; i < 6; i++)
    {
        m_dirty_box[i] = glm::vec4(width, height, 0, 0); // reset bounding box
        m_dirty_face[i] = false;
    }
    m_gpu_data = nullptr;
    m_duration = duration;
    w = width;
    h = height;
    return success;
}

bool LayerFrame::resize(int width, int height)
{
    bool loaded = gpu_load();
    glm::vec2 ratio = glm::vec2(width, height) / glm::vec2(w, h);
    for (int i = 0; i < 6; i++)
    {
        if (!m_rtt[i].resize(width, height))
            return false;
        m_dirty_box[i] = m_dirty_box[i] * glm::vec4(ratio, ratio);
    }
    w = width;
    h = height;
    if (!loaded) gpu_unload();
    return true;
}

void LayerFrame::clear(const glm::vec4& c)
{
    App::I->render_task([&]
    {
        // push clear color state
        const auto cc = query_layer_clear_color();

        bool erase = (c.a == 0.f);

        for (int i = 0; i < 6; i++)
        {
            m_rtt[i].bindFramebuffer();
            clear_layer_color_buffer(c);
            m_rtt[i].unbindFramebuffer();

            if (erase)
            {
                m_dirty_box[i] = glm::vec4(w, h, 0, 0); // reset bounding box
                m_dirty_face[i] = false;
            }
            else
            {
                m_dirty_box[i] = glm::vec4(0, 0, w, h); // reset bounding box
                m_dirty_face[i] = true;
            }
        }

        // restore clear color state
        restore_layer_clear_color(cc);
    });
}

LayerFrame LayerFrame::clone() noexcept
{
    LayerFrame dup;
    dup.m_duration = m_duration;
    dup.w = w;
    dup.h = h;
    bool loaded = gpu_load();
    for (int i = 0; i < 6; i++)
    {
        dup.m_rtt[i] = m_rtt[i].clone();
        dup.m_dirty_box[i] = m_dirty_box[i];
        dup.m_dirty_face[i] = m_dirty_face[i];
    }
    if (!loaded) gpu_unload();
    return dup;
}

void LayerFrame::restore(const Snapshot& snap)
{
    App::I->render_task([this, &snap]
    {
        clear({ 0, 0, 0, 0 });
        for (int i = 0; i < 6; i++)
        {
            if (!m_rtt[i].valid())
                m_rtt[i].create(w, h);

            if (snap.image[i] == nullptr || snap.m_dirty_face[i] == false || box_area(snap.m_dirty_box[i]) <= 0)
            {
                m_dirty_box[i] = glm::vec4(snap.width, snap.height, 0, 0);
                m_dirty_face[i] = false;
                continue;
            }

            m_dirty_box[i] = snap.m_dirty_box[i];
            m_dirty_face[i] = snap.m_dirty_face[i];

            // TODO: this should not be recreated here!
            // Sorry I messed up with this,
            // it's just a quick fix DON'T SHIP!!
            //m_rtt[i].recreate();

            glm::vec2 box_sz = zw(m_dirty_box[i]) - xy(m_dirty_box[i]);
            m_rtt[i].updateRgba8(
                static_cast<int>(m_dirty_box[i].x), static_cast<int>(m_dirty_box[i].y),
                static_cast<int>(box_sz.x), static_cast<int>(box_sz.y),
                snap.image[i].get());
            LOG("restore face %d - %d bytes (%dx%d)", i,
                (int)box_sz.x * (int)box_sz.y * 4, (int)box_sz.x, (int)box_sz.y);
        }
    });
}

LayerFrame::Snapshot LayerFrame::snapshot(std::array<glm::vec4, 6>* dirty_box /*= nullptr*/, std::array<bool, 6>* dirty_face /*= nullptr*/)
{
    Snapshot snap;
    snap.width = w;
    snap.height = h;
    App::I->render_task([this, &snap, dirty_face, dirty_box]
    {
        for (int i = 0; i < 6; i++)
        {
            snap.m_dirty_box[i] = dirty_box ? dirty_box->at(i) : m_dirty_box[i];
            snap.m_dirty_face[i] = dirty_face ? dirty_face->at(i) : m_dirty_face[i];

            if (!snap.m_dirty_face[i] || !m_rtt[i].valid())
                continue;

            snap.image[i] = std::make_unique<uint8_t[]>(m_rtt[i].bytes());

            glm::vec2 box_sz = zw(snap.m_dirty_box[i]) - xy(snap.m_dirty_box[i]);
            m_rtt[i].readPixelsRgba8(
                static_cast<int>(snap.m_dirty_box[i].x),
                static_cast<int>(snap.m_dirty_box[i].y),
                static_cast<int>(box_sz.x), static_cast<int>(box_sz.y),
                snap.image[i].get());
        }
    });
    return snap;
}

bool LayerFrame::gpu_load() noexcept
{
    if (m_gpu_data)
    {
        restore(*m_gpu_data);
    }
    else
    {
        for (auto& rtt : m_rtt)
            if (!rtt.valid())
                rtt.create(w, h);
    }
    m_gpu_data.reset();
    return true;
}

bool LayerFrame::gpu_unload() noexcept
{
    bool ret  = false;  // already unloaded from gpu
    if (!m_gpu_data)
    {
        m_gpu_data = std::make_unique<Snapshot>(snapshot());
        ret = true; // previous state was loaded on gpu
    }
    for (auto& rtt : m_rtt)
        rtt.destroy();
    return ret;
}
