//! GPU-accelerated LSH hash computation.
//!
//! Each thread projects one embedding onto hyperplanes, thresholds at 0,
//! and packs the sign bits into a hash value. Replaces `hash_embeddings_flat()`.

use std::sync::Arc;

use super::buffers::DeviceExt;
use super::{GpuContext, GpuError};

/// Uniform parameters for the hash shader.
#[repr(C)]
#[derive(Copy, Clone, bytemuck::Pod, bytemuck::Zeroable)]
struct HashParams {
    dim: u32,
    num_bits: u32,
    n_items: u32,
    _pad: u32,
}

/// GPU hash scatter pipeline.
pub struct GpuHashScatter {
    pipeline: wgpu::ComputePipeline,
    data_bind_group_layout: wgpu::BindGroupLayout,
    params_bind_group_layout: wgpu::BindGroupLayout,
    ctx: Arc<GpuContext>,
}

impl GpuHashScatter {
    /// Create the hash scatter pipeline.
    pub fn new(ctx: Arc<GpuContext>) -> Result<Self, GpuError> {
        let shader = ctx.device.create_shader_module(wgpu::ShaderModuleDescriptor {
            label: Some("hash_scatter"),
            source: wgpu::ShaderSource::Wgsl(include_str!("hash_scatter.wgsl").into()),
        });

        let data_bind_group_layout =
            ctx.device
                .create_bind_group_layout(&wgpu::BindGroupLayoutDescriptor {
                    label: Some("hash_data"),
                    entries: &[
                        wgpu::BindGroupLayoutEntry {
                            binding: 0,
                            visibility: wgpu::ShaderStages::COMPUTE,
                            ty: wgpu::BindingType::Buffer {
                                ty: wgpu::BufferBindingType::Storage { read_only: true },
                                has_dynamic_offset: false,
                                min_binding_size: None,
                            },
                            count: None,
                        },
                        wgpu::BindGroupLayoutEntry {
                            binding: 1,
                            visibility: wgpu::ShaderStages::COMPUTE,
                            ty: wgpu::BindingType::Buffer {
                                ty: wgpu::BufferBindingType::Storage { read_only: true },
                                has_dynamic_offset: false,
                                min_binding_size: None,
                            },
                            count: None,
                        },
                        wgpu::BindGroupLayoutEntry {
                            binding: 2,
                            visibility: wgpu::ShaderStages::COMPUTE,
                            ty: wgpu::BindingType::Buffer {
                                ty: wgpu::BufferBindingType::Storage { read_only: false },
                                has_dynamic_offset: false,
                                min_binding_size: None,
                            },
                            count: None,
                        },
                    ],
                });

        let params_bind_group_layout =
            ctx.device
                .create_bind_group_layout(&wgpu::BindGroupLayoutDescriptor {
                    label: Some("hash_params"),
                    entries: &[wgpu::BindGroupLayoutEntry {
                        binding: 0,
                        visibility: wgpu::ShaderStages::COMPUTE,
                        ty: wgpu::BindingType::Buffer {
                            ty: wgpu::BufferBindingType::Uniform,
                            has_dynamic_offset: false,
                            min_binding_size: None,
                        },
                        count: None,
                    }],
                });

        let pipeline_layout =
            ctx.device
                .create_pipeline_layout(&wgpu::PipelineLayoutDescriptor {
                    label: Some("hash_scatter_layout"),
                    bind_group_layouts: &[&data_bind_group_layout, &params_bind_group_layout],
                    immediate_size: 0,
                });

        let pipeline = ctx
            .device
            .create_compute_pipeline(&wgpu::ComputePipelineDescriptor {
                label: Some("hash_scatter"),
                layout: Some(&pipeline_layout),
                module: &shader,
                entry_point: Some("main"),
                compilation_options: Default::default(),
                cache: None,
            });

        Ok(GpuHashScatter {
            pipeline,
            data_bind_group_layout,
            params_bind_group_layout,
            ctx,
        })
    }

