#include "pch.h"
#include "legacy_node_stroke_preview_runtime_services.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_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>
#include <stop_token>

namespace pp::panopainter {

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

void apply_legacy_node_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_legacy_node_stroke_preview_viewport()
{
    return pp::legacy::ui_gl::query_viewport_rect("NodeStrokePreview");
}

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

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

void apply_legacy_node_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_legacy_node_stroke_preview_capability(std::uint32_t state, bool enabled)
{
    pp::legacy::ui_gl::set_capability(state, enabled, "NodeStrokePreview");
}

namespace {

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;
}

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");
}

} // namespace

void bind_legacy_node_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_legacy_node_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_legacy_node_stroke_preview_pattern_texture(const Brush& brush)
{
    set_active_texture_unit(stroke_preview_live_slots::kPattern);
    brush.m_pattern_texture ? brush.m_pattern_texture->bind() : unbind_texture_2d();
}

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

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

void unbind_legacy_node_stroke_preview_mixer_texture(RTT& mixer_rtt)
{
    set_active_texture_unit(stroke_preview_live_slots::kMixer);
    mixer_rtt.unbindTexture();
}

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

bool execute_legacy_node_stroke_preview_immediate_draw(
    const LegacyNodeStrokePreviewImmediateDrawRequest& request)
{
    if (!request.brush ||
        !request.prepare_render_target ||
        !request.finish_render_target ||
        !request.compute_frames ||
        !request.draw_samples ||
        !request.draw_mix ||
        !request.draw_checkerboard ||
        !request.draw_composite) {
        return false;
    }

    if (!ensure_legacy_node_stroke_preview_render_targets(
            LegacyNodeStrokePreviewRenderTargetSetup {
                .preview_rtt = request.preview_rtt,
                .preview_rtt_mixer = request.preview_rtt_mixer,
                .preview_stroke_texture = request.preview_stroke_texture,
                .preview_dual_texture = request.preview_dual_texture,
                .preview_background_texture = request.preview_background_texture,
                .preview_image_texture = request.preview_image_texture,
                .size = request.preview_size,
            })) {
        return false;
    }

    const auto vp = query_legacy_node_stroke_preview_viewport();
    const auto cc = query_legacy_node_stroke_preview_clear_color();
    const glm::vec2 size = {
        static_cast<float>(request.preview_rtt.getWidth()),
        static_cast<float>(request.preview_rtt.getHeight()),
    };
    const bool sequence_ok = execute_legacy_node_stroke_preview_immediate_runtime(
        LegacyNodeStrokePreviewImmediateRuntimeRequest {
            .brush = request.brush,
            .preview_size = request.preview_size,
            .zoom = request.zoom,
            .min_flow = request.min_flow,
            .stroke_max_size_override = request.stroke_max_size_override,
            .pad_override = request.pad_override,
            .camera_fov = request.camera_fov,
            .camera_rot = request.camera_rot,
            .render_device_features = request.render_device_features,
            .preview_rtt = request.preview_rtt,
            .preview_rtt_mixer = request.preview_rtt_mixer,
            .preview_stroke_texture = request.preview_stroke_texture,
            .preview_dual_texture = request.preview_dual_texture,
            .preview_background_texture = request.preview_background_texture,
            .preview_image_texture = request.preview_image_texture,
            .linear_sampler = request.linear_sampler,
            .repeat_sampler = request.repeat_sampler,
            .prepare_render_target = request.prepare_render_target,
            .finish_render_target = request.finish_render_target,
            .set_blend_enabled = [](bool enabled) {
                apply_legacy_node_stroke_preview_capability(pp::renderer::gl::blend_state(), enabled);
            },
            .setup_stroke_shader = [](const LegacyStrokeShaderSetupUniforms& uniforms) {
                setup_legacy_stroke_shader(uniforms);
            },
            .bind_dual_pass_textures = [](const Brush& dual_brush) {
                bind_legacy_node_stroke_preview_dual_pass_textures(dual_brush);
            },
            .capture_background = [&](bool colorize) {
                execute_legacy_node_stroke_preview_background_capture_pass(
                    LegacyNodeStrokePreviewBackgroundCaptureRequest {
                        .size = size,
                        .colorize = colorize,
                        .background_texture = request.preview_background_texture,
                        .draw_checkerboard = [&] {
                            request.draw_checkerboard();
                        },
                    });
            },
            .compute_frames = request.compute_frames,
            .draw_samples = request.draw_samples,
            .draw_mix = request.draw_mix,
            .unbind_mixer_texture = [&] {
                unbind_legacy_node_stroke_preview_mixer_texture(request.preview_rtt_mixer);
            },
            .bind_pattern_texture = [&] {
                bind_legacy_node_stroke_preview_pattern_texture(*request.brush);
            },
            .draw_composite = request.draw_composite,
        });
    apply_legacy_node_stroke_preview_viewport(vp.x, vp.y, vp.width, vp.height);
    apply_legacy_node_stroke_preview_clear_color(cc);
    return sequence_ok;
}

