#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_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);
}

std::array<pp::paint_renderer::CanvasStrokePoint, 4> make_stroke_preview_sample_points(
    const std::array<vertex_t, 4>& vertices)
{
    return {
        pp::paint_renderer::CanvasStrokePoint { .x = vertices[0].pos.x, .y = vertices[0].pos.y },
        pp::paint_renderer::CanvasStrokePoint { .x = vertices[1].pos.x, .y = vertices[1].pos.y },
        pp::paint_renderer::CanvasStrokePoint { .x = vertices[2].pos.x, .y = vertices[2].pos.y },
        pp::paint_renderer::CanvasStrokePoint { .x = vertices[3].pos.x, .y = vertices[3].pos.y },
    };
}

void upload_stroke_preview_brush_vertices(DynamicShape& brush_shape, std::span<const vertex_t> vertices);

pp::panopainter::LegacyStrokeSampleExecutionRequest make_stroke_preview_sample_request(
    std::array<vertex_t, 4>& vertices,
    const std::array<pp::paint_renderer::CanvasStrokePoint, 4>& sample_points,
    glm::vec2 target_size,
    Texture2D& blend_texture,
    DynamicShape& brush_shape,
    bool copy_stroke_destination)
{
    return pp::panopainter::LegacyStrokeSampleExecutionRequest {
        .context = "NodeStrokePreview::stroke_draw_samples",
        .target_size = target_size,
        .vertices = vertices,
        .sample_points = sample_points,
        .copy_stroke_destination = copy_stroke_destination,
        .bind_destination_texture = [&] {
            bind_stroke_preview_destination_texture(blend_texture);
        },
        .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_texture);
        },
        .upload_brush_vertices = [&](std::span<const vertex_t> brush_vertices) {
            upload_stroke_preview_brush_vertices(brush_shape, brush_vertices);
        },
        .draw_brush_shape = [&] {
            brush_shape.draw_fill();
        },
    };
}

void upload_stroke_preview_brush_vertices(DynamicShape& brush_shape, std::span<const vertex_t> vertices)
{
    brush_shape.update_vertices(
        const_cast<vertex_t*>(vertices.data()),
        static_cast<int>(vertices.size()));
}

glm::vec4 execute_stroke_preview_sample_pass(
    std::array<vertex_t, 4>& vertices,
    glm::vec2 target_size,
    Texture2D& blend_texture,
    DynamicShape& brush_shape,
    bool copy_stroke_destination)
{
    const auto sample_points = make_stroke_preview_sample_points(vertices);
    const auto result = pp::panopainter::execute_legacy_canvas_stroke_sample(
        make_stroke_preview_sample_request(
            vertices,
            sample_points,
            target_size,
            blend_texture,
            brush_shape,
            copy_stroke_destination));
    return result.dirty_bounds;
}

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 execute_stroke_preview_sample_pass(
        P,
        size,
        blend_tex,
        m_brush_shape,
        copy_stroke_destination);
}

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();
    
    float zoom = root()->m_zoom;

    glm::vec2 size = { m_rtt.getWidth(), m_rtt.getHeight() };
    const auto stroke_setup = pp::panopainter::plan_legacy_node_stroke_preview_stroke_setup(
        pp::panopainter::LegacyNodeStrokePreviewStrokeSetupRequest {
            .preview_size = m_size,
            .zoom = zoom,
            .brush_tip_size = m_brush->m_tip_size,
            .stroke_max_size_override = m_max_size,
            .pad_override = m_pad_override,
            .tip_size_pressure = m_brush->m_tip_size_pressure,
            .dual_enabled = m_brush->m_dual_enabled,
            .dual_size = m_brush->m_dual_size,
            .pattern_scale = m_brush->m_pattern_scale,
            .pattern_flipx = m_brush->m_pattern_flipx,
            .pattern_flipy = m_brush->m_pattern_flipy,
        });
    glm::mat4 ortho_proj = glm::ortho<float>(0, size.x, 0, size.y, -1, 1);
    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);

    const auto& b = m_brush;

    auto prepared_strokes = pp::panopainter::prepare_legacy_node_stroke_preview_strokes(
        b,
        stroke_setup,
        Canvas::I->m_cam_fov,
        Canvas::I->m_cam_rot);

    apply_stroke_preview_capability(pp::renderer::gl::blend_state(), false);
    const auto pass_orchestration = pp::panopainter::plan_legacy_node_stroke_preview_pass_orchestration(
        pp::panopainter::make_legacy_node_stroke_preview_pass_orchestration_request(
            stroke_preview_render_device_features(),
            size,
            *b,
            ortho_proj));
    pp::panopainter::setup_legacy_stroke_shader(pass_orchestration.stroke_shader);
    const bool sequence_ok = pp::panopainter::execute_legacy_node_stroke_preview_live_render_passes(
        pp::panopainter::LegacyNodeStrokePreviewLiveRenderRequest {
            .brush = *b,
            .pass_orchestration = pass_orchestration,
            .prepared_strokes = prepared_strokes,
            .stroke_texture = m_tex,
            .mixer_rtt = m_rtt_mixer,
            .render_target = m_rtt,
            .background_texture = m_tex_background,
            .dual_texture = m_tex_dual,
            .preview_texture = m_tex_preview,
            .linear_sampler = m_sampler_linear,
            .repeat_sampler = m_sampler_linear_repeat,
            .zoom = zoom,
            .min_flow = m_min_flow,
            .copy_stroke_destination = pass_orchestration.copy_stroke_destination,
            .size = size,
            .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 = [&] {
                b->m_pattern_texture ? b->m_pattern_texture->bind() : unbind_texture_2d();
            },
            .draw_composite = [&] {
                m_plane.draw_fill();
            },
        });
    assert(sequence_ok);

    m_rtt.unbindFramebuffer();

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