//! gpu_fft.rs — Dual-CSS HF Mode
//! Optional GPU-compute FFT via wgpu (vendor-neutral: Vulkan / Metal / DX12 /
//! OpenGL — no CUDA toolchain). User-selectable; the CPU rustfft path remains
//! the default and the fallback on any GPU error.
//!
//! Algorithm: an iterative Stockham radix-2 FFT in WGSL. Stockham avoids the
//! bit-reversal permutation by ping-ponging between two buffers, one radix-2
//! stage per dispatch. All supported sizes (2048/4096/8192/16384/32768) are powers of
//! two, so radix-2 covers them. A final pass converts magnitude → dBFS and
//! writes the first N/2 bins, which are copied back to the CPU.
//!
//! Design notes:
//!   - Headless: no surface, no window. We only need a Device + Queue.
//!   - One GpuFft instance is created lazily by the audio worker and reused; it
//!     rebuilds its buffers when the FFT size changes.
//!   - Every fallible step returns Result so the worker can drop to CPU.

use std::borrow::Cow;

use wgpu::util::DeviceExt;

/// Per-dispatch uniform (stage params), padded to 16 bytes for WGSL alignment.
#[repr(C)]
#[derive(Clone, Copy, bytemuck::Pod, bytemuck::Zeroable)]
struct StageParams {
    n: u32,
    stage_len: u32, // length of the sub-DFT entering this stage (1,2,4,…,n/2)
    floor: f32,     // dBFS floor (used by the magnitude pass)
    _pad: u32,
}

/// GPU-resident buffers for a given FFT size.
struct Buffers {
    n: usize,
    buf_a: wgpu::Buffer,
    buf_b: wgpu::Buffer,
    mag_buf: wgpu::Buffer,
    readback: wgpu::Buffer,
}

pub struct GpuFft {
    device: wgpu::Device,
    queue: wgpu::Queue,
    fft_pipeline: wgpu::ComputePipeline,
    mag_pipeline: wgpu::ComputePipeline,
    fft_bind_layout: wgpu::BindGroupLayout,
    mag_bind_layout: wgpu::BindGroupLayout,
    buffers: Buffers,
    adapter_name: String,
}

impl GpuFft {
    /// Initialize a headless GPU context and pipelines for FFT size `n`.
    pub fn new(n: usize) -> Result<Self, String> {
        pollster::block_on(Self::new_async(n))
    }

    pub fn adapter_name(&self) -> &str {
        &self.adapter_name
    }