    /// Compute LSH hashes for a batch of embeddings.
    ///
    /// `embeddings`: flat f32, shape (n × dim)
    /// `hyperplanes`: flat f32, shape (num_bits × dim), row-major
    /// `n`: number of embeddings
    /// `dim`: embedding dimension
    /// `num_bits`: number of hyperplanes / hash bits
    ///
    /// Returns Vec<u32> of hash values (one per embedding).
    /// Note: CPU uses u64 hashes but num_bits <= 32, so u32 suffices on GPU.
    pub fn hash_batch(
        &self,
        embeddings: &[f32],
        hyperplanes: &[f32],
        n: usize,
        dim: usize,
        num_bits: usize,
    ) -> Result<Vec<u32>, GpuError> {
        assert_eq!(embeddings.len(), n * dim);
        assert_eq!(hyperplanes.len(), num_bits * dim);
        assert!(num_bits <= 32, "GPU hash supports up to 32 bits");

        let device = &self.ctx.device;
        let queue = &self.ctx.queue;

        let emb_buf = device.create_buffer_init(&wgpu::util::BufferInitDescriptor {
            label: Some("hash_embeddings"),
            contents: bytemuck::cast_slice(embeddings),
            usage: wgpu::BufferUsages::STORAGE,
        });
        let hp_buf = device.create_buffer_init(&wgpu::util::BufferInitDescriptor {
            label: Some("hash_hyperplanes"),
            contents: bytemuck::cast_slice(hyperplanes),
            usage: wgpu::BufferUsages::STORAGE,
        });
        let hash_buf = device.create_buffer(&wgpu::BufferDescriptor {
            label: Some("hashes"),
            size: (n * std::mem::size_of::<u32>()) as u64,
            usage: wgpu::BufferUsages::STORAGE | wgpu::BufferUsages::COPY_SRC,
            mapped_at_creation: false,
        });
        let staging = device.create_buffer(&wgpu::BufferDescriptor {
            label: Some("hash_staging"),
            size: (n * std::mem::size_of::<u32>()) as u64,
            usage: wgpu::BufferUsages::MAP_READ | wgpu::BufferUsages::COPY_DST,
            mapped_at_creation: false,
        });

        let data_bg = device.create_bind_group(&wgpu::BindGroupDescriptor {
            label: Some("hash_data_bg"),
            layout: &self.data_bind_group_layout,
            entries: &[
                wgpu::BindGroupEntry { binding: 0, resource: emb_buf.as_entire_binding() },
                wgpu::BindGroupEntry { binding: 1, resource: hp_buf.as_entire_binding() },
                wgpu::BindGroupEntry { binding: 2, resource: hash_buf.as_entire_binding() },
            ],
        });

        let params = HashParams {
            dim: dim as u32,
            num_bits: num_bits as u32,
            n_items: n as u32,
            _pad: 0,
        };
        let params_buf = device.create_buffer_init(&wgpu::util::BufferInitDescriptor {
            label: Some("hash_params"),
            contents: bytemuck::bytes_of(&params),
            usage: wgpu::BufferUsages::UNIFORM,
        });
        let params_bg = device.create_bind_group(&wgpu::BindGroupDescriptor {
            label: Some("hash_params_bg"),
            layout: &self.params_bind_group_layout,
            entries: &[wgpu::BindGroupEntry {
                binding: 0,
                resource: params_buf.as_entire_binding(),
            }],
        });

        let wg_size = 256u32;
        let num_wg = (n as u32 + wg_size - 1) / wg_size;

        let mut encoder = device.create_command_encoder(&Default::default());
        {
            let mut pass = encoder.begin_compute_pass(&Default::default());
            pass.set_pipeline(&self.pipeline);
            pass.set_bind_group(0, &data_bg, &[]);
            pass.set_bind_group(1, &params_bg, &[]);
            pass.dispatch_workgroups(num_wg, 1, 1);
        }
        encoder.copy_buffer_to_buffer(&hash_buf, 0, &staging, 0, (n * 4) as u64);
        queue.submit(std::iter::once(encoder.finish()));

        let slice = staging.slice(..);
        let (tx, rx) = std::sync::mpsc::channel();
        slice.map_async(wgpu::MapMode::Read, move |r| { tx.send(r).unwrap(); });
        let _ = device.poll(wgpu::PollType::wait_indefinitely());
        rx.recv()
            .unwrap()
            .map_err(|e| GpuError::ValidationError(format!("Map failed: {e:?}")))?;

        let data = slice.get_mapped_range();
        let result: Vec<u32> = bytemuck::cast_slice(&data).to_vec();
        drop(data);
        staging.unmap();

        Ok(result)
    }
}

