#include "pch.h"
#include "canvas.h"
#include "app.h"
#include "legacy_canvas_stroke_services.h"
#include "legacy_canvas_stroke_composite_services.h"
#include "legacy_canvas_stroke_edge_services.h"
#include "legacy_canvas_stroke_execution_services.h"
#include "legacy_canvas_stroke_shader_services.h"
#include "legacy_ui_gl_dispatch.h"
#include "renderer_gl/opengl_capabilities.h"
#include "util.h"

#include <algorithm>
#include <array>
#include <cstdint>
#include <utility>

namespace {

GLenum blend_state()
{
    return static_cast<GLenum>(pp::renderer::gl::blend_state());
}

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

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

void apply_canvas_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, "Canvas");
}

pp::renderer::gl::OpenGlViewportRect query_canvas_viewport()
{
    return pp::legacy::ui_gl::query_viewport_rect("Canvas");
}

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

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

void apply_canvas_capability(std::uint32_t capability, bool enabled)
{
    pp::legacy::ui_gl::set_capability(capability, enabled, "Canvas");
}

pp::paint_renderer::CanvasStrokeRasterizationPlan canvas_stroke_rasterization_plan(
    int width,
    int height) noexcept
{
    return pp::panopainter::plan_legacy_canvas_stroke_rasterization(
        ShaderManager::render_device_features(),
        width,
        height);
}

pp::paint_renderer::CanvasStrokeMaterialPlan canvas_stroke_material_plan(
    const Brush& brush,
    bool destination_feedback_needed) noexcept
{
    return pp::panopainter::plan_legacy_canvas_stroke_material(
        pp::paint_renderer::CanvasStrokeMaterialRequest {
            .destination_feedback_needed = destination_feedback_needed,
            .pattern_enabled = brush.m_pattern_enabled,
            .pattern_eachsample = brush.m_pattern_eachsample,
            .wet_blend = brush.m_tip_wet > 0.F,
            .mix_blend = brush.m_tip_mix > 0.F,
            .noise_enabled = brush.m_tip_noise > 0.F,
            .dual_brush_enabled = brush.m_dual_enabled,
            .dual_blend_mode = brush.m_dual_blend_mode,
            .pattern_blend_mode = brush.m_pattern_blend_mode,
            .dual_alpha = brush.m_dual_opacity,
        });
}

pp::renderer::Extent2D canvas_stroke_extent(int width, int height) noexcept
{
    return pp::renderer::Extent2D {
        .width = static_cast<std::uint32_t>(std::max(width, 0)),
        .height = static_cast<std::uint32_t>(std::max(height, 0)),
    };
}

} // namespace

std::vector<Canvas::StrokeFrame> Canvas::stroke_draw_compute(Stroke& stroke) const
{
    auto samples = stroke.compute_samples();
    return pp::panopainter::plan_legacy_canvas_stroke_frames(
        pp::panopainter::LegacyCanvasStrokeComputeRequest {
            .previous_sample = stroke.m_prev_sample,
            .samples = samples,
            .zoom = App::I->zoom,
            .mixer_size = glm::vec2(App::I->width, App::I->height),
            .model_view = m_mv,
        },
        [&](std::array<vertex_t, 4>& brush_quad, bool project_3d, glm::mat4 model_view) {
            return stroke_draw_project(brush_quad, project_3d, model_view);
        },
        [](
            glm::vec4 mixer_rect,
            glm::vec4 color,
            float flow,
            float opacity,
            std::array<std::vector<vertex_t>, 6>&& shapes) {
            return StrokeFrame {
                .col = color,
                .flow = flow,
                .opacity = opacity,
                .shapes = std::move(shapes),
                .m_mixer_rect = mixer_rect,
            };
        });
}