    async fn new_async(n: usize) -> Result<Self, String> {
        if !n.is_power_of_two() || n < 2 {
            return Err(format!("FFT size {n} is not a power of two"));
        }
        let instance = wgpu::Instance::default();
        let adapter = instance
            .request_adapter(&wgpu::RequestAdapterOptions {
                power_preference: wgpu::PowerPreference::HighPerformance,
                compatible_surface: None,
                force_fallback_adapter: false,
            })
            .await
            .map_err(|e| format!("no GPU adapter: {e}"))?;
        let adapter_name = adapter.get_info().name;

        let (device, queue) = adapter
            .request_device(&wgpu::DeviceDescriptor {
                label: Some("dual_css_gpu_fft"),
                required_features: wgpu::Features::empty(),
                required_limits: wgpu::Limits::downlevel_defaults(),
                experimental_features: wgpu::ExperimentalFeatures::disabled(),
                memory_hints: wgpu::MemoryHints::Performance,
                trace: wgpu::Trace::Off,
            })
            .await
            .map_err(|e| format!("GPU device request failed: {e}"))?;

        // Two separate shader modules — WGSL forbids two globals sharing a
        // group/binding slot within one module, and each pass binds different
        // buffer roles, so we keep them isolated.
        let fft_module = device.create_shader_module(wgpu::ShaderModuleDescriptor {
            label: Some("fft_stage_shader"),
            source: wgpu::ShaderSource::Wgsl(Cow::Borrowed(FFT_SHADER)),
        });
        let mag_module = device.create_shader_module(wgpu::ShaderModuleDescriptor {
            label: Some("magnitude_shader"),
            source: wgpu::ShaderSource::Wgsl(Cow::Borrowed(MAG_SHADER)),
        });

        let fft_bind_layout = device.create_bind_group_layout(&wgpu::BindGroupLayoutDescriptor {
            label: Some("fft_bind_layout"),
            entries: &[
                storage_entry(0, true),
                storage_entry(1, false),
                uniform_entry(2),
            ],
        });
        let mag_bind_layout = device.create_bind_group_layout(&wgpu::BindGroupLayoutDescriptor {
            label: Some("mag_bind_layout"),
            entries: &[
                storage_entry(0, true),
                storage_entry(1, false),
                uniform_entry(2),
            ],
        });

        let fft_pipeline = device.create_compute_pipeline(&wgpu::ComputePipelineDescriptor {
            label: Some("fft_pipeline"),
            layout: Some(
                &device.create_pipeline_layout(&wgpu::PipelineLayoutDescriptor {
                    label: Some("fft_pl_layout"),
                    bind_group_layouts: &[Some(&fft_bind_layout)],
                    immediate_size: 0,
                }),
            ),
            module: &fft_module,
            entry_point: Some("fft_stage"),
            compilation_options: Default::default(),
            cache: None,
        });
        let mag_pipeline = device.create_compute_pipeline(&wgpu::ComputePipelineDescriptor {
            label: Some("mag_pipeline"),
            layout: Some(
                &device.create_pipeline_layout(&wgpu::PipelineLayoutDescriptor {
                    label: Some("mag_pl_layout"),
                    bind_group_layouts: &[Some(&mag_bind_layout)],
                    immediate_size: 0,
                }),
            ),
            module: &mag_module,
            entry_point: Some("magnitude"),
            compilation_options: Default::default(),
            cache: None,
        });

        let buffers = Buffers::new(&device, n);

        Ok(Self {
            device,
            queue,
            fft_pipeline,
            mag_pipeline,
            fft_bind_layout,
            mag_bind_layout,
            buffers,
            adapter_name,
        })
    }

