#include "pch.h"
#include "log.h"
#include "node_stroke_preview.h"
#include "texture.h"
#include "shader.h"
#include "bezier.h"
#include "canvas.h"
#include "app.h"
#include "legacy_canvas_draw_merge_services.h"
#include "legacy_canvas_stroke_composite_services.h"
#include "legacy_canvas_stroke_execution_services.h"
#include "legacy_canvas_stroke_preview_services.h"
#include "legacy_canvas_stroke_shader_services.h"
#include "legacy_canvas_stroke_services.h"
#include "legacy_node_stroke_preview_execution_services.h"
#include "legacy_node_stroke_preview_runtime_services.h"
#include "legacy_node_stroke_preview_sample_services.h"
#include "legacy_ui_gl_dispatch.h"
#include "paint_renderer/compositor.h"
#include "renderer_gl/opengl_capabilities.h"
#include "util.h"
#include <array>
#include <cstdint>

namespace {

pp::renderer::RenderDeviceFeatures stroke_preview_render_device_features() noexcept
{
    return ShaderManager::render_device_features();
}

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

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

void apply_stroke_preview_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, "NodeStrokePreview");
}

pp::renderer::gl::OpenGlViewportRect query_stroke_preview_viewport()
{
    return pp::legacy::ui_gl::query_viewport_rect("NodeStrokePreview");
}

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

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

void apply_stroke_preview_scissor(std::int32_t x, std::int32_t y, std::int32_t width, std::int32_t height)
{
    pp::legacy::ui_gl::apply_scissor_rect(x, y, width, height, "NodeStrokePreview");
}

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

namespace stroke_preview_composite_slots {
constexpr std::uint32_t kBackground = 0U;
constexpr std::uint32_t kStroke = 1U;
constexpr std::uint32_t kDual = 3U;
constexpr std::uint32_t kPattern = 4U;
}

namespace stroke_preview_live_slots {
constexpr std::uint32_t kTip = 0U;
constexpr std::uint32_t kDestination = 1U;
constexpr std::uint32_t kPattern = 2U;
constexpr std::uint32_t kMixer = 3U;
constexpr std::uint32_t kReservedLinear = 4U;
}

pp::panopainter::LegacyNodeStrokePreviewMixPassRequest make_stroke_preview_mix_pass_request(
    const Brush& brush,
    glm::vec2 resolution) noexcept
{
    return {
        .resolution = resolution,
        .pattern_scale = brush.m_pattern_scale,
        .pattern_flipx = brush.m_pattern_flipx,
        .pattern_flipy = brush.m_pattern_flipy,
        .pattern_invert = brush.m_pattern_invert,
        .pattern_brightness = brush.m_pattern_brightness,
        .pattern_contrast = brush.m_pattern_contrast,
        .pattern_depth = brush.m_pattern_depth,
        .pattern_rand_offset = brush.m_pattern_rand_offset,
        .pattern_enabled = brush.m_pattern_enabled,
        .pattern_eachsample = brush.m_pattern_eachsample,
        .tip_wet = brush.m_tip_wet,
        .tip_mix = brush.m_tip_mix,
        .tip_noise = brush.m_tip_noise,
        .dual_enabled = brush.m_dual_enabled,
        .dual_blend_mode = brush.m_dual_blend_mode,
        .pattern_blend_mode = brush.m_pattern_blend_mode,
        .dual_opacity = brush.m_dual_opacity,
        .blend_mode = brush.m_blend_mode,
    };
}

pp::panopainter::LegacyStrokeCompositeUniforms make_stroke_preview_mix_composite_uniforms(
    const pp::panopainter::LegacyNodeStrokePreviewMixPassPlan::ShaderPlan& shader_plan) noexcept
{
    return {
        .resolution = shader_plan.resolution,
        .pattern = {
            .scale = shader_plan.pattern_scale,
            .invert = shader_plan.pattern_invert,
            .brightness = shader_plan.pattern_brightness,
            .contrast = shader_plan.pattern_contrast,
            .depth = shader_plan.pattern_depth,
            .blend_mode = shader_plan.pattern_blend_mode,
            .offset = shader_plan.pattern_offset,
        },
        .mvp = glm::ortho(-.5f, .5f, -.5f, .5f, -1.f, 1.f),
        .layer_alpha = 1.0f,
        .alpha_lock = false,
        .mask_enabled = false,
        .use_fragcoord = false,
        .blend_mode = shader_plan.blend_mode,
        .use_dual = shader_plan.use_dual,
        .dual_blend_mode = shader_plan.dual_blend_mode,
        .dual_alpha = shader_plan.dual_alpha,
        .use_pattern = shader_plan.use_pattern,
    };
}

