/* SPDX-FileCopyrightText: 2024 Blender Authors
 *
 * SPDX-License-Identifier: GPL-2.0-or-later */

#include "BLI_color_types.hh"
#include "BLI_math_geom.h"
#include "BLI_math_matrix.hh"
#include "BLI_math_vector.hh"

#include "BKE_attribute.hh"
#include "BKE_camera.h"
#include "BKE_curves.hh"
#include "BKE_grease_pencil.hh"
#include "BKE_image.hh"
#include "BKE_material.hh"

#include "DNA_gpencil_legacy_types.h"
#include "DNA_material_types.h"
#include "DNA_object_types.h"
#include "DNA_scene_types.h"
#include "DNA_userdef_types.h"
#include "DNA_view3d_types.h"

#include "ED_grease_pencil.hh"
#include "ED_view3d.hh"

#include "IMB_imbuf.hh"
#include "IMB_imbuf_types.hh"

#include "GPU_debug.hh"
#include "GPU_framebuffer.hh"
#include "GPU_immediate.hh"
#include "GPU_matrix.hh"
#include "GPU_primitive.hh"
#include "GPU_shader_builtin.hh"
#include "GPU_shader_shared.hh"
#include "GPU_state.hh"
#include "GPU_texture.hh"
#include "GPU_uniform_buffer.hh"
#include "GPU_vertex_format.hh"

