[nzsl_version("1.0")]
module;

struct VertOut
{
	[location(0)] uv: vec2[f32],
    [builtin(frag_coord)] fragcoord: vec4[f32],
}

struct FragOut
{
	[location(0)] color: vec4[f32]
}

[layout(std140)]
struct Data
{
    inv_res: vec2[f32],
    underwater: i32,
}

external
{
    [set(0), binding(0)] u_texture: sampler2D[f32],
    [set(0), binding(1)] u_depth: sampler2D[f32],
    [set(0), binding(2)] u_data: uniform[Data],
}

fn Mixf32(a: f32, b: f32, t: f32) -> f32
{
    return a + (b - a) * t;
}

fn MixVec4f32(a: vec4[f32], b: vec4[f32], t: f32) -> vec4[f32]
{
    return vec4[f32](
        Mixf32(a.x, b.x, t),
        Mixf32(a.y, b.y, t),
        Mixf32(a.z, b.z, t),
        Mixf32(a.w, b.w, t)
    );
}

option EnableDepthOfField: bool = false;
fn SampleDepthOfField(uv: vec2[f32], depth: f32) -> vec4[f32]
{
    const focal_depth: f32 = 0.5;
    const focal_range: f32 = 0.15;
    const blur_amount: f32 = 0.0005;
    const samples: i32 = 16;

    let blur: f32 = clamp(abs(depth - focal_depth) / focal_range, 0.0, 1.0) * blur_amount;

    let color = vec4[f32](0.0, 0.0, 0.0, 1.0);
    for i in 0 -> samples
    {
        let angle: f32 = f32(i) * 3.14159265 * 2.0 / f32(samples);
        let offset: vec2[f32] = vec2[f32](cos(angle), sin(angle)) * blur;
        color += u_texture.Sample(uv + offset);
    }
    return color / f32(samples);
}

option EnableRadialBlur: bool = true;
fn SampleRadialBlur(uv: vec2[f32]) -> vec4[f32]
{
    const center = vec2[f32](0.5, 0.5);
    const blur_strength: f32 = 0.01;
    const samples: i32 = 16;
    const fade_start: f32 = -0.1;

    let dir: vec2[f32] = uv - center;
    let dist: f32 = length(dir);
    let dynamic_strength: f32 = blur_strength * dist;
    let step: vec2[f32] = dir * dynamic_strength / f32(samples);

    let accum_color: vec4[f32] = vec4[f32](0.0, 0.0, 0.0, 0.0);
    let sample_uv: vec2[f32] = uv;

    for i in 0 -> samples
    {
        accum_color += u_texture.Sample(sample_uv);
        sample_uv -= step;
    }
    let color = accum_color / f32(samples);
    let fade: f32 = clamp((dist - fade_start) / (1.0 - fade_start), 0.0, 1.0);
    return color * (1.0 - fade);
}

[entry(frag)]
fn main(input: VertOut) -> FragOut
{
    let depth: f32 = u_depth.Sample(input.uv).r;
    let output: FragOut;

    const if(EnableDepthOfField)
        output.color = SampleDepthOfField(input.uv, depth);
    else
        output.color = u_texture.Sample(input.uv);

    const if(EnableRadialBlur)
        output.color = SampleRadialBlur(input.uv);

    if(u_data.underwater != 0)
    {
        const fog_near: f32 = 0.9;
        const fog_far: f32 = 1.0;
        const fog_color: vec4[f32] = vec4[f32](0.0, 0.0, 0.25, 1.0);
        let fog_factor: f32 = (fog_far - depth) / (fog_far - fog_near);
        fog_factor = clamp(fog_factor, 0.0, 1.0);
        output.color = MixVec4f32(fog_color, output.color, fog_factor);
    }

    output.color.w = 1.0;
    return output;
}