pp::panopainter::LegacyCanvasStrokeMixPassRequest make_stroke_preview_mix_pass_execution_request(
    const pp::panopainter::LegacyNodeStrokePreviewMixPassPlan::ShaderPlan& shader,
    RTT& mixer_rtt,
    const Brush& brush,
    Sampler& linear_sampler,
    Texture2D& background_texture,
    Texture2D& stroke_texture,
    Texture2D& dual_texture,
    std::span<const pp::panopainter::LegacyCanvasStrokeMixPassPlane> mix_planes,
    std::function<void()> draw_mix)
{
    return pp::panopainter::make_legacy_canvas_stroke_mix_pass_request(
        "NodeStrokePreview::stroke_draw_mix",
        glm::vec2(static_cast<float>(mixer_rtt.getWidth()), static_cast<float>(mixer_rtt.getHeight())),
        mix_planes,
        [&] {
            linear_sampler.bind(stroke_preview_composite_slots::kBackground);
            linear_sampler.bind(stroke_preview_composite_slots::kStroke);
            linear_sampler.bind(stroke_preview_composite_slots::kDual);
            linear_sampler.bind(stroke_preview_composite_slots::kPattern);
            set_active_texture_unit(stroke_preview_composite_slots::kBackground);
            background_texture.bind();
            set_active_texture_unit(stroke_preview_composite_slots::kStroke);
            stroke_texture.bind();
            set_active_texture_unit(stroke_preview_composite_slots::kDual);
            dual_texture.bind();
            set_active_texture_unit(stroke_preview_composite_slots::kPattern);
            brush.m_pattern_texture ? brush.m_pattern_texture->bind() : unbind_texture_2d();
        },
        [] {},
        [&](int, const glm::mat4& plane_mvp) {
            auto uniforms = make_stroke_preview_mix_composite_uniforms(shader);
            uniforms.mvp = plane_mvp;
            pp::panopainter::setup_legacy_stroke_composite_shader(uniforms);
        },
        [&](int) {
            set_active_texture_unit(stroke_preview_composite_slots::kBackground);
            background_texture.bind();
        },
        [&](int) {
            set_active_texture_unit(stroke_preview_composite_slots::kStroke);
            stroke_texture.bind();
        },
        [&](int) {
            set_active_texture_unit(stroke_preview_composite_slots::kDual);
            dual_texture.bind();
        },
        std::move(draw_mix),
        [&](int) {
            set_active_texture_unit(stroke_preview_composite_slots::kDual);
            dual_texture.unbind();
        },
        [&](int) {
            set_active_texture_unit(stroke_preview_composite_slots::kStroke);
            stroke_texture.unbind();
        },
        [&](int) {
            set_active_texture_unit(stroke_preview_composite_slots::kBackground);
            background_texture.unbind();
    });
}

void bind_stroke_preview_live_samplers(
    Sampler& mipmap_sampler,
    Sampler& linear_sampler,
    Sampler& repeat_sampler)
{
    mipmap_sampler.bind(stroke_preview_live_slots::kTip);
    linear_sampler.bind(stroke_preview_live_slots::kDestination);
    repeat_sampler.bind(stroke_preview_live_slots::kPattern);
    linear_sampler.bind(stroke_preview_live_slots::kMixer);
    linear_sampler.bind(stroke_preview_live_slots::kReservedLinear);
}

void bind_stroke_preview_dual_pass_textures(const Brush& dual_brush)
{
    set_active_texture_unit(stroke_preview_live_slots::kTip);
    dual_brush.m_tip_texture ?
        dual_brush.m_tip_texture->bind() :
        unbind_texture_2d();
}

void bind_stroke_preview_destination_texture(Texture2D& texture)
{
    set_active_texture_unit(stroke_preview_live_slots::kDestination);
    texture.bind();
}

void unbind_stroke_preview_destination_texture(Texture2D& texture)
{
    set_active_texture_unit(stroke_preview_live_slots::kDestination);
    texture.unbind();
}

void copy_stroke_preview_destination_texture_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);
}

void execute_stroke_preview_background_capture_pass(
    glm::vec2 size,
    bool colorize,
    Texture2D& background_texture,
    const std::function<void()>& draw_checkerboard)
{
    const float aspect = size.x / size.y;
    pp::panopainter::setup_legacy_canvas_draw_merge_checkerboard_shader(
        pp::panopainter::LegacyCanvasDrawMergeCheckerboardUniforms {
            .mvp = glm::ortho(-.5f, .5f, -.5f / aspect, .5f / aspect, -1.f, 1.f),
            .colorize = colorize,
        });
    draw_checkerboard();
    const auto copy_status = pp::paint_renderer::copy_stroke_preview_result_to_texture(
        [&] {
            background_texture.bind();
        },
        [](
            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);
        },
        pp::paint_renderer::StrokePreviewCopySize {
            .width = static_cast<int>(size.x),
            .height = static_cast<int>(size.y),
        });
    assert(copy_status.ok());
}

}

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