void Canvas::stroke_draw_pad_pass(
    const std::array<pp::panopainter::LegacyCanvasStrokePadFace, 6>& pad_faces,
    bool copy_stroke_destination,
    const std::array<pp::panopainter::LegacyCanvasStrokeTextureBinding, 1>& pad_destination_texture_binding,
    const pp::panopainter::LegacyCanvasStrokeTextureInputDispatch& pad_destination_texture_dispatch,
    const pp::renderer::Extent2D& stroke_extent,
    const glm::vec4& pad_color)
{
    pp::panopainter::setup_legacy_stroke_pad_shader(
        pp::panopainter::LegacyStrokePadUniforms {
            .color = pad_color,
            .uses_destination_feedback = copy_stroke_destination,
        });
    [[maybe_unused]] const auto pad_result = pp::panopainter::execute_legacy_canvas_stroke_pad_face_callbacks(
        pad_faces,
        stroke_extent,
        copy_stroke_destination,
        [&](std::span<const vertex_t> pad_quad) {
            m_brush_shape.update_vertices(
                const_cast<vertex_t*>(pad_quad.data()),
                static_cast<int>(pad_quad.size()));
        },
        [&](int face_index) {
            m_tmp[face_index].bindFramebuffer();
        },
        [&](int) {
            pp::panopainter::bind_legacy_canvas_stroke_texture_inputs(
                pad_destination_texture_binding,
                pad_destination_texture_dispatch);
        },
        [&](const pp::paint_renderer::CanvasStrokeCopyRegion& copy_region) {
            pp::panopainter::execute_legacy_canvas_stroke_pad_copy_region(
                copy_region,
                [](int src_x, int src_y, int dst_x, int dst_y, int width, int height) {
                    copy_framebuffer_to_texture_2d(src_x, src_y, dst_x, dst_y, width, height);
                });
        },
        [&](int) {
            pp::panopainter::unbind_legacy_canvas_stroke_texture_inputs(
                pad_destination_texture_binding,
                pad_destination_texture_dispatch);
        },
        [&] {
            m_brush_shape.draw_fill();
        },
        [&](int face_index) {
            m_tmp[face_index].unbindFramebuffer();
        });
}

void Canvas::stroke_draw_pad_face_orchestration(
    const std::array<bool, 6>& box_dirty,
    const std::array<glm::vec4, 6>& box_face,
    bool copy_stroke_destination,
    const pp::renderer::Extent2D& stroke_extent,
    const glm::vec4& pad_color)
{
    stroke_draw_pad_face_callback_body(
        box_dirty,
        box_face,
        copy_stroke_destination,
        stroke_extent,
        pad_color);
}

void Canvas::stroke_draw_pad_face_callback_body(
    const std::array<bool, 6>& box_dirty,
    const std::array<glm::vec4, 6>& box_face,
    bool copy_stroke_destination,
    const pp::renderer::Extent2D& stroke_extent,
    const glm::vec4& pad_color)
{
    pp::panopainter::setup_legacy_stroke_pad_shader(
        pp::panopainter::LegacyStrokePadUniforms {
            .color = pad_color,
            .uses_destination_feedback = copy_stroke_destination,
        });
    const auto pad_faces = pp::panopainter::make_legacy_canvas_stroke_pad_faces(box_dirty, box_face);
    constexpr std::array pad_destination_texture_binding {
        pp::panopainter::LegacyCanvasStrokeTextureBinding {
            .input = pp::panopainter::LegacyCanvasStrokeTextureInput::stroke_destination,
            .slot = 1,
        },
    };
    const auto pad_destination_texture_dispatch =
        pp::panopainter::make_legacy_canvas_stroke_pad_destination_texture_dispatch(
            [&](int texture_slot) {
                set_active_texture_unit(texture_slot);
            },
            [&](int dst_face_index) {
                m_tex[dst_face_index].bind();
            },
            [&](int dst_face_index) {
                m_tex[dst_face_index].unbind();
            },
            0);
    stroke_draw_pad_pass(
        pad_faces,
        copy_stroke_destination,
        pad_destination_texture_binding,
        pad_destination_texture_dispatch,
        stroke_extent,
        pad_color);
}

void Canvas::stroke_draw_dual_pass(
    const std::vector<StrokeFrame>& frames_dual,
    const std::array<pp::panopainter::LegacyCanvasStrokeTextureBinding, 1>& dual_pass_texture_bindings,
    const pp::panopainter::LegacyCanvasStrokeTextureInputDispatch& dual_pass_brush_tip_dispatch,
    const pp::renderer::Extent2D& stroke_extent,
    const std::array<bool, 6>& include_dual_dirty,
    bool uses_pattern,
    bool copy_stroke_destination)
{
    [[maybe_unused]] const auto dual_result =
        pp::panopainter::execute_legacy_canvas_stroke_dual_pass(
            make_stroke_draw_dual_pass_request(
                frames_dual,
                dual_pass_texture_bindings,
                dual_pass_brush_tip_dispatch,
                stroke_extent,
                include_dual_dirty,
                uses_pattern,
                copy_stroke_destination));
}

