#version 450

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

layout(binding = 0) uniform sampler2D video_y;
layout(binding = 1) uniform sampler2D video_u;
layout(binding = 2) uniform sampler2D video_v;
layout(binding = 3) uniform sampler2D overlay_tex;
layout(binding = 4, rgba8) uniform writeonly image2D output_img;

layout(push_constant) uniform PushConstants {
    ivec2 output_size;
    ivec2 overlay_origin;
    ivec2 overlay_size;
    int chroma_mode;
} push_constants;

vec3 yuv420p_to_rgb(vec2 uv) {
    if (push_constants.chroma_mode == 2) {
        return texture(video_y, uv).rgb;
    }
    float y = texture(video_y, uv).r;
    float u;
    float v;
    if (push_constants.chroma_mode == 1) {
        vec2 uv_pair = texture(video_u, uv).rg;
        u = uv_pair.r - 0.5;
        v = uv_pair.g - 0.5;
    } else {
        u = texture(video_u, uv).r - 0.5;
        v = texture(video_v, uv).r - 0.5;
    }

    float yy = 1.16438356 * max(y - 0.0625, 0.0);
    float r = yy + 1.79274107 * v;
    float g = yy - 0.21324861 * u - 0.53290933 * v;
    float b = yy + 2.11240179 * u;
    return clamp(vec3(r, g, b), 0.0, 1.0);
}

void main() {
    ivec2 pixel = ivec2(gl_GlobalInvocationID.xy);
    if (pixel.x >= push_constants.output_size.x || pixel.y >= push_constants.output_size.y) {
        return;
    }

    vec2 uv = (vec2(pixel) + vec2(0.5)) / vec2(push_constants.output_size);
    vec3 base_rgb = yuv420p_to_rgb(uv);
    vec4 out_color = vec4(base_rgb, 1.0);

    ivec2 overlay_pixel = pixel - push_constants.overlay_origin;
    if (overlay_pixel.x >= 0 && overlay_pixel.y >= 0 &&
        overlay_pixel.x < push_constants.overlay_size.x &&
        overlay_pixel.y < push_constants.overlay_size.y) {
        vec2 overlay_uv = (vec2(overlay_pixel) + vec2(0.5)) / vec2(push_constants.overlay_size);
        vec4 overlay = texture(overlay_tex, overlay_uv);
        out_color.rgb = mix(out_color.rgb, overlay.rgb, overlay.a);
    }

    imageStore(output_img, pixel, out_color);
}