    /// Compute the power spectrum (dBFS, first N/2 bins) of a real input that the
    /// caller has already windowed. Values past `n` are ignored; short inputs
    /// are zero-padded.
    pub fn power_spectrum(
        &mut self,
        samples: &[f32],
        n: usize,
        floor: f32,
    ) -> Result<Vec<f32>, String> {
        if !n.is_power_of_two() || n < 2 {
            return Err(format!("FFT size {n} is not a power of two"));
        }
        if self.buffers.n != n {
            self.buffers = Buffers::new(&self.device, n);
        }

        // Upload interleaved complex input (im = 0) into buf_a.
        let mut interleaved = vec![0.0f32; 2 * n];
        for i in 0..n {
            interleaved[2 * i] = samples.get(i).copied().unwrap_or(0.0);
        }
        self.queue
            .write_buffer(&self.buffers.buf_a, 0, bytemuck::cast_slice(&interleaved));

        let workgroups = (n as u32).div_ceil(64);
        let mut encoder = self
            .device
            .create_command_encoder(&wgpu::CommandEncoderDescriptor {
                label: Some("fft_encoder"),
            });

        // Build a bind group per Stockham stage (stage_len = 1,2,…,n/2),
        // ping-ponging src↔dst. After n_stages swaps the result lands back in
        // buf_a when n_stages is even, buf_b when odd — tracked by `src`.
        let n_stages = (n as f32).log2() as u32;
        let mut src = &self.buffers.buf_a;
        let mut dst = &self.buffers.buf_b;
        let mut keep = Vec::with_capacity(n_stages as usize);

        let mut stage_len = 1u32;
        for _ in 0..n_stages {
            let params = StageParams {
                n: n as u32,
                stage_len,
                floor,
                _pad: 0,
            };
            let params_buf = self
                .device
                .create_buffer_init(&wgpu::util::BufferInitDescriptor {
                    label: Some("fft_params"),
                    contents: bytemuck::bytes_of(&params),
                    usage: wgpu::BufferUsages::UNIFORM,
                });
            let bg = self.device.create_bind_group(&wgpu::BindGroupDescriptor {
                label: Some("fft_bg"),
                layout: &self.fft_bind_layout,
                entries: &[bind(0, src), bind(1, dst), bind(2, &params_buf)],
            });
            keep.push((bg, params_buf));
            stage_len *= 2;
            std::mem::swap(&mut src, &mut dst);
        }

        {
            let mut pass = encoder.begin_compute_pass(&wgpu::ComputePassDescriptor {
                label: Some("fft_pass"),
                timestamp_writes: None,
            });
            pass.set_pipeline(&self.fft_pipeline);
            for (bg, _) in &keep {
                pass.set_bind_group(0, bg, &[]);
                pass.dispatch_workgroups(workgroups, 1, 1);
            }
        }

        let result_buf = src; // holds the final spectrum after the swaps

        // Magnitude → dBFS over the first N/2 bins.
        let mag_params = StageParams {
            n: n as u32,
            stage_len: 0,
            floor,
            _pad: 0,
        };
        let mag_params_buf = self
            .device
            .create_buffer_init(&wgpu::util::BufferInitDescriptor {
                label: Some("mag_params"),
                contents: bytemuck::bytes_of(&mag_params),
                usage: wgpu::BufferUsages::UNIFORM,
            });
        let mag_bg = self.device.create_bind_group(&wgpu::BindGroupDescriptor {
            label: Some("mag_bg"),
            layout: &self.mag_bind_layout,
            entries: &[
                bind(0, result_buf),
                bind(1, &self.buffers.mag_buf),
                bind(2, &mag_params_buf),
            ],
        });
        {
            let mut pass = encoder.begin_compute_pass(&wgpu::ComputePassDescriptor {
                label: Some("mag_pass"),
                timestamp_writes: None,
            });
            pass.set_pipeline(&self.mag_pipeline);
            pass.set_bind_group(0, &mag_bg, &[]);
            pass.dispatch_workgroups(((n / 2) as u32).div_ceil(64), 1, 1);
        }

        let mag_bytes = (n / 2 * std::mem::size_of::<f32>()) as u64;
        encoder.copy_buffer_to_buffer(
            &self.buffers.mag_buf,
            0,
            &self.buffers.readback,
            0,
            mag_bytes,
        );
        self.queue.submit(Some(encoder.finish()));

        // Map readback and wait.
        let slice = self.buffers.readback.slice(..);
        let (tx, rx) = std::sync::mpsc::channel();
        slice.map_async(wgpu::MapMode::Read, move |res| {
            let _ = tx.send(res);
        });
        // Bounded wait: a wedged driver returns PollError::Timeout instead of
        // hanging the FFT worker forever, and the worker falls back to CPU.
        self.device
            .poll(wgpu::PollType::Wait {
                submission_index: None,
                timeout: Some(std::time::Duration::from_secs(5)),
            })
            .map_err(|e| format!("GPU poll failed: {e}"))?;
        match rx.recv() {
            Ok(Ok(())) => {}
            Ok(Err(e)) => return Err(format!("GPU buffer map failed: {e}")),
            Err(e) => return Err(format!("GPU map channel closed: {e}")),
        }

        let data = slice.get_mapped_range();
        let out: Vec<f32> = bytemuck::cast_slice(&data).to_vec();
        drop(data);
        self.buffers.readback.unmap();
        Ok(out)
    }
}