namespace blender::ed::greasepencil::image_render {

/* Enable GPU debug capture (needs WITH_RENDERDOC option). */
constexpr const bool enable_debug_gpu_capture = true;

RegionViewData region_init(ARegion &region, const int2 &win_size)
{
  RegionView3D &rv3d = *static_cast<RegionView3D *>(region.regiondata);

  const RegionViewData data = {
      int2{region.winx, region.winy}, region.winrct, ED_view3d_mats_rv3d_backup(&rv3d)};

  /* Resize region. */
  region.winrct.xmin = 0;
  region.winrct.ymin = 0;
  region.winrct.xmax = win_size.x;
  region.winrct.ymax = win_size.y;
  region.winx = short(win_size.x);
  region.winy = short(win_size.y);

  return data;
}

void region_reset(ARegion &region, const RegionViewData &data)
{
  RegionView3D &rv3d = *static_cast<RegionView3D *>(region.regiondata);

  region.winx = data.winsize.x;
  region.winy = data.winsize.y;
  region.winrct = data.winrct;

  ED_view3d_mats_rv3d_restore(&rv3d, data.rv3d_store);
  ED_view3D_mats_rv3d_free(data.rv3d_store);
}

GPUOffScreen *image_render_begin(const int2 &win_size)
{
  if (enable_debug_gpu_capture) {
    GPU_debug_capture_begin("Grease Pencil Image Render");
  }

  char err_out[256] = "unknown";
  GPUOffScreen *offscreen = GPU_offscreen_create(win_size.x,
                                                 win_size.y,
                                                 true,
                                                 gpu::TextureFormat::UNORM_8_8_8_8,
                                                 GPU_TEXTURE_USAGE_HOST_READ,
                                                 false,
                                                 err_out);
  if (offscreen == nullptr) {
    return nullptr;
  }

  GPU_offscreen_bind(offscreen, true);

  GPU_matrix_push_projection();
  GPU_matrix_identity_projection_set();
  GPU_matrix_push();
  GPU_matrix_identity_set();

  GPU_clear_color(0.0f, 0.0f, 0.0f, 0.0f);
  GPU_clear_depth(1.0f);

  return offscreen;
}

Image *image_render_end(Main &bmain, GPUOffScreen *buffer)
{
  GPU_matrix_pop_projection();
  GPU_matrix_pop();

  const int2 win_size = {GPU_offscreen_width(buffer), GPU_offscreen_height(buffer)};
  ImBuf *ibuf = IMB_allocImBuf(win_size.x, win_size.y, ImBufFlags::ByteData);
  GPU_offscreen_read_color(buffer, GPU_DATA_UBYTE, ibuf->byte_data_for_write());

  Image *ima = BKE_image_add_from_imbuf(&bmain, ibuf, "Grease Pencil Fill");
  ima->id.tag |= ID_TAG_DOIT;

  BKE_image_release_ibuf(ima, ibuf, nullptr);

  /* Switch back to regular frame-buffer. */
  GPU_offscreen_unbind(buffer, true);
  GPU_offscreen_free(buffer);

  if (enable_debug_gpu_capture) {
    GPU_debug_capture_end();
  }

  return ima;
}

void compute_view_matrices(const ViewContext &view_context,
                           const Scene &scene,
                           const int2 &win_size,
                           const float2 &zoom,
                           const float2 &offset)
{
  rctf viewplane;
  float clip_start, clip_end;
  const bool is_ortho = ED_view3d_viewplane_get(view_context.depsgraph,
                                                view_context.v3d,
                                                view_context.rv3d,
                                                win_size.x,
                                                win_size.y,
                                                &viewplane,
                                                &clip_start,
                                                &clip_end,
                                                nullptr);

  /* Rescale `viewplane` to fit all strokes. */
  const float2 view_min = float2(viewplane.xmin, viewplane.ymin);
  const float2 view_max = float2(viewplane.xmax, viewplane.ymax);
  const float2 view_extent = view_max - view_min;
  const float2 view_center = 0.5f * (view_max + view_min);
  const float2 offset_abs = offset * view_extent;
  const float2 view_min_new = (view_min - view_center) * zoom + view_center + offset_abs;
  const float2 view_max_new = (view_max - view_center) * zoom + view_center + offset_abs;
  viewplane.xmin = view_min_new.x;
  viewplane.ymin = view_min_new.y;
  viewplane.xmax = view_max_new.x;
  viewplane.ymax = view_max_new.y;

  float4x4 winmat;
  if (is_ortho) {
    orthographic_m4(winmat.ptr(),
                    viewplane.xmin,
                    viewplane.xmax,
                    viewplane.ymin,
                    viewplane.ymax,
                    -clip_end,
                    clip_end);
  }
  else {
    perspective_m4(winmat.ptr(),
                   viewplane.xmin,
                   viewplane.xmax,
                   viewplane.ymin,
                   viewplane.ymax,
                   clip_start,
                   clip_end);
  }

  ED_view3d_update_viewmat(view_context.depsgraph,
                           &scene,
                           view_context.v3d,
                           view_context.region,
                           nullptr,
                           winmat.ptr(),
                           nullptr,
                           true);
}

void set_view_matrix(const RegionView3D &rv3d)
{
  GPU_matrix_set(rv3d.viewmat);
}

void clear_view_matrix()
{
  GPU_matrix_identity_set();
}

void set_projection_matrix(const RegionView3D &rv3d)
{
  GPU_matrix_projection_set(rv3d.winmat);
}

void clear_projection_matrix()
{
  GPU_matrix_identity_projection_set();
}

void draw_dot(const float4x4 &transform,
              const float3 &position,
              const float point_size,
              const ColorGeometry4f &color)
{
  GPUVertFormat *format = immVertexFormat();
  uint attr_pos = GPU_vertformat_attr_add(format, "pos", gpu::VertAttrType::SFLOAT_32_32_32);
  uint attr_size = GPU_vertformat_attr_add(format, "size", gpu::VertAttrType::SFLOAT_32);
  uint attr_color = GPU_vertformat_attr_add(
      format, "color", gpu::VertAttrType::SFLOAT_32_32_32_32);

  GPU_program_point_size(true);
  immBindBuiltinProgram(GPU_SHADER_3D_POINT_VARYING_SIZE_VARYING_COLOR);
  immBegin(GPU_PRIM_POINTS, 1);
  immAttr1f(attr_size, point_size * M_SQRT2);
  immAttr4fv(attr_color, color);
  immVertex3fv(attr_pos, math::transform_point(transform, position));
  immEnd();
  immUnbindProgram();
  GPU_program_point_size(false);
}

void draw_polyline(const float4x4 &transform,
                   const IndexRange indices,
                   Span<float3> positions,
                   const VArray<ColorGeometry4f> &colors,
                   const bool cyclic,
                   const float line_width)
{
  GPUVertFormat *format = immVertexFormat();
  const uint attr_pos = GPU_vertformat_attr_add(format, "pos", gpu::VertAttrType::SFLOAT_32_32_32);
  const uint attr_color = GPU_vertformat_attr_add(
      format, "color", gpu::VertAttrType::SFLOAT_32_32_32_32);
  immBindBuiltinProgram(GPU_SHADER_3D_FLAT_COLOR);

  GPU_line_width(line_width);
  /* If cyclic the curve needs one more vertex. */
  const int cyclic_add = (cyclic && indices.size() > 2) ? 1 : 0;
  immBeginAtMost(GPU_PRIM_LINE_STRIP, indices.size() + cyclic_add);

  for (const int point_i : indices) {
    immAttr4fv(attr_color, colors[point_i]);
    immVertex3fv(attr_pos, math::transform_point(transform, positions[point_i]));
  }

  if (cyclic && indices.size() > 2) {
    const int point_i = indices[0];
    immAttr4fv(attr_color, colors[point_i]);
    immVertex3fv(attr_pos, math::transform_point(transform, positions[point_i]));
  }

  immEnd();
  immUnbindProgram();
}

static gpu::UniformBuf *create_shader_ubo(const RegionView3D &rv3d,
                                          const int2 &win_size,
                                          const Object &object,
                                          const eGPDstroke_Caps cap_start,
                                          const eGPDstroke_Caps cap_end,
                                          const bool is_fill_stroke)
{
  GPencilStrokeData data;
  copy_v2_v2(data.viewport, float2(win_size));
  data.pixsize = rv3d.pixsize;
  data.objscale = math::average(float3(object.scale));
  data.caps_start = cap_start;
  data.caps_end = cap_end;
  data.fill_stroke = is_fill_stroke;

  return GPU_uniformbuf_create_ex(sizeof(GPencilStrokeData), &data, __func__);
}

static void draw_grease_pencil_stroke(const float4x4 &transform,
                                      const RegionView3D &rv3d,
                                      const int2 &win_size,
                                      const Object &object,
                                      const IndexRange indices,
                                      Span<float3> positions,
                                      const VArray<float> &radii,
                                      const VArray<ColorGeometry4f> &colors,
                                      const bool cyclic,
                                      const eGPDstroke_Caps cap_start,
                                      const eGPDstroke_Caps cap_end,
                                      const bool fill_stroke,
                                      const float radius_scale)
{
  if (indices.is_empty()) {
    return;
  }

  GPUVertFormat *format = immVertexFormat();
  /* Format is matching shader manual load. Keep in sync with #GreasePencilStrokeData.
   * Only the name of the first attribute is important. */
  const uint attr_pos = GPU_vertformat_attr_add(
      format, "gp_vert_data", gpu::VertAttrType::SFLOAT_32_32_32);
  const uint attr_thickness = GPU_vertformat_attr_add(
      format, "thickness", gpu::VertAttrType::SFLOAT_32);
  const uint attr_color = GPU_vertformat_attr_add(
      format, "color", gpu::VertAttrType::SFLOAT_32_32_32_32);

  immBindBuiltinProgram(GPU_SHADER_GPENCIL_STROKE);
  gpu::UniformBuf *ubo = create_shader_ubo(
      rv3d, win_size, object, cap_start, cap_end, fill_stroke);
  immBindUniformBuf("gpencil_stroke_data", ubo);

  /* If cyclic the curve needs one more vertex. */
  const int cyclic_add = (cyclic && indices.size() > 2) ? 1 : 0;

  gpu::Batch *batch = immBeginBatchAtMost(GPU_PRIM_LINE_STRIP_ADJ,
                                          indices.size() + cyclic_add + 2);

  auto draw_point = [&](const int point_i) {
    const float thickness = 2.0f * radii[point_i] * radius_scale;

    immAttr4fv(attr_color, colors[point_i]);
    immAttr1f(attr_thickness, std::max(thickness, 0.0f));
    immVertex3fv(attr_pos, math::transform_point(transform, positions[point_i]));
  };

  /* First point for adjacency (not drawn). */
  if (cyclic && indices.size() > 2) {
    draw_point(indices.last() - 1);
  }
  else {
    draw_point(indices.first() + 1);
  }

  for (const int point_i : indices) {
    draw_point(point_i);
  }

  if (cyclic && indices.size() > 2) {
    draw_point(indices.first());
    draw_point(indices.first() + 1);
  }
  /* Last adjacency point (not drawn). */
  else {
    draw_point(indices.last() - 1);
  }

  immEnd();

  /* Expanded `drawcall`. */
  GPUPrimType expand_prim_type = GPUPrimType::GPU_PRIM_TRIS;
  /* Hard-coded in shader. */
  const uint expand_prim_len = 12;
  /* Do not count adjacency info for start and end primitives. */
  const uint final_vert_len = ((batch->vertex_count_get() - 2) * expand_prim_len) * 3;

  if (final_vert_len > 0) {
    GPU_batch_bind_as_resources(batch, batch->shader);

    /* TODO(fclem): get rid of this dummy VBO. */
    GPUVertFormat format = {0};
    GPU_vertformat_attr_add(&format, "dummy", gpu::VertAttrType::SFLOAT_32);
    gpu::VertBuf *vbo = GPU_vertbuf_create_with_format(format);
    GPU_vertbuf_data_alloc(*vbo, 1);

    gpu::Batch *gpu_batch = GPU_batch_create_ex(
        expand_prim_type, vbo, nullptr, GPU_BATCH_OWNS_VBO);

    GPU_batch_set_shader(gpu_batch, batch->shader);
    GPU_batch_draw_advanced(gpu_batch, 0, final_vert_len, 0, 1);

    GPU_batch_discard(gpu_batch);
  }
  GPU_batch_discard(batch);

  immUnbindProgram();

  GPU_uniformbuf_free(ubo);
}

static void draw_dots(const float4x4 &transform,
                      const RegionView3D &rv3d,
                      const Object &object,
                      const IndexRange indices,
                      Span<float3> positions,
                      const VArray<float> &radii,
                      const VArray<ColorGeometry4f> &colors,
                      const float radius_scale)
{
  if (indices.is_empty()) {
    return;
  }

  GPUVertFormat *format = immVertexFormat();
  const uint attr_pos = GPU_vertformat_attr_add(format, "pos", gpu::VertAttrType::SFLOAT_32_32_32);
  const uint attr_size = GPU_vertformat_attr_add(format, "size", gpu::VertAttrType::SFLOAT_32);
  const uint attr_color = GPU_vertformat_attr_add(
      format, "color", gpu::VertAttrType::SFLOAT_32_32_32_32);

  immBindBuiltinProgram(GPU_SHADER_3D_POINT_VARYING_SIZE_VARYING_COLOR);
  GPU_program_point_size(true);

  immBegin(GPU_PRIM_POINTS, indices.size());

  /* Includes viewport zoom factor and perspective projection. The point shader does not include
   * these factors internally, unlike the line shader. */
  const float objscale = math::average(float3(object.scale));
  const float pixel_scale = 4.0f * radius_scale * objscale / rv3d.pixsize;
  for (const int point_i : indices) {
    const float3 &position = positions[point_i];
    float perspective_factor = 1.0f;
    if (rv3d.is_persp) {
      const float3 view_position = math::transform_point(float4x4(rv3d.persmat), position);
      perspective_factor = math::safe_rcp(view_position.z);
    }
    const float thickness = std::max(radii[point_i] * pixel_scale * perspective_factor, 1.0f);

    immAttr4fv(attr_color, colors[point_i]);
    immAttr1f(attr_size, thickness);
    immVertex3fv(attr_pos, math::transform_point(transform, position));
  }

  immEnd();
  immUnbindProgram();
  GPU_program_point_size(false);
}

void draw_circles(const float4x4 &transform,
                  const IndexRange indices,
                  Span<float3> centers,
                  const VArray<float> &radii,
                  const VArray<ColorGeometry4f> &colors,
                  const float2 &viewport_size,
                  const float line_width,
                  const bool fill)
{
  if (indices.is_empty()) {
    return;
  }

  constexpr const int segments_num = 32;
  static const float2 coords[] = {
      {1.0000f, 0.0000f},   {0.9808f, 0.1951f},   {0.9239f, 0.3827f},   {0.8315f, 0.5556f},
      {0.7071f, 0.7071f},   {0.5556f, 0.8315f},   {0.3827f, 0.9239f},   {0.1951f, 0.9808f},
      {0.0000f, 1.0000f},   {-0.1951f, 0.9808f},  {-0.3827f, 0.9239f},  {-0.5556f, 0.8315f},
      {-0.7071f, 0.7071f},  {-0.8315f, 0.5556f},  {-0.9239f, 0.3827f},  {-0.9808f, 0.1951f},
      {-1.0000f, 0.0000f},  {-0.9808f, -0.1951f}, {-0.9239f, -0.3827f}, {-0.8315f, -0.5556f},
      {-0.7071f, -0.7071f}, {-0.5556f, -0.8315f}, {-0.3827f, -0.9239f}, {-0.1951f, -0.9808f},
      {-0.0000f, -1.0000f}, {0.1951f, -0.9808f},  {0.3827f, -0.9239f},  {0.5556f, -0.8315f},
      {0.7071f, -0.7071f},  {0.8315f, -0.5556f},  {0.9239f, -0.3827f},  {0.9808f, -0.1951f},
  };

  GPUVertFormat *format = immVertexFormat();
  const uint attr_pos = GPU_vertformat_attr_add(format, "pos", gpu::VertAttrType::SFLOAT_32_32_32);
  const uint attr_color = GPU_vertformat_attr_add(
      format, "color", gpu::VertAttrType::SFLOAT_32_32_32_32);

  const float scale = math::average(math::to_scale(transform));

  if (fill) {
    immBindBuiltinProgram(GPU_SHADER_3D_FLAT_COLOR);

    for (const int point_i : indices) {
      const float radius = radii[point_i];
      const ColorGeometry4f color = colors[point_i];
      const float3 center = math::transform_point(transform, centers[point_i]);

      immBegin(GPU_PRIM_TRI_STRIP, segments_num);

      for (const int i : IndexRange(segments_num / 2)) {
        immAttr4fv(attr_color, color);
        immVertex3fv(attr_pos, center + float3(radius * scale * coords[i], 0.0f));
        if (segments_num - 1 - i > i) {
          immAttr4fv(attr_color, color);
          immVertex3fv(attr_pos,
                       center + float3(radius * scale * coords[segments_num - 1 - i], 0.0f));
        }
      }

      immEnd();
    }

    immUnbindProgram();
  }
  else {
    immBindBuiltinProgram(GPU_SHADER_3D_POLYLINE_FLAT_COLOR);

    immUniform2fv("viewportSize", viewport_size);
    immUniform1f("lineWidth", line_width * U.pixelsize);

    for (const int point_i : indices) {
      const float radius = radii[point_i];
      const ColorGeometry4f color = colors[point_i];
      const float3 center = math::transform_point(transform, centers[point_i]);

      immBegin(GPU_PRIM_LINE_STRIP, segments_num + 1);

      for (const int i : IndexRange(segments_num)) {
        immAttr4fv(attr_color, color);
        immVertex3fv(attr_pos, center + float3(radius * scale * coords[i], 0.0f));
      }
      immAttr4fv(attr_color, color);
      immVertex3fv(attr_pos, center + float3(radius * scale * coords[0], 0.0f));

      immEnd();
    }

    immUnbindProgram();
  }
}

void draw_lines(const float4x4 &transform,
                IndexRange indices,
                Span<float3> start_positions,
                Span<float3> end_positions,
                const VArray<ColorGeometry4f> &colors,
                float line_width)
{
  GPUVertFormat *format = immVertexFormat();
  const uint attr_pos = GPU_vertformat_attr_add(format, "pos", gpu::VertAttrType::SFLOAT_32_32_32);
  const uint attr_color = GPU_vertformat_attr_add(
      format, "color", gpu::VertAttrType::SFLOAT_32_32_32_32);
  immBindBuiltinProgram(GPU_SHADER_3D_FLAT_COLOR);

  GPU_line_width(line_width);
  immBeginAtMost(GPU_PRIM_LINES, 2 * indices.size());

  for (const int point_i : indices) {
    immAttr4fv(attr_color, colors[point_i]);
    immVertex3fv(attr_pos, math::transform_point(transform, start_positions[point_i]));

    immAttr4fv(attr_color, colors[point_i]);
    immVertex3fv(attr_pos, math::transform_point(transform, end_positions[point_i]));
  }

  immEnd();
  immUnbindProgram();
}

static GVArray attribute_interpolate(const GSpan &input, const bke::CurvesGeometry &curves)
{
  if (curves.is_single_type(CURVE_TYPE_POLY)) {
    return GVArray::from_span(input);
  }

  const CPPType &type = input.type();
  GArray<> out(type, curves.evaluated_points_num());
  curves.interpolate_to_evaluated(input, out.as_mutable_span());
  return GVArray::from_garray(std::move(out));
};

void draw_grease_pencil_strokes(const RegionView3D &rv3d,
                                const int2 &win_size,
                                const Object &object,
                                const bke::greasepencil::Drawing &drawing,
                                const float4x4 &transform,
                                const IndexMask &strokes_mask,
                                const VArray<ColorGeometry4f> &colors,
                                const bool use_xray,
                                const float radius_scale)
{
  set_view_matrix(rv3d);

  GPU_program_point_size(true);
  /* Do not write to depth (avoid self-occlusion). */
  const bool prev_depth_mask = GPU_depth_mask_get();
  GPU_depth_mask(false);
  if (!use_xray) {
    GPU_depth_test(GPU_DEPTH_LESS_EQUAL);
    /* First arg is normally rv3d->dist, but this isn't
     * available here and seems to work quite well without. */
    GPU_polygon_offset(1.0f, 1.0f);
  }

  const bke::CurvesGeometry &curves = drawing.strokes();
  const OffsetIndices points_by_curve = curves.evaluated_points_by_curve();
  const Span<float3> positions = curves.evaluated_positions();
  const bke::AttributeAccessor attributes = curves.attributes();
  const VArray<bool> cyclic = curves.cyclic();

  curves.ensure_can_interpolate_to_evaluated();

  const VArray<float> radii =
      attribute_interpolate(VArraySpan(drawing.radii()), curves).typed<float>();
  const VArray<ColorGeometry4f> eval_colors =
      attribute_interpolate(VArraySpan(colors), curves).typed<ColorGeometry4f>();

  const VArray<int8_t> stroke_start_caps = *attributes.lookup_or_default<int8_t>(
      "start_cap", bke::AttrDomain::Curve, GP_STROKE_CAP_ROUND);
  const VArray<int8_t> stroke_end_caps = *attributes.lookup_or_default<int8_t>(
      "end_cap", bke::AttrDomain::Curve, GP_STROKE_CAP_ROUND);
  const VArray<int> materials = *attributes.lookup_or_default<int>(
      "material_index", bke::AttrDomain::Curve, 0);

  /* Note: Serial loop since immediate mode drawing can't happen in worker
   * threads, has to be from the main thread. */
  strokes_mask.foreach_index([&](const int stroke_i) {
    /* Check if the color is visible. */
    const int material_index = materials[stroke_i];
    const Material *mat = BKE_object_material_get(const_cast<Object *>(&object),
                                                  material_index + 1);
    const eMaterialGPencilStyle_Mode stroke_mode = mat && mat->gp_style ?
                                                       eMaterialGPencilStyle_Mode(
                                                           mat->gp_style->mode) :
                                                       GP_MATERIAL_MODE_LINE;

    if (mat == nullptr || (mat->gp_style->flag & GP_MATERIAL_HIDE)) {
      return;
    }

    switch (eMaterialGPencilStyle_Mode(stroke_mode)) {
      case GP_MATERIAL_MODE_LINE:
        draw_grease_pencil_stroke(transform,
                                  rv3d,
                                  win_size,
                                  object,
                                  points_by_curve[stroke_i],
                                  positions,
                                  radii,
                                  eval_colors,
                                  cyclic[stroke_i],
                                  eGPDstroke_Caps(stroke_start_caps[stroke_i]),
                                  eGPDstroke_Caps(stroke_end_caps[stroke_i]),
                                  false,
                                  radius_scale);
        break;
      case GP_MATERIAL_MODE_DOT:
      case GP_MATERIAL_MODE_SQUARE:
        /* NOTE: Squares don't have their own shader, render as dots too. */
        draw_dots(transform,
                  rv3d,
                  object,
                  points_by_curve[stroke_i],
                  positions,
                  radii,
                  eval_colors,
                  radius_scale);
        break;
    }
  });

  if (!use_xray) {
    GPU_depth_test(GPU_DEPTH_NONE);

    GPU_polygon_offset(0.0f, 0.0f);
  }
  GPU_depth_mask(prev_depth_mask);
  GPU_program_point_size(false);
  clear_view_matrix();
}

}  // namespace blender::ed::greasepencil::image_render