void NodeStrokePreview::clone_copy(Node* dest) const
{
    NodeBorder::clone_copy(dest);
}

void NodeStrokePreview::clone_children(Node* dest) const
{
    // stop children cloning
}

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

void NodeStrokePreview::init_controls()
{
//     TextureManager::load("data/thumbs/Round-Hard.png");
//     Canvas::I->m_current_brush.m_tex_id = const_hash("data/thumbs/Round-Hard.png");
}

void NodeStrokePreview::stroke_draw_mix(const glm::vec2& bb_min, const glm::vec2& bb_sz)
{
    const auto& b = m_brush;
    const auto mix_pass = pp::panopainter::plan_legacy_node_stroke_preview_mix_pass(
        make_stroke_preview_mix_pass_request(*b, m_size));
    const auto mix_planes = std::array {
        pp::panopainter::LegacyCanvasStrokeMixPassPlane {
            .index = 0,
            .visible = true,
            .has_target = true,
            .opacity = 1.0f,
            .mvp = glm::ortho(-.5f, .5f, -.5f, .5f, -1.f, 1.f),
        },
    };
    gl_state gl;
    const auto mix_result = pp::panopainter::execute_legacy_canvas_stroke_mix_pass_with_setup(
        [&] {
            apply_stroke_preview_viewport(0, 0, m_rtt_mixer.getWidth(), m_rtt_mixer.getHeight());
            apply_stroke_preview_capability(pp::renderer::gl::depth_test_state(), false);
            apply_stroke_preview_capability(pp::renderer::gl::scissor_test_state(), true);
            apply_stroke_preview_capability(pp::renderer::gl::blend_state(), false);
            apply_stroke_preview_scissor(
                static_cast<std::int32_t>(bb_min.x),
                static_cast<std::int32_t>(bb_min.y),
                static_cast<std::int32_t>(bb_sz.x),
                static_cast<std::int32_t>(bb_sz.y));
            gl.save();
            m_rtt_mixer.bindFramebuffer();
        },
        [&] {
            m_rtt_mixer.unbindFramebuffer();
            gl.restore();
        },
        make_stroke_preview_mix_pass_execution_request(
            mix_pass.shader,
            m_rtt_mixer,
            *m_brush,
            m_sampler_linear,
            m_tex_background,
            m_tex,
            m_tex_dual,
            mix_planes,
            [this] {
                m_plane.draw_fill();
            }));
    assert(mix_result.ok);
}

glm::vec4 NodeStrokePreview::stroke_draw_samples(
    std::array<vertex_t, 4>& P,
    Texture2D& blend_tex,
    bool copy_stroke_destination)
{
    const glm::vec2 size = { m_rtt.getWidth(), m_rtt.getHeight() };
    return pp::panopainter::execute_legacy_node_stroke_preview_sample_pass(
        pp::panopainter::LegacyNodeStrokePreviewSamplePassRequest {
            .target_size = size,
            .vertices = P,
            .brush_shape = m_brush_shape,
            .copy_stroke_destination = copy_stroke_destination,
            .bind_destination_texture = [&] {
                bind_stroke_preview_destination_texture(blend_tex);
            },
            .copy_framebuffer_to_destination_texture = [](
                int src_x,
                int src_y,
                int dst_x,
                int dst_y,
                int width,
                int height) {
                copy_stroke_preview_destination_texture_region(
                    src_x,
                    src_y,
                    dst_x,
                    dst_y,
                    width,
                    height);
            },
            .unbind_destination_texture = [&] {
                unbind_stroke_preview_destination_texture(blend_tex);
            },
        });
}

std::vector<NodeStrokePreview::StrokeFrame> NodeStrokePreview::stroke_draw_compute(const Stroke& stroke, float zoom) const
{
    auto samples = const_cast<Stroke&>(stroke).compute_samples();
    StrokeSample previous_sample = stroke.m_prev_sample;
    previous_sample.size *= zoom;
    for (auto& sample : samples) {
        sample.size *= zoom;
    }

    return pp::panopainter::plan_legacy_canvas_stroke_frames(
        pp::panopainter::LegacyCanvasStrokeComputeRequest {
            .previous_sample = previous_sample,
            .samples = samples,
            .zoom = 1.0f,
            .mixer_size = glm::vec2(m_rtt.getWidth(), m_rtt.getHeight()),
        },
        [](
            std::array<vertex_t, 4>& brush_quad,
            bool /*project_3d*/,
            glm::mat4 /*model_view*/) {
            return brush_quad;
        },
        [](
            glm::vec4 mixer_rect,
            glm::vec4 color,
            float flow,
            float opacity,
            std::array<vertex_t, 4>&& shapes) -> StrokeFrame {
            return StrokeFrame {
                .col = color,
                .flow = flow,
                .opacity = opacity,
                .shapes = std::move(shapes),
                .mixer_rect = mixer_rect,
            };
        });
}