impl Buffers {
    fn new(device: &wgpu::Device, n: usize) -> Self {
        let complex_bytes = (2 * n * std::mem::size_of::<f32>()) as u64;
        let mag_bytes = (n / 2 * std::mem::size_of::<f32>()) as u64;
        let storage = wgpu::BufferUsages::STORAGE | wgpu::BufferUsages::COPY_DST;
        Self {
            n,
            buf_a: device.create_buffer(&wgpu::BufferDescriptor {
                label: Some("fft_buf_a"),
                size: complex_bytes,
                usage: storage,
                mapped_at_creation: false,
            }),
            buf_b: device.create_buffer(&wgpu::BufferDescriptor {
                label: Some("fft_buf_b"),
                size: complex_bytes,
                usage: storage,
                mapped_at_creation: false,
            }),
            mag_buf: device.create_buffer(&wgpu::BufferDescriptor {
                label: Some("fft_mag"),
                size: mag_bytes,
                usage: wgpu::BufferUsages::STORAGE | wgpu::BufferUsages::COPY_SRC,
                mapped_at_creation: false,
            }),
            readback: device.create_buffer(&wgpu::BufferDescriptor {
                label: Some("fft_readback"),
                size: mag_bytes,
                usage: wgpu::BufferUsages::MAP_READ | wgpu::BufferUsages::COPY_DST,
                mapped_at_creation: false,
            }),
        }
    }
}

// ── bind-group / layout helpers ──────────────────────────────────────────────

fn storage_entry(binding: u32, read_only: bool) -> wgpu::BindGroupLayoutEntry {
    wgpu::BindGroupLayoutEntry {
        binding,
        visibility: wgpu::ShaderStages::COMPUTE,
        ty: wgpu::BindingType::Buffer {
            ty: wgpu::BufferBindingType::Storage { read_only },
            has_dynamic_offset: false,
            min_binding_size: None,
        },
        count: None,
    }
}

fn uniform_entry(binding: u32) -> wgpu::BindGroupLayoutEntry {
    wgpu::BindGroupLayoutEntry {
        binding,
        visibility: wgpu::ShaderStages::COMPUTE,
        ty: wgpu::BindingType::Buffer {
            ty: wgpu::BufferBindingType::Uniform,
            has_dynamic_offset: false,
            min_binding_size: None,
        },
        count: None,
    }
}

fn bind(binding: u32, buffer: &wgpu::Buffer) -> wgpu::BindGroupEntry<'_> {
    wgpu::BindGroupEntry {
        binding,
        resource: buffer.as_entire_binding(),
    }
}

const FFT_SHADER: &str = r#"
struct Params {
    n: u32,
    stage_len: u32,
    floor: f32,
    _pad: u32,
};

@group(0) @binding(0) var<storage, read>       src: array<f32>;
@group(0) @binding(1) var<storage, read_write> dst: array<f32>;
@group(0) @binding(2) var<uniform>             p: Params;

const PI: f32 = 3.14159265358979323846;

// One Stockham radix-2 stage. Each invocation handles one butterfly index in
// [0, n/2). Complex values are interleaved (re,im) f32 pairs.
@compute @workgroup_size(64)
fn fft_stage(@builtin(global_invocation_id) gid: vec3<u32>) {
    let half = p.n / 2u;
    let j = gid.x;
    if (j >= half) { return; }

    let len = p.stage_len;       // sub-DFT length entering this stage
    let k = j % len;             // index within the sub-DFT
    let block = j / len;         // which pair of sub-DFTs

    let in0 = block * len + k;
    let in1 = in0 + half;

    let a_re = src[2u * in0];
    let a_im = src[2u * in0 + 1u];
    let b_re = src[2u * in1];
    let b_im = src[2u * in1 + 1u];

    // Twiddle exp(-i*pi*k/len) — radix-2 forward FFT.
    let ang = -PI * f32(k) / f32(len);
    let wr = cos(ang);
    let wi = sin(ang);
    let tr = wr * b_re - wi * b_im;
    let ti = wr * b_im + wi * b_re;

    let out0 = block * (2u * len) + k;
    let out1 = out0 + len;

    dst[2u * out0]      = a_re + tr;
    dst[2u * out0 + 1u] = a_im + ti;
    dst[2u * out1]      = a_re - tr;
    dst[2u * out1 + 1u] = a_im - ti;
}
"#;

const MAG_SHADER: &str = r#"
struct Params {
    n: u32,
    stage_len: u32,
    floor: f32,
    _pad: u32,
};

