// Copyright (c) 2024 Epic Games Tools
// Licensed under the MIT license (https://opensource.org/license/mit/)

////////////////////////////////
//~ rjf: Pipeline Tables

@table(name, source, ilay_table)
R_D3D11_VShadTable:
{
  {Rect                r_d3d11_g_rect_shader_src               r_d3d11_g_rect_ilay_elements }
  {Blur                r_d3d11_g_blur_shader_src               0                            }
  {Mesh                r_d3d11_g_mesh_shader_src               r_d3d11_g_mesh_ilay_elements }
  {Geo3DComposite      r_d3d11_g_geo3dcomposite_shader_src     0                            }
  {Finalize            r_d3d11_g_finalize_shader_src           0                            }
}

@table(name, source)
R_D3D11_PShadTable:
{
  {Rect                r_d3d11_g_rect_shader_src              }
  {Blur                r_d3d11_g_blur_shader_src              }
  {Mesh                r_d3d11_g_mesh_shader_src              }
  {Geo3DComposite      r_d3d11_g_geo3dcomposite_shader_src    }
  {Finalize            r_d3d11_g_finalize_shader_src          }
}

@table(name)
R_D3D11_UniformTypeTable:
{
  {Rect}
  {Blur}
  {Mesh}
}

////////////////////////////////
//~ rjf: UI Rectangle Shaders