pp::panopainter::LegacyCanvasStrokeMainPassExecutionRequest Canvas::make_stroke_draw_main_pass_request(
    const Brush& brush,
    std::function<void()> bind_samplers,
    std::function<void()> execute_frame_pass,
    std::function<void()> unbind_samplers)
{
    constexpr std::array main_pass_texture_bindings {
        pp::panopainter::LegacyCanvasStrokeTextureBinding {
            .input = pp::panopainter::LegacyCanvasStrokeTextureInput::brush_tip,
            .slot = 0,
        },
        pp::panopainter::LegacyCanvasStrokeTextureBinding {
            .input = pp::panopainter::LegacyCanvasStrokeTextureInput::stroke_destination,
            .slot = 1,
        },
        pp::panopainter::LegacyCanvasStrokeTextureBinding {
            .input = pp::panopainter::LegacyCanvasStrokeTextureInput::pattern,
            .slot = 2,
        },
        pp::panopainter::LegacyCanvasStrokeTextureBinding {
            .input = pp::panopainter::LegacyCanvasStrokeTextureInput::mixer,
            .slot = 3,
        },
    };
    constexpr std::array main_pass_texture_unbindings {
        pp::panopainter::LegacyCanvasStrokeTextureBinding {
            .input = pp::panopainter::LegacyCanvasStrokeTextureInput::mixer,
            .slot = 3,
        },
        pp::panopainter::LegacyCanvasStrokeTextureBinding {
            .input = pp::panopainter::LegacyCanvasStrokeTextureInput::stroke_destination,
            .slot = 1,
        },
        pp::panopainter::LegacyCanvasStrokeTextureBinding {
            .input = pp::panopainter::LegacyCanvasStrokeTextureInput::brush_tip,
            .slot = 0,
        },
    };
    const auto main_pass_texture_dispatch = pp::panopainter::make_legacy_canvas_stroke_main_pass_texture_dispatch(
        [&](int texture_slot) {
            set_active_texture_unit(texture_slot);
        },
        [&] {
            brush.m_tip_texture->bind();
        },
        [&] {
            brush.m_tip_texture->unbind();
        },
        [&] {
            brush.m_pattern_texture ?
                brush.m_pattern_texture->bind() :
                unbind_texture_2d();
        },
        [&] {
            m_mixer.bindTexture();
        },
        [&] {
            m_mixer.unbindTexture();
        });

    return pp::panopainter::make_legacy_canvas_stroke_main_pass_execution_request(
        "Canvas::stroke_draw",
        std::move(bind_samplers),
        [&] {
            pp::panopainter::bind_legacy_canvas_stroke_texture_inputs(
                main_pass_texture_bindings,
                main_pass_texture_dispatch);
        },
        std::move(execute_frame_pass),
        [&] {
            pp::panopainter::unbind_legacy_canvas_stroke_texture_inputs(
                main_pass_texture_unbindings,
                main_pass_texture_dispatch);
        },
        std::move(unbind_samplers));
}

pp::panopainter::LegacyCanvasStrokeDualPassRequest Canvas::make_stroke_draw_dual_pass_request(
    const std::vector<StrokeFrame>& frames_dual,
    const std::array<pp::panopainter::LegacyCanvasStrokeTextureBinding, 1>& dual_pass_texture_bindings,
    const pp::panopainter::LegacyCanvasStrokeTextureInputDispatch& dual_pass_brush_tip_dispatch,
    const pp::renderer::Extent2D& stroke_extent,
    const std::array<bool, 6>& include_dual_dirty,
    bool uses_pattern,
    bool copy_stroke_destination)
{
    return pp::panopainter::LegacyCanvasStrokeDualPassRequest {
        .context = "Canvas::stroke_draw",
        .bind_brush_tip = [&] {
            pp::panopainter::bind_legacy_canvas_stroke_texture_inputs(
                dual_pass_texture_bindings,
                dual_pass_brush_tip_dispatch);
        },
        .unbind_brush_tip = [&] {
            pp::panopainter::unbind_legacy_canvas_stroke_texture_inputs(
                dual_pass_texture_bindings,
                dual_pass_brush_tip_dispatch);
        },
        .setup_dual_shader = [&] {
            pp::panopainter::setup_legacy_canvas_stroke_dual_shader(
                uses_pattern);
        },
        .execute_frame_pass = [&] {
            stroke_draw_dual_pass_frame_pass(
                frames_dual,
                stroke_extent,
                include_dual_dirty,
                copy_stroke_destination);
        },
    };
}