@group(0) @binding(0) var<storage, read>       cin: array<f32>;
@group(0) @binding(1) var<storage, read_write> mag: array<f32>;
@group(0) @binding(2) var<uniform>             mp: Params;

// magnitude → dBFS for the first n/2 bins: 20*log10(|X|/n + eps).
@compute @workgroup_size(64)
fn magnitude(@builtin(global_invocation_id) gid: vec3<u32>) {
    let half = mp.n / 2u;
    let i = gid.x;
    if (i >= half) { return; }
    let re = cin[2u * i];
    let im = cin[2u * i + 1u];
    let m = sqrt(re * re + im * im) / f32(mp.n) + 1e-12;
    mag[i] = 20.0 * log(m) / log(10.0);
}
"#;

#[cfg(test)]
mod tests {
    use super::*;
    use rustfft::num_complex::Complex;
    use rustfft::FftPlanner;

    /// Two tests creating wgpu instances concurrently can segfault inside the
    /// EGL/Mesa driver (observed on AMD + Mesa: SIGSEGV during parallel
    /// `cargo test`). Serialize all GPU-touching tests through this lock.
    static GPU_LOCK: std::sync::Mutex<()> = std::sync::Mutex::new(());

    /// Reference CPU power spectrum (dBFS, first n/2 bins) for parity checks.
    fn cpu_reference(windowed: &[f32], n: usize) -> Vec<f32> {
        let mut planner = FftPlanner::<f32>::new();
        let fft = planner.plan_fft_forward(n);
        let mut buf: Vec<Complex<f32>> =
            windowed.iter().map(|&re| Complex { re, im: 0.0 }).collect();
        fft.process(&mut buf);
        buf.iter()
            .take(n / 2)
            .map(|c| 20.0 * (c.norm() / n as f32 + 1e-12).log10())
            .collect()
    }

    #[test]
    fn gpu_matches_cpu_when_available() {
        let _gpu = GPU_LOCK.lock().unwrap();
        // No GPU in CI/headless? Skip rather than fail — the CPU path is the
        // guaranteed one; this only validates the WGSL FFT when hardware exists.
        let n = 2048usize;
        let mut gpu = match GpuFft::new(n) {
            Ok(g) => g,
            Err(e) => {
                eprintln!("skipping GPU parity test (no GPU): {e}");
                return;
            }
        };

        // A couple of tones plus DC so several bins are non-trivial.
        let input: Vec<f32> = (0..n)
            .map(|i| {
                let t = i as f32;
                0.6 * (2.0 * std::f32::consts::PI * 64.0 * t / n as f32).sin()
                    + 0.3 * (2.0 * std::f32::consts::PI * 200.0 * t / n as f32).sin()
            })
            .collect();

        let gpu_spec = gpu.power_spectrum(&input, n, -100.0).expect("gpu fft");
        let cpu_spec = cpu_reference(&input, n);
        assert_eq!(gpu_spec.len(), cpu_spec.len());

        // Compare the strong bins; near-zero bins have huge dB variance from
        // tiny float differences, so only check bins with real energy.
        let peak = cpu_spec.iter().cloned().fold(f32::NEG_INFINITY, f32::max);
        let mut compared = 0;
        for (g, c) in gpu_spec.iter().zip(cpu_spec.iter()) {
            if *c > peak - 40.0 {
                assert!(
                    (g - c).abs() < 1.0,
                    "GPU {g:.2} dB vs CPU {c:.2} dB diverged"
                );
                compared += 1;
            }
        }
        assert!(compared > 0, "no significant bins compared");
    }

    #[test]
    fn rejects_non_power_of_two() {
        let _gpu = GPU_LOCK.lock().unwrap();
        if let Ok(mut gpu) = GpuFft::new(2048) {
            let r = gpu.power_spectrum(&[0.0; 1000], 1000, -100.0);
            assert!(r.is_err());
        }
        // If no GPU, nothing to assert — construction itself would have failed.
    }
}
