// This file is part of the 64k demo project.
// It defines WGSL shader code for various effects.

#include "../demo_effects.h"

#include "gpu/effects/shader_composer.h"

void InitShaderComposer() {
  auto& sc = ShaderComposer::Get();

  sc.RegisterSnippet("common_uniforms", R"(
struct GlobalUniforms {
    view_proj: mat4x4<f32>,
    camera_pos_time: vec4<f32>,
    params: vec4<f32>,
};
struct ObjectData {
    model: mat4x4<f32>,
    inv_model: mat4x4<f32>,
    color: vec4<f32>,
    params: vec4<f32>,
};
struct ObjectsBuffer {
    objects: array<ObjectData>,
};
)");

  sc.RegisterSnippet("sdf_primitives", R"(
fn sdSphere(p: vec3<f32>, r: f32) -> f32 {
    return length(p) - r;
}
fn sdBox(p: vec3<f32>, b: vec3<f32>) -> f32 {
    let q = abs(p) - b;
    return length(max(q, vec3<f32>(0.0))) + min(max(q.x, max(q.y, q.z)), 0.0);
}
fn sdTorus(p: vec3<f32>, t: vec2<f32>) -> f32 {
    let q = vec2<f32>(length(p.xz) - t.x, p.y);
    return length(q) - t.y;
}
fn sdPlane(p: vec3<f32>, n: vec3<f32>, h: f32) -> f32 {
  return dot(p, n) + h;
}
)");

  sc.RegisterSnippet("lighting", R"(
fn get_normal_basic(p: vec3<f32>, obj_type: f32) -> vec3<f32> {
    if (obj_type == 1.0) { return normalize(p); }
    let e = vec2<f32>(0.001, 0.0);
    return normalize(vec3<f32>(
        get_dist(p + e.xyy, obj_type) - get_dist(p - e.xyy, obj_type),
        get_dist(p + e.yxy, obj_type) - get_dist(p - e.yxy, obj_type),
        get_dist(p + e.yyx, obj_type) - get_dist(p - e.yyx, obj_type)
    ));
}

fn calc_shadow(ro: vec3<f32>, rd: vec3<f32>, tmin: f32, tmax: f32, skip_idx: u32) -> f32 {
    var res = 1.0;
    var t = tmin;
    if (t < 0.05) { t = 0.05; } 
    for (var i = 0; i < 32; i = i + 1) {
        let h = map_scene(ro + rd * t, skip_idx);
        if (h < 0.001) { return 0.0; }
        res = min(res, 16.0 * h / t); 
        t = t + clamp(h, 0.02, 0.4);
        if (t > tmax) { break; }
    }
    return clamp(res, 0.0, 1.0);
}
)");

  sc.RegisterSnippet("ray_box", R"(
struct RayBounds {
    t_entry: f32,
    t_exit: f32,
    hit: bool,
};

fn ray_box_intersection(ro: vec3<f32>, rd: vec3<f32>, extent: vec3<f32>) -> RayBounds {
    let inv_rd = 1.0 / rd;
    let t0 = (-extent - ro) * inv_rd;
    let t1 = (extent - ro) * inv_rd;
    let tmin_vec = min(t0, t1);
    let tmax_vec = max(t0, t1);
    let t_entry = max(0.0, max(tmin_vec.x, max(tmin_vec.y, tmin_vec.z)));
    let t_exit = min(tmax_vec.x, min(tmax_vec.y, tmax_vec.z));
    return RayBounds(t_entry, t_exit, t_entry <= t_exit);
}
)");
}

const char* main_shader_wgsl = R"(
struct Uniforms {
    audio_peak: f32,
    aspect_ratio: f32,
    time: f32,
};

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

@vertex fn vs_main(@builtin(vertex_index) i: u32) -> @builtin(position) vec4<f32> {
    let PI = 3.14159265;
    let num_sides = 7.0;
    let scale = 0.5 + 0.3 * uniforms.audio_peak;
    let tri_idx = f32(i / 3u);
    let sub_idx = i % 3u;
    if (sub_idx == 0u) {
        return vec4<f32>(0.0, 0.0, 0.0, 1.0);
    }
    let angle = (tri_idx + f32(sub_idx - 1u)) * 2.0 * PI / num_sides + uniforms.time * 0.5;
    return vec4<f32>(scale * cos(angle) / uniforms.aspect_ratio, scale * sin(angle), 0.0, 1.0);
}