void execute_legacy_node_stroke_preview_background_capture_pass(
    const LegacyNodeStrokePreviewBackgroundCaptureRequest& request)
{
    if (!request.draw_checkerboard) {
        return;
    }

    const float aspect = request.size.x / request.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 = request.colorize,
        });
    request.draw_checkerboard();
    const auto copy_status = pp::paint_renderer::copy_stroke_preview_result_to_texture(
        [&] {
            request.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>(request.size.x),
            .height = static_cast<int>(request.size.y),
        });
    assert(copy_status.ok());
}

std::vector<LegacyNodeStrokePreviewFrame> plan_legacy_node_stroke_preview_stroke_frames(
    const LegacyNodeStrokePreviewStrokeComputeRequest& request)
{
    auto samples = const_cast<Stroke&>(request.stroke).compute_samples();
    StrokeSample previous_sample = request.stroke.m_prev_sample;
    previous_sample.size *= request.zoom;
    for (auto& sample : samples) {
        sample.size *= request.zoom;
    }

    return pp::panopainter::plan_legacy_canvas_stroke_frames(
        pp::panopainter::LegacyCanvasStrokeComputeRequest {
            .previous_sample = previous_sample,
            .samples = samples,
            .zoom = 1.0f,
            .mixer_size = request.mixer_size,
        },
        [](
            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) -> LegacyNodeStrokePreviewFrame {
            return LegacyNodeStrokePreviewFrame {
                .col = color,
                .flow = flow,
                .opacity = opacity,
                .shapes = std::move(shapes),
                .mixer_rect = mixer_rect,
            };
        });
}

std::vector<LegacyNodeStrokePreviewFrame> plan_legacy_node_stroke_preview_stroke_frames(
    const Stroke& stroke,
    float zoom,
    glm::vec2 mixer_size)
{
    return plan_legacy_node_stroke_preview_stroke_frames(
        LegacyNodeStrokePreviewStrokeComputeRequest {
            .stroke = stroke,
            .zoom = zoom,
            .mixer_size = mixer_size,
        });
}

bool execute_legacy_node_stroke_preview_immediate_runtime(
    const LegacyNodeStrokePreviewImmediateRuntimeRequest& request)
{
    if (!request.brush ||
        !request.prepare_render_target ||
        !request.finish_render_target ||
        !request.set_blend_enabled ||
        !request.setup_stroke_shader) {
        return false;
    }

    const glm::vec2 render_target_size = {
        static_cast<float>(request.preview_rtt.getWidth()),
        static_cast<float>(request.preview_rtt.getHeight()),
    };
    const auto stroke_setup = plan_legacy_node_stroke_preview_stroke_setup(
        LegacyNodeStrokePreviewStrokeSetupRequest {
            .preview_size = request.preview_size,
            .zoom = request.zoom,
            .brush_tip_size = request.brush->m_tip_size,
            .stroke_max_size_override = request.stroke_max_size_override,
            .pad_override = request.pad_override,
            .tip_size_pressure = request.brush->m_tip_size_pressure,
            .dual_enabled = request.brush->m_dual_enabled,
            .dual_size = request.brush->m_dual_size,
            .pattern_scale = request.brush->m_pattern_scale,
            .pattern_flipx = request.brush->m_pattern_flipx,
            .pattern_flipy = request.brush->m_pattern_flipy,
        });
    const auto prepared_strokes = prepare_legacy_node_stroke_preview_strokes(
        request.brush,
        stroke_setup,
        request.camera_fov,
        request.camera_rot);
    const glm::mat4 ortho_proj = glm::ortho<float>(0, render_target_size.x, 0, render_target_size.y, -1, 1);

    request.prepare_render_target();
    request.set_blend_enabled(false);

    const auto pass_orchestration = plan_legacy_node_stroke_preview_pass_orchestration(
        make_legacy_node_stroke_preview_pass_orchestration_request(
            request.render_device_features,
            render_target_size,
            *request.brush,
            ortho_proj));
    request.setup_stroke_shader(pass_orchestration.stroke_shader);
    const bool sequence_ok = execute_legacy_node_stroke_preview_live_render_passes(
        LegacyNodeStrokePreviewLiveRenderRequest {
            .brush = *request.brush,
            .pass_orchestration = pass_orchestration,
            .prepared_strokes = prepared_strokes,
            .stroke_texture = request.preview_stroke_texture,
            .mixer_rtt = request.preview_rtt_mixer,
            .render_target = request.preview_rtt,
            .background_texture = request.preview_background_texture,
            .dual_texture = request.preview_dual_texture,
            .preview_texture = request.preview_image_texture,
            .linear_sampler = request.linear_sampler,
            .repeat_sampler = request.repeat_sampler,
            .zoom = request.zoom,
            .min_flow = request.min_flow,
            .copy_stroke_destination = pass_orchestration.copy_stroke_destination,
            .size = render_target_size,
            .bind_dual_pass_textures = request.bind_dual_pass_textures,
            .capture_background = request.capture_background,
            .compute_frames = request.compute_frames,
            .draw_samples = request.draw_samples,
            .draw_mix = request.draw_mix,
            .unbind_mixer_texture = request.unbind_mixer_texture,
            .bind_pattern_texture = request.bind_pattern_texture,
            .draw_composite = request.draw_composite,
        });
    request.finish_render_target();
    return sequence_ok;
}