void Canvas::stroke_draw_dual_pass_frame_pass(
    const std::vector<StrokeFrame>& frames_dual,
    const pp::renderer::Extent2D& stroke_extent,
    const std::array<bool, 6>& include_dual_dirty,
    bool copy_stroke_destination)
{
    pp::panopainter::execute_legacy_canvas_stroke_dual_pass_frame_callbacks(
        frames_dual,
        stroke_extent,
        std::span<glm::vec4>(m_dirty_box),
        std::span<glm::vec4>(),
        std::span<const bool>(include_dual_dirty),
        [&](auto& f) {
            pp::panopainter::apply_legacy_stroke_sample_uniforms(
                pp::panopainter::LegacyStrokeSampleUniforms {
                    .color = f.col,
                    .alpha = f.flow,
                    .opacity = f.opacity,
                });
        },
        [](auto&, int, auto&) {},
        [&](auto&, int i, auto&& P) {
            auto polygon = std::move(P);
            return stroke_draw_samples(i, polygon, copy_stroke_destination);
        },
        m_tmp_dual,
        true);
}

void Canvas::stroke_draw_live()
{
    if (!(m_current_stroke && m_current_stroke->has_sample()))
    {
        stroke_commit_timelapse();
        //stroke_draw_mix({ 0,0 }, { m_mixer.getWidth(), m_mixer.getHeight() });
        return;
    }

    m_dirty = true;

    const auto vp = query_canvas_viewport();
    const auto cc = query_canvas_clear_color();
    
    const auto& brush = m_current_stroke->m_brush;
    auto ortho_proj = glm::ortho(0.f, (float)m_width, 0.f, (float)m_height, -1.f, 1.f);

    apply_canvas_viewport(0, 0, m_width, m_height);

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

    const auto stroke_rasterization = canvas_stroke_rasterization_plan(m_width, m_height);
    const bool copy_stroke_destination = stroke_rasterization.copy_stroke_destination;
    const auto stroke_material = canvas_stroke_material_plan(*brush, copy_stroke_destination);
    const auto stroke_extent = canvas_stroke_extent(m_width, m_height);

    apply_canvas_capability(blend_state(), false);
    pp::panopainter::setup_legacy_stroke_shader(
        pp::panopainter::LegacyStrokeShaderSetupUniforms {
            .resolution = glm::vec2(static_cast<float>(m_width), static_cast<float>(m_height)),
            .pattern = {
                .scale = patt_scale,
                .invert = static_cast<float>(brush->m_pattern_invert),
                .brightness = brush->m_pattern_brightness,
                .contrast = brush->m_pattern_contrast,
                .depth = brush->m_pattern_depth,
                .blend_mode = brush->m_pattern_blend_mode,
                .offset = m_pattern_offset,
            },
            .mvp = ortho_proj,
            .uses_destination_feedback = stroke_material.stroke_pass.uses_destination_feedback,
            .uses_pattern = stroke_material.stroke_pass.uses_pattern,
            .mix_alpha = brush->m_tip_mix,
            .wet = brush->m_tip_wet,
            .noise = brush->m_tip_noise,
            .set_opacity = true,
            .opacity = brush->m_tip_opacity,
        });

    // DRAW MAIN BRUSH
    constexpr std::array live_pass_sampler_bindings {
        pp::panopainter::LegacyCanvasStrokeTextureBinding {
            .input = pp::panopainter::LegacyCanvasStrokeTextureInput::brush_tip,
            .slot = 0,
        },
        pp::panopainter::LegacyCanvasStrokeTextureBinding {
            .input = pp::panopainter::LegacyCanvasStrokeTextureInput::stroke_destination,
            .slot = 1,
        },
        pp::panopainter::LegacyCanvasStrokeTextureBinding {
            .input = pp::panopainter::LegacyCanvasStrokeTextureInput::pattern,
            .slot = 2,
        },
        pp::panopainter::LegacyCanvasStrokeTextureBinding {
            .input = pp::panopainter::LegacyCanvasStrokeTextureInput::mixer,
            .slot = 3,
        },
    };
    const pp::panopainter::LegacyCanvasStrokeSamplerDispatch live_pass_sampler_dispatch =
        pp::panopainter::make_legacy_canvas_stroke_live_pass_sampler_dispatch(
            [&](int slot) {
                m_sampler_brush.bind(slot);
            },
            [&] {
                m_sampler_brush.unbind();
            },
            [&](int slot) {
                m_sampler_nearest.bind(slot);
            },
            [&] {
                m_sampler_nearest.unbind();
            },
            [&](int slot) {
                m_sampler_stencil.bind(slot);
            },
            [&] {
                m_sampler_stencil.unbind();
            },
            [&](int slot) {
                m_sampler.bind(slot);
            },
                [&] {
                    m_sampler.unbind();
                });
    auto frames = stroke_draw_compute(*m_current_stroke);
    
    std::array<glm::vec4, 6> box_face = SIXPLETTE(glm::vec4(m_size, 0, 0));
    std::array<bool, 6> box_dirty = SIXPLETTE(false);
    const std::array<bool, 6> include_main_dirty = SIXPLETTE(true);
    glm::vec4 pad_color;
    [[maybe_unused]] const auto main_pass_result =
        pp::panopainter::execute_legacy_canvas_stroke_main_pass(
            make_stroke_draw_main_pass_request(
                *brush,
                [&] {
                    pp::panopainter::bind_legacy_canvas_stroke_sampler_inputs(
                        live_pass_sampler_bindings,
                        live_pass_sampler_dispatch);
                },
                [&] {
                    pp::panopainter::execute_legacy_canvas_stroke_main_pass_frame_callbacks(
                        frames,
                        stroke_extent,
                        std::span<glm::vec4>(m_dirty_box),
                        std::span<glm::vec4>(box_face),
                        std::span<const bool>(include_main_dirty),
                        [&](auto& f) {
                            if (brush->m_tip_mix > 0.f)
                            {
                                stroke_draw_mix(xy(f.m_mixer_rect), zw(f.m_mixer_rect));
                            }
                            pad_color = f.col;
                        },
                        [&](auto&, int i, auto&) {
                            m_dirty_face[i] = true;
                            box_dirty[i] = true;
                        },
                        [&](auto& f, int i, auto& P) {
                            pp::panopainter::use_legacy_stroke_shader();
                            pp::panopainter::apply_legacy_stroke_sample_uniforms(
                                pp::panopainter::LegacyStrokeSampleUniforms {
                                    .color = f.col,
                                    .alpha = f.flow,
                                    .opacity = f.opacity,
                            });
                            return stroke_draw_samples(i, P, copy_stroke_destination);
                        },
                        m_tmp);
                },
                [&] {
                    pp::panopainter::unbind_legacy_canvas_stroke_sampler_inputs(
                        live_pass_sampler_bindings,
                        live_pass_sampler_dispatch);
                }
            ));

    // pad stroke
    // In order to mitigate color bleeding at the edge of shapes in transparent layers
    // we need to fill the area around the stroke with the same color because by default
    // the transparent area may have black or other undefined color.
    // This step is only useful for previewing the stroke because on commit the dilate
    // algorithm fixes this issue.
    // NOTE: at the moment this works on the whole canvas, but it can be optimized 
    // to only affect the current dirty box. In this case it may be necessary to do this
    // work on documents that doesn't have the padding, so on document loading.
    stroke_draw_pad_face_orchestration(
        box_dirty,
        box_face,
        copy_stroke_destination,
        stroke_extent,
        pad_color);

    // DRAW DUAL BRUSH

    if (stroke_material.dual_pass.enabled && m_dual_stroke)
    {
        const auto& dual_brush = m_dual_stroke->m_brush;
        constexpr std::array dual_pass_texture_bindings {
            pp::panopainter::LegacyCanvasStrokeTextureBinding {
                .input = pp::panopainter::LegacyCanvasStrokeTextureInput::brush_tip,
                .slot = 0,
            },
        };
        const auto dual_pass_brush_tip_dispatch =
            pp::panopainter::make_legacy_canvas_stroke_brush_tip_texture_dispatch(
                [&](int texture_slot) {
                    set_active_texture_unit(texture_slot);
                },
                [&](int) {
                    dual_brush->m_tip_texture ?
                        dual_brush->m_tip_texture->bind() :
                        unbind_texture_2d();
                },
                [&](int) {
                    dual_brush->m_tip_texture ?
                        dual_brush->m_tip_texture->unbind() :
                        unbind_texture_2d();
                },
                0);
        auto frames_dual = stroke_draw_compute(*m_dual_stroke);
        const std::array<bool, 6> include_dual_dirty =
            SIXPLETTE(stroke_material.composite_pass.dual_blend_mode == 0);
        stroke_draw_dual_pass(
            frames_dual,
            dual_pass_texture_bindings,
            dual_pass_brush_tip_dispatch,
            stroke_extent,
            include_dual_dirty,
            stroke_material.dual_pass.uses_pattern,
            copy_stroke_destination);
    }

    pp::panopainter::unbind_legacy_canvas_stroke_sampler_inputs(
        live_pass_sampler_bindings,
        live_pass_sampler_dispatch);
    
    apply_canvas_viewport(vp.x, vp.y, vp.width, vp.height);
    apply_canvas_clear_color(cc);

    if (m_commit_delayed)
    {
        m_show_tmp = false;
        m_commit_delayed = false;
        stroke_commit();
        m_current_stroke = nullptr;
    }
}

void Canvas::stroke_draw()
{
    stroke_draw_live();
}