@fragment fn fs_main() -> @location(0) vec4<f32> {
    let h = uniforms.time * 2.0 + uniforms.audio_peak * 3.0;
    let r = sin(h) * 0.5 + 0.5;
    let g = sin(h + 2.0) * 0.9 + 0.3;
    let b = sin(h + 4.0) * 0.5 + 0.5;
    let boost = uniforms.audio_peak * 0.5;
    return vec4<f32>(r + boost, g + boost, b + boost, 1.0);
})";

const char* particle_compute_wgsl = R"(
struct Particle {
    pos: vec4<f32>,
    vel: vec4<f32>,
    rot: vec4<f32>,
    color: vec4<f32>,
};

struct Uniforms {
    audio_peak: f32,
    aspect_ratio: f32,
    time: f32,
    beat: f32,
};

@group(0) @binding(0) var<storage, read_write> particles: array<Particle>;
@group(0) @binding(1) var<uniform> uniforms: Uniforms;

@compute @workgroup_size(64)
fn main(@builtin(global_invocation_id) id: vec3<u32>) {
    let i = id.x;
    if (i >= arrayLength(&particles)) {
        return;
    }
    var p = particles[i];
    let new_pos = p.pos.xyz + p.vel.xyz * 0.016;
    p.pos = vec4<f32>(new_pos, p.pos.w);
    p.vel.y = p.vel.y - 0.01 * (1.0 + uniforms.audio_peak * 5.0);
    p.rot.x = p.rot.x + p.rot.y * 0.016;
    if (p.pos.y < -1.5) {
        p.pos.y = 1.5;
        p.pos.x = (f32(i % 100u) / 50.0) - 1.0 + (uniforms.audio_peak * 0.5);
        p.vel.y = 0.0;
    }
    particles[i] = p;
})";

const char* particle_render_wgsl = R"(
struct Particle {
    pos: vec4<f32>,
    vel: vec4<f32>,
    rot: vec4<f32>,
    color: vec4<f32>,
};

struct Uniforms {
    audio_peak: f32,
    aspect_ratio: f32,
    time: f32,
    beat: f32,
};

@group(0) @binding(0) var<storage, read> particles: array<Particle>;
@group(0) @binding(1) var<uniform> uniforms: Uniforms;

struct VSOut {
    @builtin(position) pos: vec4<f32>,
    @location(0) color: vec4<f32>,
};

@vertex fn vs_main(@builtin(vertex_index) vi: u32, @builtin(instance_index) ii: u32) -> VSOut {
    let p = particles[ii];
    let size = 0.02 + p.pos.z * 0.01 + uniforms.audio_peak * 0.02;
    var offsets = array<vec2<f32>, 6>(
        vec2<f32>(-1, -1),
        vec2<f32>(1, -1),
        vec2<f32>(-1, 1),
        vec2<f32>(-1, 1),
        vec2<f32>(1, -1),
        vec2<f32>(1, 1)
    );
    let offset = offsets[vi];
    let c = cos(p.rot.x);
    let s = sin(p.rot.x);
    let rotated_offset = vec2<f32>(offset.x * c - offset.y * s, offset.x * s + offset.y * c);
    let pos = vec2<f32>(p.pos.x + rotated_offset.x * size / uniforms.aspect_ratio, p.pos.y + rotated_offset.y * size);
    return VSOut(vec4<f32>(pos, 0.0, 1.0), p.color * (0.5 + 0.5 * uniforms.audio_peak));
}

@fragment fn fs_main(@location(0) color: vec4<f32>) -> @location(0) vec4<f32> {
    return color;
}
)";

const char* passthrough_shader_wgsl = R"(
@group(0) @binding(0) var smplr: sampler;
@group(0) @binding(1) var txt: texture_2d<f32>;

struct Uniforms {
    time: f32,
    beat: f32,
    intensity: f32,
    aspect_ratio: f32,
    resolution: vec2<f32>,
};
@group(0) @binding(2) var<uniform> uniforms: Uniforms;

@vertex fn vs_main(@builtin(vertex_index) i: u32) -> @builtin(position) vec4<f32> {
    var pos = array<vec2<f32>, 3>(
        vec2<f32>(-1, -1),
        vec2<f32>(3, -1),
        vec2<f32>(-1, 3)
    );
    return vec4<f32>(pos[i], 0.0, 1.0);
}

@fragment fn fs_main(@builtin(position) p: vec4<f32>) -> @location(0) vec4<f32> {
    return textureSample(txt, smplr, p.xy / uniforms.resolution);
})";

const char* ellipse_shader_wgsl = R"(
struct Uniforms {
    time: f32,
    beat: f32,
    intensity: f32,
    aspect_ratio: f32,
    resolution: vec2<f32>,
};

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