#[cfg(test)]
mod tests {
    use super::*;

    /// CPU reference: hash_single from tree.rs (duplicated here for test independence).
    fn hash_single_cpu(embedding: &[f32], hyperplanes: &[f32], num_bits: usize, dim: usize) -> u32 {
        let mut hash = 0u32;
        for bit in 0..num_bits {
            let hp_start = bit * dim;
            let dot: f32 = embedding
                .iter()
                .zip(&hyperplanes[hp_start..hp_start + dim])
                .map(|(a, b)| a * b)
                .sum();
            if dot > 0.0 {
                hash |= 1 << bit;
            }
        }
        hash
    }

    #[test]
    fn gpu_hash_matches_cpu() {
        use rand::prelude::*;
        use rand_chacha::ChaCha8Rng;

        let dim = 64;
        let n = 1000;
        let num_bits = 8;
        let seed = 77u64;

        let mut rng = ChaCha8Rng::seed_from_u64(seed);

        // Random embeddings (L2-normalized)
        let mut embeddings = vec![0.0f32; n * dim];
        for i in 0..n {
            let start = i * dim;
            let mut norm = 0.0f32;
            for j in 0..dim {
                let v: f32 = rng.gen_range(-1.0..1.0);
                embeddings[start + j] = v;
                norm += v * v;
            }
            let norm = norm.sqrt();
            for j in 0..dim {
                embeddings[start + j] /= norm;
            }
        }

        // Random hyperplanes (L2-normalized)
        let mut hyperplanes = vec![0.0f32; num_bits * dim];
        for bit in 0..num_bits {
            let start = bit * dim;
            let mut norm = 0.0f32;
            for j in 0..dim {
                let v: f32 = rng.gen_range(-1.0..1.0);
                hyperplanes[start + j] = v;
                norm += v * v;
            }
            let norm = norm.sqrt();
            for j in 0..dim {
                hyperplanes[start + j] /= norm;
            }
        }

        // CPU hashes
        let cpu_hashes: Vec<u32> = (0..n)
            .map(|i| {
                hash_single_cpu(
                    &embeddings[i * dim..(i + 1) * dim],
                    &hyperplanes,
                    num_bits,
                    dim,
                )
            })
            .collect();

        // GPU hashes
        let ctx = match crate::gpu::GpuContext::new() {
            Ok(ctx) => Arc::new(ctx),
            Err(_) => {
                eprintln!("No GPU available, skipping test");
                return;
            }
        };
        let hasher = GpuHashScatter::new(ctx).unwrap();
        let gpu_hashes = hasher
            .hash_batch(&embeddings, &hyperplanes, n, dim, num_bits)
            .unwrap();

        // Compare bit-for-bit
        let mut mismatches = 0;
        for i in 0..n {
            if cpu_hashes[i] != gpu_hashes[i] {
                mismatches += 1;
            }
        }
        assert_eq!(
            mismatches, 0,
            "GPU vs CPU: {mismatches}/{n} hashes differ"
        );
    }
}
