#import bevy_sprite::mesh2d_vertex_output::VertexOutput

@group(2) @binding(0) var<uniform> u_time: f32;
@group(2) @binding(1) var<uniform> u_resolution: vec2f;
@group(2) @binding(2) var<uniform> u_mouse_position: vec2f;
@group(2) @binding(3) var fft_texture: texture_2d<f32>;
@group(2) @binding(4) var fft_texture_sampler: sampler;


fn lerp(v0: f32, v1: f32, t: f32) -> f32 {
  return v0 + t * (v1 - v0);
}

// Utility: clamp to 0..1 range
fn saturate(x: f32) -> f32 { return clamp(x, 0.0, 1.0); }

// Cosine palette (like IQ's)
fn palette(t: f32, a: vec3f, b: vec3f, c: vec3f, d: vec3f) -> vec3f {
    return a + b * cos(6.2831853 * (c * t + d));
}

// Sample the FFT texture at normalized x in [0,1]
fn sample_fft(x: f32) -> f32 {
    // 5-tap Gaussian blur in texture space using actual texel width
    let size = textureDimensions(fft_texture);
    let tex_w = max(f32(size.x), 1.0);
    let one_texel = 1.0 / tex_w;
    let u = saturate(x);
    let v = 0.5;
    let w0 = 0.0625; // gaussian-ish weights: 1,4,6,4,1 normalized
    let w1 = 0.25;
    let w2 = 0.375;
    let s0 = textureSample(fft_texture, fft_texture_sampler, vec2f(u - 2.0 * one_texel, v)).r;
    let s1 = textureSample(fft_texture, fft_texture_sampler, vec2f(u - 1.0 * one_texel, v)).r;
    let s2 = textureSample(fft_texture, fft_texture_sampler, vec2f(u, v)).r;
    let s3 = textureSample(fft_texture, fft_texture_sampler, vec2f(u + 1.0 * one_texel, v)).r;
    let s4 = textureSample(fft_texture, fft_texture_sampler, vec2f(u + 2.0 * one_texel, v)).r;
    return w0 * s0 + w1 * s1 + w2 * s2 + w1 * s3 + w0 * s4;
}

@fragment
fn fragment(mesh: VertexOutput) -> @location(0) vec4f {
    // Normalize UV to -1..1 with aspect correction
    let aspect = u_resolution.x / max(u_resolution.y, 1.0);
    let p = (mesh.uv * 2.0 - vec2f(1.0, 1.0)) * vec2f(aspect, 1.0);

    // Polar coords
    let r = length(p);
    let ang = atan2(p.y, p.x); // [-pi, pi]
    var a01 = ang / 6.2831853 + 0.5; // [0,1]

    // Scroll the angle slowly for motion
    a01 = fract(a01 + 0.05 * u_time);

    // Log-like frequency remap for better distribution of low freqs
    let freq = pow(a01, 0.8);

    // Smooth FFT sample with a few taps
    let amp0 = sample_fft(freq);
    let amp1 = sample_fft(saturate(freq + 1.0 / 512.0));
    let amp2 = sample_fft(saturate(freq - 1.0 / 512.0));
    var amp = (amp0 + amp1 + amp2) / 3.0;
    // Optional shaping
    amp = pow(amp, 0.6);

    // Target ring radius driven by amplitude
    let base_r = 0.35;
    let ring_r = base_r + amp * 0.45;

    // Ring intensity with derivative-based anti-aliasing
    let band_w = mix(0.02, 0.06, 1.0 - amp);
    let aa = fwidth(r);
    let dist = abs(r - ring_r);
    let ring = 1.0 - smoothstep(band_w - aa, band_w + aa, dist);

    // Radial spokes modulated by FFT at harmonics
    let spokes = 0.5 + 0.5 * cos(ang * (8.0 + 24.0 * amp) + u_time * (1.0 + 2.0 * amp));
    let spokes_mask = pow(spokes, 3.0);

    // Glow falloff
    let glow = 0.015 / (r * r + 0.03);

    // Background subtle gradient with slow drift
    let bg = 0.08 + 0.06 * cos(6.28318 * (p.x * 0.15 + p.y * 0.12 + 0.03 * u_time));

    // Color palette over angle
    let col = palette(
        a01,
        vec3f(0.35, 0.3, 0.35), // a
        vec3f(0.35, 0.35, 0.35), // b
        vec3f(1.0, 1.0, 1.0),    // c
        vec3f(0.0, 0.33, 0.67)   // d
    );

    // Combine layers
    var color = vec3f(bg) * 0.6;
    color += col * ring * (0.6 + 0.6 * spokes_mask);
    color += vec3f(0.9, 0.95, 1.0) * glow * (0.3 + 1.5 * amp);

    // Subtle vignette
    let vig = smoothstep(1.2, 0.3, r);
    color *= vig;

    // Final boost
    color = clamp(color, vec3f(0.0), vec3f(1.0));
    return vec4f(color, 1.0);
}