void NodeStrokePreview::draw_stroke_immediate()
{
    if (m_size.x == 0 || m_size.y == 0)
        return;

    if (!pp::panopainter::ensure_legacy_node_stroke_preview_render_targets(
            pp::panopainter::LegacyNodeStrokePreviewRenderTargetSetup {
                .preview_rtt = m_rtt,
                .preview_rtt_mixer = m_rtt_mixer,
                .preview_stroke_texture = m_tex,
                .preview_dual_texture = m_tex_dual,
                .preview_background_texture = m_tex_background,
                .preview_image_texture = m_tex_preview,
                .size = m_preview_size,
            })) {
        return;
    }

    const auto vp = query_stroke_preview_viewport();
    const auto cc = query_stroke_preview_clear_color();
    const glm::vec2 size = { m_rtt.getWidth(), m_rtt.getHeight() };
    const float zoom = root()->m_zoom;
    const bool sequence_ok = pp::panopainter::execute_legacy_node_stroke_preview_immediate_runtime(
        pp::panopainter::LegacyNodeStrokePreviewImmediateRuntimeRequest {
            .brush = m_brush,
            .preview_size = m_size,
            .zoom = zoom,
            .min_flow = m_min_flow,
            .stroke_max_size_override = m_max_size,
            .pad_override = m_pad_override,
            .camera_fov = Canvas::I->m_cam_fov,
            .camera_rot = Canvas::I->m_cam_rot,
            .render_device_features = stroke_preview_render_device_features(),
            .preview_rtt = m_rtt,
            .preview_rtt_mixer = m_rtt_mixer,
            .preview_stroke_texture = m_tex,
            .preview_dual_texture = m_tex_dual,
            .preview_background_texture = m_tex_background,
            .preview_image_texture = m_tex_preview,
            .linear_sampler = m_sampler_linear,
            .repeat_sampler = m_sampler_linear_repeat,
            .prepare_render_target = [&] {
                apply_stroke_preview_viewport(0, 0, m_rtt.getWidth(), m_rtt.getHeight());
                m_rtt.bindFramebuffer();
                m_rtt.clear();
                bind_stroke_preview_live_samplers(
                    m_sampler_mipmap,
                    m_sampler_linear,
                    m_sampler_linear_repeat);
            },
            .finish_render_target = [&] {
                m_rtt.unbindFramebuffer();
            },
            .set_blend_enabled = [&](bool enabled) {
                apply_stroke_preview_capability(pp::renderer::gl::blend_state(), enabled);
            },
            .setup_stroke_shader = [](const pp::panopainter::LegacyStrokeShaderSetupUniforms& uniforms) {
                pp::panopainter::setup_legacy_stroke_shader(uniforms);
            },
            .bind_dual_pass_textures = [](const Brush& dual_brush) {
                bind_stroke_preview_dual_pass_textures(dual_brush);
            },
            .capture_background = [&](bool colorize) {
                execute_stroke_preview_background_capture_pass(
                    size,
                    colorize,
                    m_tex_background,
                    [&] {
                        m_plane.draw_fill();
                    });
            },
            .compute_frames = [&](const Stroke& stroke, float frame_zoom) {
                return stroke_draw_compute(stroke, frame_zoom);
            },
            .draw_samples = [&](std::array<vertex_t, 4>& shapes, Texture2D& texture, bool copy_stroke_destination) {
                return stroke_draw_samples(shapes, texture, copy_stroke_destination);
            },
            .draw_mix = [&](const glm::vec2& bb_min, const glm::vec2& bb_sz) {
                stroke_draw_mix(bb_min, bb_sz);
            },
            .unbind_mixer_texture = [&] {
                set_active_texture_unit(3U);
                m_rtt_mixer.unbindTexture();
            },
            .bind_pattern_texture = [&] {
                m_brush->m_pattern_texture ? m_brush->m_pattern_texture->bind() : unbind_texture_2d();
            },
            .draw_composite = [&] {
                m_plane.draw_fill();
            },
        });
    assert(sequence_ok);

    apply_stroke_preview_viewport(vp.x, vp.y, vp.width, vp.height);
    apply_stroke_preview_clear_color(cc);
}