#version 450

layout(local_size_x = 16, local_size_y = 16) in;

layout(set = 0, binding = 0) uniform sampler2D rgb_in;
layout(set = 0, binding = 1) uniform writeonly image2D y_out;
layout(set = 0, binding = 2) uniform writeonly image2D uv_out;
layout(set = 0, binding = 3) uniform sampler2D cursor_in;

layout(push_constant) uniform PushConstants {
  vec4 color_vec_y;
  vec4 color_vec_u;
  vec4 color_vec_v;
  vec2 range_y;
  vec2 range_uv;
  ivec2 src_offset;
  ivec2 src_size;
  ivec2 dst_offset;  // top-left of the aspect-preserved viewport
  ivec2 dst_size;  // effective viewport (src after aspect-fit)
  ivec2 dst_full_size;  // full encoder frame size
  ivec2 cursor_pos;
  ivec2 cursor_size;  // w=0 means no cursor
  int y_invert;
} pc;

vec3 blend_cursor(vec3 rgb, ivec2 pos) {
  ivec2 cp = pos - pc.cursor_pos;
  if (cp.x >= 0 && cp.y >= 0 && cp.x < pc.cursor_size.x && cp.y < pc.cursor_size.y) {
    vec4 c = texture(cursor_in, (vec2(cp) + 0.5) / vec2(pc.cursor_size));
    rgb = mix(rgb, c.bgr, c.a);
  }
  return rgb;
}

void main() {
  ivec2 pos = ivec2(gl_GlobalInvocationID.xy);
  if (pos.x >= pc.dst_full_size.x || pos.y >= pc.dst_full_size.y)
    return;

  ivec2 d = pos - pc.dst_offset;
  if (d.x < 0 || d.y < 0 || d.x >= pc.dst_size.x || d.y >= pc.dst_size.y) {
    // Outside the aspect-fit viewport: write encoded black.
    // Use the same conversion formula as RGB=(0,0,0), where the dot product
    // is zero and the encoded value is color_vec_*.w adjusted by range_*.
    float y_black = pc.color_vec_y.w * pc.range_y.x + pc.range_y.y;
    float u_black = pc.color_vec_u.w * pc.range_uv.x + pc.range_uv.y;
    float v_black = pc.color_vec_v.w * pc.range_uv.x + pc.range_uv.y;
    imageStore(y_out, pos, vec4(y_black, 0, 0, 0));
    if ((pos.x & 1) == 0 && (pos.y & 1) == 0)
      imageStore(uv_out, pos >> 1, vec4(u_black, v_black, 0, 0));
    return;
  }

  vec2 inv_tex = 1.0 / vec2(textureSize(rgb_in, 0));
  vec2 scale = vec2(pc.src_size) / vec2(pc.dst_size);

  vec2 uv = (vec2(pc.src_offset) + (vec2(d) + 0.5) * scale) * inv_tex;
  if (pc.y_invert != 0)
    uv.y = 1.0 - uv.y;
  vec3 rgb = texture(rgb_in, uv).rgb;

  if (pc.cursor_size.x > 0)
    rgb = blend_cursor(rgb, pos);

  // Y plane
  float y = dot(pc.color_vec_y.xyz, rgb) + pc.color_vec_y.w;
  imageStore(y_out, pos, vec4(y * pc.range_y.x + pc.range_y.y, 0, 0, 0));

  // UV plane (half resolution, one thread per 2x2 block)
  if ((pos.x & 1) == 0 && (pos.y & 1) == 0) {
    vec2 step = scale * inv_tex;
    if (pc.y_invert != 0)
      step.y = -step.y;

    vec3 rgb_r  = texture(rgb_in, uv + vec2(step.x, 0)).rgb;
    vec3 rgb_b  = texture(rgb_in, uv + vec2(0, step.y)).rgb;
    vec3 rgb_br = texture(rgb_in, uv + vec2(step.x, step.y)).rgb;
    if (pc.cursor_size.x > 0) {
      rgb_r  = blend_cursor(rgb_r,  pos + ivec2(1, 0));
      rgb_b  = blend_cursor(rgb_b,  pos + ivec2(0, 1));
      rgb_br = blend_cursor(rgb_br, pos + ivec2(1, 1));
    }

    vec3 avg = (rgb + rgb_r + rgb_b + rgb_br) * 0.25;

    float cb = dot(pc.color_vec_u.xyz, avg) + pc.color_vec_u.w;
    float cr = dot(pc.color_vec_v.xyz, avg) + pc.color_vec_v.w;

    imageStore(uv_out, pos >> 1, vec4(cb * pc.range_uv.x + pc.range_uv.y,
                                      cr * pc.range_uv.x + pc.range_uv.y, 0, 0));
  }
}