struct VertexInput {
    @location(0) position: vec3<f32>,
    @location(1) normal: vec3<f32>,
    @location(2) uv: vec2<f32>,
}

struct VertexOutput {
    @builtin(position) clip_position: vec4<f32>,
    @location(0) world_position: vec3<f32>,
    @location(1) world_normal: vec3<f32>,
}

struct Uniforms {
    model: mat4x4<f32>,
    view: mat4x4<f32>,
    projection: mat4x4<f32>,
}

@group(0) @binding(0)
var<uniform> uniforms: Uniforms;

@vertex
fn vs_main(input: VertexInput) -> VertexOutput {
    var output: VertexOutput;

    let world_pos = uniforms.model * vec4<f32>(input.position, 1.0);
    output.world_position = world_pos.xyz;
    output.world_normal = (uniforms.model * vec4<f32>(input.normal, 0.0)).xyz;
    output.clip_position = uniforms.projection * uniforms.view * world_pos;

    return output;
}

fn bayer_2x2_dither(value: f32, screen_pos: vec2<f32>) -> f32 {
    let pattern = array<f32, 4>(
        0.0/4.0, 2.0/4.0,
        3.0/4.0, 1.0/4.0
    );
    let x = i32(screen_pos.x) % 2;
    let y = i32(screen_pos.y) % 2;
    let index = y * 2 + x;
    return select(0.0, 1.0, value > pattern[index]);
}

fn bayer_4x4_dither(value: f32, screen_pos: vec2<f32>) -> f32 {
    let pattern = array<f32, 16>(
        0.0/16.0,  8.0/16.0,  2.0/16.0, 10.0/16.0,
        12.0/16.0, 4.0/16.0, 14.0/16.0,  6.0/16.0,
        3.0/16.0, 11.0/16.0,  1.0/16.0,  9.0/16.0,
        15.0/16.0, 7.0/16.0, 13.0/16.0,  5.0/16.0
    );
    let x = i32(screen_pos.x) % 4;
    let y = i32(screen_pos.y) % 4;
    let index = y * 4 + x;
    return select(0.0, 1.0, value > pattern[index]);
}

fn bayer_8x8_dither(value: f32, screen_pos: vec2<f32>) -> f32 {
    let pattern = array<f32, 64>(
        0.0/64.0,  32.0/64.0,  8.0/64.0,  40.0/64.0,  2.0/64.0,  34.0/64.0, 10.0/64.0, 42.0/64.0,
        48.0/64.0, 16.0/64.0, 56.0/64.0, 24.0/64.0, 50.0/64.0, 18.0/64.0, 58.0/64.0, 26.0/64.0,
        12.0/64.0, 44.0/64.0,  4.0/64.0, 36.0/64.0, 14.0/64.0, 46.0/64.0,  6.0/64.0, 38.0/64.0,
        60.0/64.0, 28.0/64.0, 52.0/64.0, 20.0/64.0, 62.0/64.0, 30.0/64.0, 54.0/64.0, 22.0/64.0,
        3.0/64.0,  35.0/64.0, 11.0/64.0, 43.0/64.0,  1.0/64.0, 33.0/64.0,  9.0/64.0, 41.0/64.0,
        51.0/64.0, 19.0/64.0, 59.0/64.0, 27.0/64.0, 49.0/64.0, 17.0/64.0, 57.0/64.0, 25.0/64.0,
        15.0/64.0, 47.0/64.0,  7.0/64.0, 39.0/64.0, 13.0/64.0, 45.0/64.0,  5.0/64.0, 37.0/64.0,
        63.0/64.0, 31.0/64.0, 55.0/64.0, 23.0/64.0, 61.0/64.0, 29.0/64.0, 53.0/64.0, 21.0/64.0
    );
    let x = i32(screen_pos.x) % 8;
    let y = i32(screen_pos.y) % 8;
    let index = y * 8 + x;
    return select(0.0, 1.0, value > pattern[index]);
}


@fragment
fn fs_main(input: VertexOutput) -> @location(0) vec4<f32> {
    let light_pos = vec3<f32>(5.0, 5.0, 5.0);
    let light_color = vec3<f32>(1.0, 1.0, 1.0);
    let object_color = vec3<f32>(1.0, 1.0, 1.0);

    let ambient_strength = 0.3;
    let ambient = ambient_strength * light_color;

    let norm = normalize(input.world_normal);
    let light_dir = normalize(vec3<f32>(1.0, -1.0, 1.0));
    let diff = max(dot(norm, light_dir), 0.0);
    let diffuse = diff * light_color;

    let result = (ambient + diffuse) * object_color;

    let dithered_r = bayer_8x8_dither(result.r, input.clip_position.xy);
    let dithered_g = bayer_8x8_dither(result.g, input.clip_position.xy);
    let dithered_b = bayer_8x8_dither(result.b, input.clip_position.xy);

    return vec4<f32>(dithered_r, dithered_g, dithered_b, 1.0);
}