void initialize_legacy_node_stroke_preview_clone(Node* dest)
{
    static_cast<NodeStrokePreview*>(dest)->init_controls();
}

bool execute_legacy_node_stroke_preview_immediate_draw(NodeStrokePreview& preview)
{
    if (preview.m_size.x == 0 || preview.m_size.y == 0) {
        return false;
    }

    const bool sequence_ok = pp::panopainter::execute_legacy_node_stroke_preview_immediate_draw(
        pp::panopainter::LegacyNodeStrokePreviewImmediateDrawRequest {
            .brush = preview.m_brush,
            .preview_size = preview.m_size,
            .zoom = preview.root()->m_zoom,
            .min_flow = preview.m_min_flow,
            .stroke_max_size_override = preview.m_max_size,
            .pad_override = preview.m_pad_override,
            .camera_fov = Canvas::I->m_cam_fov,
            .camera_rot = Canvas::I->m_cam_rot,
            .render_device_features = pp::panopainter::stroke_preview_render_device_features(),
            .preview_rtt = preview.m_rtt,
            .preview_rtt_mixer = preview.m_rtt_mixer,
            .preview_stroke_texture = preview.m_tex,
            .preview_dual_texture = preview.m_tex_dual,
            .preview_background_texture = preview.m_tex_background,
            .preview_image_texture = preview.m_tex_preview,
            .linear_sampler = preview.m_sampler_linear,
            .repeat_sampler = preview.m_sampler_linear_repeat,
            .prepare_render_target = [&] {
                pp::panopainter::apply_legacy_node_stroke_preview_viewport(
                    0,
                    0,
                    preview.m_rtt.getWidth(),
                    preview.m_rtt.getHeight());
                preview.m_rtt.bindFramebuffer();
                preview.m_rtt.clear();
                pp::panopainter::bind_legacy_node_stroke_preview_live_samplers(
                    preview.m_sampler_mipmap,
                    preview.m_sampler_linear,
                    preview.m_sampler_linear_repeat);
            },
            .finish_render_target = [&] {
                preview.m_rtt.unbindFramebuffer();
            },
            .compute_frames = [&](const Stroke& stroke, float frame_zoom) {
                return plan_legacy_node_stroke_preview_stroke_frames(
                    stroke,
                    frame_zoom,
                    glm::vec2(preview.m_rtt.getWidth(), preview.m_rtt.getHeight()));
            },
            .draw_samples = [&](std::array<vertex_t, 4>& shapes, Texture2D& texture, bool copy_stroke_destination) {
                return pp::panopainter::execute_legacy_node_stroke_preview_sample_pass(
                    preview.m_rtt,
                    shapes,
                    preview.m_brush_shape,
                    texture,
                    copy_stroke_destination);
            },
            .draw_mix = [&](const glm::vec2& bb_min, const glm::vec2& bb_sz) {
                pp::panopainter::execute_legacy_node_stroke_preview_mix_pass(
                    *preview.m_brush,
                    preview.m_size,
                    preview.m_rtt_mixer,
                    bb_min,
                    bb_sz,
                    preview.m_sampler_linear,
                    preview.m_tex_background,
                    preview.m_tex,
                    preview.m_tex_dual,
                    [&] {
                        preview.m_plane.draw_fill();
                    });
            },
            .draw_checkerboard = [&] {
                preview.m_plane.draw_fill();
            },
            .draw_composite = [&] {
                preview.m_plane.draw_fill();
            },
        });
    assert(sequence_ok);
    return sequence_ok;
}

} // namespace pp::panopainter

std::atomic_int NodeStrokePreview::s_instances{ 0 };
std::atomic_bool NodeStrokePreview::s_running{ false };
std::mutex NodeStrokePreview::s_render_mutex;
BlockingQueue<std::shared_ptr<NodeStrokePreview>> NodeStrokePreview::s_queue;