@vertex fn vs_main(@builtin(vertex_index) i: u32) -> @builtin(position) vec4<f32> {
    var pos = array<vec2<f32>, 3>(
        vec2<f32>(-1.0, -1.0),
        vec2<f32>(3.0, -1.0),
        vec2<f32>(-1.0, 3.0)
    );
    return vec4<f32>(pos[i], 0.0, 1.0);
}

fn sdEllipse(p: vec2<f32>, ab: vec2<f32>) -> f32 {
    var p_abs = abs(p);
    if (p_abs.x > p_abs.y) {
        p_abs = vec2<f32>(p_abs.y, p_abs.x);
    }
    let l = ab.y * ab.y - ab.x * ab.x;
    let m = ab.x * p_abs.x / l;
    let n = ab.y * p_abs.y / l;
    let m2 = m * m;
    let n2 = n * n;
    let c = (m2 + n2 - 1.0) / 3.0;
    let c3 = c * c * c;
    let d = c3 + m2 * n2;
    let g = m + m * n2;
    var co: f32;
    if (d < 0.0) {
        let h = acos((c3 + m2 * n2 * 2.0) / c3) / 3.0;
        let s = cos(h);
        let t = sin(h) * sqrt(3.0);
        co = (sqrt(-c * (s + t * 2.0) + m2) + sign(l) * sqrt(-c * (s - t * 2.0) + m2) + abs(g) / (sqrt(-c * (s + t * 2.0) + m2) * sqrt(-c * (s - t * 2.0) + m2)) - m) / 2.0;
    } else {
        let h = 2.0 * m * n * sqrt(d);
        let s = sign(c3 + m2 * n2 + h) * pow(abs(c3 + m2 * n2 + h), 1.0 / 3.0);
        let u = sign(c3 + m2 * n2 - h) * pow(abs(c3 + m2 * n2 - h), 1.0 / 3.0);
        let rx = -s - u + m2 * 2.0;
        let ry = (s - u) * sqrt(3.0);
        co = (ry / sqrt(sqrt(rx * rx + ry * ry) - rx) + 2.0 * g / sqrt(rx * rx + ry * ry) - m) / 2.0;
    }
    let si = sqrt(max(0.0, 1.0 - co * co));
    return length(p_abs - vec2<f32>(ab.x * co, ab.y * si)) * sign(p_abs.y * ab.x * co - p_abs.x * ab.y * si);
}

@fragment fn fs_main(@builtin(position) p: vec4<f32>) -> @location(0) vec4<f32> {
    let uv = (p.xy / uniforms.resolution - 0.5) * 2.0;
    let movement = vec2<f32>(sin(uniforms.time * 0.7), cos(uniforms.time * 0.5));
    let d = sdEllipse((uv * vec2<f32>(uniforms.aspect_ratio, 1.0)) - movement, vec2<f32>(0.5, 0.3) * (1.0 + uniforms.beat * 0.2));
    return mix(vec4<f32>(0.2, 0.8, 0.4, 1.0), vec4<f32>(0.0), smoothstep(0.0, 0.01, d));
})";

const char* particle_spray_compute_wgsl = R"(
struct Particle {
    pos: vec4<f32>,
    vel: vec4<f32>,
    rot: vec4<f32>,
    color: vec4<f32>,
};

struct Uniforms {
    intensity: f32,
    aspect_ratio: f32,
    time: f32,
    beat: f32,
};

@group(0) @binding(0) var<storage, read_write> particles: array<Particle>;
@group(0) @binding(1) var<uniform> uniforms: Uniforms;

fn hash(p: f32) -> f32 {
    return fract(sin(p) * 43758.5453);
}

@compute @workgroup_size(64)
fn main(@builtin(global_invocation_id) id: vec3<u32>) {
    let i = id.x;
    if (i >= arrayLength(&particles)) {
        return;
    }
    var p = particles[i];
    if (p.pos.w <= 0.0) {
        let r = hash(f32(i) + uniforms.time);
        let angle = r * 6.28318;
        p.pos = vec4<f32>(0.0, 0.0, 0.0, 1.0);
        p.vel = vec4<f32>(cos(angle), sin(angle), 0.0, 0.0) * (0.5 + hash(r) * 0.5) * (1.0 + uniforms.intensity * 2.0);
        p.color = vec4<f32>(hash(r + 0.1), hash(r + 0.2), 1.0, 1.0);
    }
    let new_pos = p.pos.xyz + p.vel.xyz * 0.016;
    p.pos = vec4<f32>(new_pos, p.pos.w - 0.01 * (1.0 + uniforms.beat));
    p.vel.y = p.vel.y - 0.01;
    particles[i] = p;
})";

const char* gaussian_blur_shader_wgsl = R"(
@group(0) @binding(0) var smplr: sampler;
@group(0) @binding(1) var txt: texture_2d<f32>;