@embed_string r_d3d11_g_rect_shader_src:
"""
cbuffer Globals : register(b0)
{
  float2 viewport_size_px;
  float opacity;
  row_major float4x4 texture_sample_channel_map;
  float2 texture_t2d_size_px;
  row_major float3x3 xform;
  float2 xform_scale;
}

struct CPU2Vertex
{
  float4 dst_rect_px     : POS;
  float4 src_rect_px     : TEX;
  float4 color00         : COL0;
  float4 color01         : COL1;
  float4 color10         : COL2;
  float4 color11         : COL3;
  float4 corner_radii_px : CRAD;
  float4 style_params    : STY; // x: border_thickness_px, y: softness_px, z: omit_texture, w: unused
  uint vertex_id         : SV_VertexID;
};

struct Vertex2Pixel
{
  float4 position           : SV_POSITION;
  nointerpolation float2 rect_half_size_px  : PSIZE;
  float2 texcoord_pct       : TEX;
  float2 sdf_sample_pos     : SDF;
  float4 tint               : TINT;
  float corner_radius_px    : CRAD;
  nointerpolation float border_thickness_px : BTHC;
  nointerpolation float softness_px         : SFT;
  nointerpolation float omit_texture        : OTX;
};

Texture2D    main_t2d     : register(t0);
SamplerState main_sampler : register(s0);

float rect_sdf(float2 sample_pos, float2 rect_half_size, float r)
{
  return length(max(abs(sample_pos) - rect_half_size + r, 0.0)) - r;
}

float linear_from_srgb_f32(float x)
{
  return x < 0.0404482362771082 ? x / 12.92 : pow((x + 0.055) / 1.055, 2.4);
}

float4 linear_from_srgba(float4 v)
{
  float4 result = float4(linear_from_srgb_f32(v.x),
                         linear_from_srgb_f32(v.y),
                         linear_from_srgb_f32(v.z),
                         v.w);
  return result;
}

//- rjf: vertex shader

Vertex2Pixel
vs_main(CPU2Vertex cpu2vertex)
{
  //- rjf: unpack & xform rectangle src/dst vertices
  float2 dst_p0_px  = cpu2vertex.dst_rect_px.xy;
  float2 dst_p1_px  = cpu2vertex.dst_rect_px.zw;
  float2 src_p0_px  = cpu2vertex.src_rect_px.xy;
  float2 src_p1_px  = cpu2vertex.src_rect_px.zw;
  float2 dst_size_px = abs(dst_p1_px - dst_p0_px);
  
  //- rjf: unpack style params
  float border_thickness_px = cpu2vertex.style_params.x;
  float softness_px         = cpu2vertex.style_params.y;
  float omit_texture        = cpu2vertex.style_params.z;
  
  //- rjf: prep per-vertex arrays to sample from (p: position, t: texcoord, c: colorcoord, r: cornerradius)
  float2 dst_p_verts_px[] =
  {
    float2(dst_p0_px.x, dst_p1_px.y),
    float2(dst_p0_px.x, dst_p0_px.y),
    float2(dst_p1_px.x, dst_p1_px.y),
    float2(dst_p1_px.x, dst_p0_px.y),
  };
  float2 src_p_verts_px[] =
  {
    float2(src_p0_px.x, src_p1_px.y),
    float2(src_p0_px.x, src_p0_px.y),
    float2(src_p1_px.x, src_p1_px.y),
    float2(src_p1_px.x, src_p0_px.y),
  };
  float dst_r_verts_px[] =
  {
    cpu2vertex.corner_radii_px.y,
    cpu2vertex.corner_radii_px.x,
    cpu2vertex.corner_radii_px.w,
    cpu2vertex.corner_radii_px.z,
  };
  float4 src_color[] =
  {
    linear_from_srgba(cpu2vertex.color01),
    linear_from_srgba(cpu2vertex.color00),
    linear_from_srgba(cpu2vertex.color11),
    linear_from_srgba(cpu2vertex.color10),
  };
  float2 dst_verts_pct = float2((cpu2vertex.vertex_id >> 1) ? 1.f : 0.f,
                                (cpu2vertex.vertex_id & 1)  ? 0.f : 1.f);
  
  // rjf: fill vertex -> pixel data
  Vertex2Pixel vertex2pixel;
  {
    float2 xformed_pos = mul(xform, float3(dst_p_verts_px[cpu2vertex.vertex_id], 1.f)).xy;
    xformed_pos.y = viewport_size_px.y - xformed_pos.y;
    vertex2pixel.position.xy          = 2.f * xformed_pos/viewport_size_px - 1.f;
    vertex2pixel.position.z           = 0.f;
    vertex2pixel.position.w           = 1.f;
    vertex2pixel.rect_half_size_px    = dst_size_px / 2.f * xform_scale;
    vertex2pixel.texcoord_pct         = src_p_verts_px[cpu2vertex.vertex_id] / texture_t2d_size_px;
    vertex2pixel.sdf_sample_pos       = (2.f * dst_verts_pct - 1.f) * vertex2pixel.rect_half_size_px;
    vertex2pixel.tint                 = src_color[cpu2vertex.vertex_id];
    vertex2pixel.corner_radius_px     = dst_r_verts_px[cpu2vertex.vertex_id];
    vertex2pixel.border_thickness_px  = border_thickness_px;
    vertex2pixel.softness_px          = softness_px;
    vertex2pixel.omit_texture         = omit_texture;
  }
  return vertex2pixel;
}

//- rjf: pixel shader

float4
ps_main(Vertex2Pixel vertex2pixel) : SV_TARGET
{
  // rjf: blend corner colors to produce final tint
  float4 tint = vertex2pixel.tint;
  
  // rjf: sample texture
  float4 albedo_sample = float4(1, 1, 1, 1);
  if(vertex2pixel.omit_texture < 1)
  {
    albedo_sample = mul(main_t2d.Sample(main_sampler, vertex2pixel.texcoord_pct), texture_sample_channel_map);
    albedo_sample = linear_from_srgba(albedo_sample);
  }
  
  // rjf: determine SDF sample position
  float2 sdf_sample_pos = vertex2pixel.sdf_sample_pos;
  
  // rjf: sample for borders
  float border_sdf_t = 1;
  if(vertex2pixel.border_thickness_px > 0)
  {
    float border_sdf_s = rect_sdf(sdf_sample_pos,
                                  vertex2pixel.rect_half_size_px - float2(vertex2pixel.softness_px*2.f, vertex2pixel.softness_px*2.f) - vertex2pixel.border_thickness_px,
                                  max(vertex2pixel.corner_radius_px-vertex2pixel.border_thickness_px, 0));
    border_sdf_t = smoothstep(0, 2*vertex2pixel.softness_px, border_sdf_s);
  }
  if(border_sdf_t < 0.001f)
  {
    discard;
  }
  
  // rjf: sample for corners
  float corner_sdf_t = 1;
  if(vertex2pixel.corner_radius_px > 0 || vertex2pixel.softness_px > 0.75f)
  {
    float corner_sdf_s = rect_sdf(sdf_sample_pos,
                                  vertex2pixel.rect_half_size_px - float2(vertex2pixel.softness_px*2.f, vertex2pixel.softness_px*2.f),
                                  vertex2pixel.corner_radius_px);
    corner_sdf_t = 1-smoothstep(0, 2*vertex2pixel.softness_px, corner_sdf_s);
  }
  
  // rjf: form+return final color
  float4 final_color = albedo_sample;
  final_color *= tint;
  final_color.a *= opacity;
  final_color.a *= corner_sdf_t;
  final_color.a *= border_sdf_t;
  return final_color;
}
"""

////////////////////////////////
//~ rjf: Blur Shaders