RTT NodeStrokePreview::m_rtt;
RTT NodeStrokePreview::m_rtt_mixer;
Texture2D NodeStrokePreview::m_tex; // blending tmp texture
Texture2D NodeStrokePreview::m_tex_dual;
Texture2D NodeStrokePreview::m_tex_background;
Sampler NodeStrokePreview::m_sampler_linear;
Sampler NodeStrokePreview::m_sampler_linear_repeat;
Sampler NodeStrokePreview::m_sampler_mipmap;
DynamicShape NodeStrokePreview::m_brush_shape;
std::jthread NodeStrokePreview::s_renderer;

void NodeStrokePreview::terminate_renderer()
{
    if (!s_renderer.joinable())
        return;

    s_running = false;
    s_renderer.request_stop();
    s_queue.UnlockGetters();
    s_renderer.join();
}

void NodeStrokePreview::empty_queue()
{
    s_queue.q.clear();
}

void NodeStrokePreview::restore_context()
{
    NodeBorder::restore_context();
    init_controls();
    if (m_size.x > 0 && m_size.y > 0)
        m_tex_preview.create(static_cast<int>(m_size.x), static_cast<int>(m_size.y));
    draw_stroke();
}

void NodeStrokePreview::clear_context()
{
    NodeBorder::clear_context();
    m_tex_preview.destroy();
}

Image NodeStrokePreview::render_to_image()
{
    std::lock_guard<std::mutex> _lock(s_render_mutex);

    App::I->render_task([this] {
        draw_stroke_immediate();
    });
    return m_tex_preview.get_image();
}

void NodeStrokePreview::draw_stroke()
{
    if (m_size.x == 0 || m_size.y == 0)
        return;
    std::unique_lock<std::mutex> queue_lock(s_queue.mutex);
    if (!s_renderer.joinable())
    {
        s_running = true;
        s_renderer = std::jthread([](std::stop_token stop_token) {
            BT_SetTerminate();

            m_sampler_linear.create();
            m_sampler_linear_repeat.create(
                pp::renderer::gl::linear_texture_filter(),
                pp::renderer::gl::repeat_texture_wrap());
            m_sampler_mipmap.create();
            m_sampler_mipmap.set_filter(
                pp::renderer::gl::linear_mipmap_linear_texture_filter(),
                pp::renderer::gl::linear_texture_filter());
            m_brush_shape.create();
            while (s_running && !stop_token.stop_requested())
            {
                auto node = s_queue.Get();
                if (node)
                {
                    std::lock_guard<std::mutex> _lock(s_render_mutex);

                    // if the brush is not already loaded, load it and then destroy it
                    bool to_unload = (node->m_brush->m_tip_texture == nullptr);
                    node->m_brush->preload();

                    App::I->render_task([node, to_unload]
                    {
                        gl_state gl;
                        gl.save();

                        node->m_brush->load();
                        node->draw_stroke_immediate();
                        if (to_unload)
                            node->m_brush->unload();

                        gl.restore();
                    });
                    
                    node->app_redraw();

                    //std::this_thread::sleep_for(std::chrono::milliseconds(30));
                    std::this_thread::yield();
                }
            }
            m_rtt.destroy();
            m_rtt_mixer.destroy();
            m_tex.destroy();
            m_tex_dual.destroy();
            m_tex_background.destroy();
            m_brush_shape.destroy();
            s_running = false;
        });
    }
    queue_lock.unlock();
    s_queue.PostUnique(std::static_pointer_cast<NodeStrokePreview>(shared_from_this()), m_draw_first);
}

void NodeStrokePreview::draw()
{
    pp::panopainter::setup_legacy_canvas_draw_merge_texture_shader(
        pp::panopainter::LegacyCanvasDrawMergeTextureUniforms {
            .mvp = m_mvp,
            .texture_slot = 0,
        });
    m_tex_preview.bind();
    m_sampler_linear.bind(0);
    m_plane.draw_fill();
    m_sampler_linear.unbind();
    m_tex_preview.unbind();
}

void NodeStrokePreview::handle_resize(glm::vec2 old_size, glm::vec2 new_size, float zoom)
{
    if (m_preview_size == (new_size * root()->m_zoom) || !m_brush)
        return;

    m_preview_size = new_size * root()->m_zoom;

    if (m_on_screen)
        draw_stroke();
}

void NodeStrokePreview::destroy()
{
    m_tex_preview.destroy();
    Node::destroy();
}

void NodeStrokePreview::handle_on_screen(bool old_visibility, bool new_visibility)
{
    parent::handle_on_screen(old_visibility, new_visibility);
    if (new_visibility)
    {
        draw_stroke();
    }
    else
    {
        s_queue.Remove(std::static_pointer_cast<NodeStrokePreview>(shared_from_this()));
        m_tex_preview.destroy();
    }
}