struct Uniforms {
    time: f32,
    beat: f32,
    intensity: f32,
    aspect_ratio: f32,
    resolution: vec2<f32>,
};

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

@vertex fn vs_main(@builtin(vertex_index) i: u32) -> @builtin(position) vec4<f32> {
    var pos = array<vec2<f32>, 3>(
        vec2<f32>(-1, -1),
        vec2<f32>(3, -1),
        vec2<f32>(-1, 3)
    );
    return vec4<f32>(pos[i], 0.0, 1.0);
}

@fragment fn fs_main(@builtin(position) p: vec4<f32>) -> @location(0) vec4<f32> {
    let uv = p.xy / uniforms.resolution;
    var res = vec4<f32>(0.0);
    let size = 5.0 * uniforms.intensity;
    for (var x: f32 = -2.0; x <= 2.0; x += 1.0) {
        for (var y: f32 = -2.0; y <= 2.0; y += 1.0) {
            res += textureSample(txt, smplr, uv + vec2<f32>(x, y) * size / uniforms.resolution.x);
        }
    }
    return res / 25.0;
})";

const char* solarize_shader_wgsl = R"(
@group(0) @binding(0) var smplr: sampler;
@group(0) @binding(1) var txt: texture_2d<f32>;

struct Uniforms {
    time: f32,
    beat: f32,
    intensity: f32,
    aspect_ratio: f32,
    resolution: vec2<f32>,
};

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

@vertex fn vs_main(@builtin(vertex_index) i: u32) -> @builtin(position) vec4<f32> {
    var pos = array<vec2<f32>, 3>(
        vec2<f32>(-1, -1),
        vec2<f32>(3, -1),
        vec2<f32>(-1, 3)
    );
    return vec4<f32>(pos[i], 0.0, 1.0);
}

@fragment fn fs_main(@builtin(position) p: vec4<f32>) -> @location(0) vec4<f32> {
    let uv = p.xy / uniforms.resolution;
    var col = textureSample(txt, smplr, uv);
    let thr = 0.5 + 0.3 * sin(uniforms.time);
    if (col.r < thr) {
        col.r = 1.0 - col.r;
    }
    if (col.g < thr) {
        col.g = 1.0 - col.g;
    }
    if (col.b < thr) {
        col.b = 1.0 - col.b;
    }
    return col;
})";

const char* distort_shader_wgsl = R"(
@group(0) @binding(0) var smplr: sampler;
@group(0) @binding(1) var txt: texture_2d<f32>;

struct Uniforms {
    time: f32,
    beat: f32,
    intensity: f32,
    aspect_ratio: f32,
    resolution: vec2<f32>,
};

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

@vertex fn vs_main(@builtin(vertex_index) i: u32) -> @builtin(position) vec4<f32> {
    var pos = array<vec2<f32>, 3>(
        vec2<f32>(-1, -1),
        vec2<f32>(3, -1),
        vec2<f32>(-1, 3)
    );
    return vec4<f32>(pos[i], 0.0, 1.0);
}

@fragment fn fs_main(@builtin(position) p: vec4<f32>) -> @location(0) vec4<f32> {
    let uv = p.xy / uniforms.resolution;
    let dist = 0.1 * uniforms.intensity * sin(uv.y * 20.0 + uniforms.time * 5.0);
    return textureSample(txt, smplr, uv + vec2<f32>(dist, 0.0));
})";

const char* chroma_aberration_shader_wgsl = R"(
@group(0) @binding(0) var smplr: sampler;
@group(0) @binding(1) var txt: texture_2d<f32>;

struct Uniforms {
    time: f32,
    beat: f32,
    intensity: f32,
    aspect_ratio: f32,
    resolution: vec2<f32>,
};

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

@vertex fn vs_main(@builtin(vertex_index) i: u32) -> @builtin(position) vec4<f32> {
    var pos = array<vec2<f32>, 3>(
        vec2<f32>(-1, -1),
        vec2<f32>(3, -1),
        vec2<f32>(-1, 3)
    );
    return vec4<f32>(pos[i], 0.0, 1.0);
}

@fragment fn fs_main(@builtin(position) p: vec4<f32>) -> @location(0) vec4<f32> {
    let uv = p.xy / uniforms.resolution;
    let off = 0.02 * uniforms.intensity;
    let r = textureSample(txt, smplr, uv + vec2<f32>(off, 0.0)).r;
    let g = textureSample(txt, smplr, uv).g;
    let b = textureSample(txt, smplr, uv - vec2<f32>(off, 0.0)).b;
    return vec4<f32>(r, g, b, 1.0);
})";