@embed_string r_d3d11_g_blur_shader_src:
"""
cbuffer Globals : register(b0)
{
  float4 rect;
  float4 corner_radii_px;
  float2 direction;
  float2 viewport_size;
  uint blur_count;
}

cbuffer Kernel : register(b1)
{
  float4 kernel[32];
}

struct CPU2Vertex
{
  uint vertex_id         : SV_VertexID;
};

struct Vertex2Pixel
{
  float4 position        : SV_POSITION;
  float2 texcoord        : TEX;
  float2 sdf_sample_pos  : SDF;
  nointerpolation float2 rect_half_size : RHS;
  float corner_radius    : RAD;
};

Texture2D    stage_t2d     : register(t0);
SamplerState stage_sampler : register(s0);

float rect_sdf(float2 sample_pos, float2 rect_half_size, float r)
{
  return length(max(abs(sample_pos) - rect_half_size + r, 0.0)) - r;
}

//- rjf: vertex shader

Vertex2Pixel
vs_main(CPU2Vertex c2v)
{
  float2 vertex_positions__scrn[] =
  {
    rect.xw,
    rect.xy,
    rect.zw,
    rect.zy,
  };
  float corner_radii__px[] =
  {
    corner_radii_px.y,
    corner_radii_px.x,
    corner_radii_px.w,
    corner_radii_px.z,
  };
  float2 cornercoords__pct = float2(
                                    (c2v.vertex_id >> 1) ? 1.f : 0.f,
                                    (c2v.vertex_id & 1)  ? 0.f : 1.f);
  
  float2 vertex_position__pct = vertex_positions__scrn[c2v.vertex_id] / viewport_size;
  float2 vertex_position__scr = 2.f * vertex_position__pct - 1.f;
  
  float2 rect_half_size = float2((rect.z-rect.x)/2, (rect.w-rect.y)/2);
  
  Vertex2Pixel v2p;
  {
    v2p.position = float4(vertex_position__scr.x, -vertex_position__scr.y, 0.f, 1.f);
    v2p.texcoord = vertex_position__pct;
    v2p.sdf_sample_pos = (2.f * cornercoords__pct - 1.f) * rect_half_size;
    v2p.rect_half_size = rect_half_size - 2.f;
    v2p.corner_radius = corner_radii__px[c2v.vertex_id];
  }
  return v2p;
}

//- rjf: pixel shader

float4
ps_main(Vertex2Pixel v2p) : SV_TARGET
{
  // rjf: blend weighted texture samples into color
  float3 color = kernel[0].x * stage_t2d.Sample(stage_sampler, v2p.texcoord).rgb;
  
  for(uint i = 1; i < blur_count; i += 1)
  {
    float weight = kernel[i].x;
    float offset = kernel[i].y;
    color += weight * stage_t2d.Sample(stage_sampler, v2p.texcoord - offset * direction).rgb;
    color += weight * stage_t2d.Sample(stage_sampler, v2p.texcoord + offset * direction).rgb;
  }
  
  // rjf: sample for corners
  float corner_sdf_s = rect_sdf(v2p.sdf_sample_pos, v2p.rect_half_size, v2p.corner_radius);
  float corner_sdf_t = 1-smoothstep(0, 2, corner_sdf_s);
  
  // rjf: weight output color by sdf
  // this is doing alpha testing, leave blurring only where mostly opaque pixels are
  if (corner_sdf_t < 0.9f)
  {
    discard;
  }
  
  return float4(color, 1.f);
}
"""

////////////////////////////////
//~ rjf: Mesh Shaders

@embed_string r_d3d11_g_mesh_shader_src:
"""
cbuffer Uniforms : register(b0)
{
  row_major float4x4 xform;
}

struct CPU2Vertex
{
  float3 position : POS;
  float3 normal   : NOR;
  float2 texcoord : TEX;
  float3 color    : COL;
};

struct Vertex2Pixel
{
  float4 position : SV_POSITION;
  float2 texcoord : TEX;
  float4 color    : COL;
};

Vertex2Pixel vs_main(CPU2Vertex c2v)
{
  Vertex2Pixel v2p;
  v2p.position = mul(float4(c2v.position, 1.f), xform);
  v2p.texcoord = c2v.texcoord;
  v2p.color    = float4(c2v.color, 1.f);
  return v2p;
}

float4 ps_main(Vertex2Pixel v2p) : SV_TARGET
{
  return v2p.color;
}
""";

////////////////////////////////
//~ rjf: Geo3D Composition Shaders

@embed_string r_d3d11_g_geo3dcomposite_shader_src:
"""
struct CPU2Vertex
{
  uint vertex_id         : SV_VertexID;
};

struct Vertex2Pixel
{
  float4 position        : SV_POSITION;
  float2 texcoord        : TEX;
};

Texture2D    stage_t2d     : register(t0);
SamplerState stage_sampler : register(s0);

//- rjf: vertex shader

Vertex2Pixel
vs_main(CPU2Vertex c2v)
{
  float4 vertex_positions__modl[] =
  {
    float4(0, 0, 0, 1),
    float4(0, 1, 0, 1),
    float4(1, 0, 0, 1),
    float4(1, 1, 0, 1),
  };
  float4 vertex_position__modl = vertex_positions__modl[c2v.vertex_id];
  float4 vertex_position__clip = float4(2*vertex_position__modl.x - 1, 2*vertex_position__modl.y - 1, 0, 1);
  float2 texcoord              = float2(vertex_position__modl.x, vertex_position__modl.y);
  texcoord.y = 1-texcoord.y;
  Vertex2Pixel v2p;
  {
    v2p.position = vertex_position__clip;
    v2p.texcoord = texcoord;
  }
  return v2p;
}

//- rjf: pixel shader

float4
ps_main(Vertex2Pixel v2p) : SV_TARGET
{
  float4 final_color = stage_t2d.Sample(stage_sampler, v2p.texcoord);
  return final_color;
}
"""

////////////////////////////////
//~ rjf: Finalize Shaders

@embed_string r_d3d11_g_finalize_shader_src:
"""
struct CPU2Vertex
{
  uint vertex_id         : SV_VertexID;
};

struct Vertex2Pixel
{
  float4 position        : SV_POSITION;
  float2 texcoord        : TEX;
};

Texture2D    stage_t2d     : register(t0);
SamplerState stage_sampler : register(s0);

//- rjf: vertex shader

Vertex2Pixel
vs_main(CPU2Vertex c2v)
{
  float4 vertex_positions__modl[] =
  {
    float4(0, 0, 0, 1),
    float4(0, 1, 0, 1),
    float4(1, 0, 0, 1),
    float4(1, 1, 0, 1),
  };
  float4 vertex_position__modl = vertex_positions__modl[c2v.vertex_id];
  float4 vertex_position__clip = float4(2*vertex_position__modl.x - 1, 2*vertex_position__modl.y - 1, 0, 1);
  float2 texcoord              = float2(vertex_position__modl.x, vertex_position__modl.y);
  texcoord.y = 1-texcoord.y;
  Vertex2Pixel v2p;
  {
    v2p.position = vertex_position__clip;
    v2p.texcoord = texcoord;
  }
  return v2p;
}

//- rjf: pixel shader

float4
ps_main(Vertex2Pixel v2p) : SV_TARGET
{
  float4 final_color = stage_t2d.Sample(stage_sampler, v2p.texcoord);
  final_color.a = 1;
  return final_color;
}
"""

////////////////////////////////
//~ rjf: Table Generators

@enum R_D3D11_VShadKind:
{
  @expand(R_D3D11_VShadTable a) `$(a.name)`,
  COUNT,
}

@enum R_D3D11_PShadKind:
{
  @expand(R_D3D11_PShadTable a) `$(a.name)`,
  COUNT,
}

@enum R_D3D11_UniformTypeKind:
{
  @expand(R_D3D11_UniformTypeTable a) `$(a.name)`,
  COUNT,
}

@c_file @data(`String8*`)
r_d3d11_g_vshad_kind_source_table:
{
  @expand(R_D3D11_VShadTable a) `&$(a.source)`;
}

@c_file @data(String8)
r_d3d11_g_vshad_kind_source_name_table:
{
  @expand(R_D3D11_VShadTable a) `str8_lit_comp("$(a.source)")`;
}

@c_file @data(`D3D11_INPUT_ELEMENT_DESC *`)
r_d3d11_g_vshad_kind_elements_ptr_table:
{
  @expand(R_D3D11_VShadTable a) `$(a.ilay_table)`;
}

@c_file @data(U64)
r_d3d11_g_vshad_kind_elements_count_table:
{
  @expand(R_D3D11_VShadTable a) `$(a.ilay_table != 0 -> "ArrayCount("..a.ilay_table..")") $(a.ilay_table == 0 -> "0")`;
}

@c_file @data(`String8*`)
r_d3d11_g_pshad_kind_source_table:
{
  @expand(R_D3D11_PShadTable a) `&$(a.source)`;
}

@c_file @data(String8)
r_d3d11_g_pshad_kind_source_name_table:
{
  @expand(R_D3D11_PShadTable a) `str8_lit_comp("$(a.source)")`;
}

@c_file @data(U64)
r_d3d11_g_uniform_type_kind_size_table:
{
  @expand(R_D3D11_UniformTypeTable a) `sizeof(R_D3D11_Uniforms_$(a.name))`;
}