//! GPU-accelerated tree routing.
//!
//! Flattens the tree into SoA (struct-of-arrays) layout for GPU upload,
//! then dispatches a compute shader that routes each embedding root→leaf
//! independently.
//!
//! Per-node metadata (num_bits, hp_offset, bucket_offset, bucket_count,
//! children_offset, children_count) is packed into a single buffer to stay
//! within the 8-storage-buffer-per-stage limit.

use std::sync::Arc;

use super::buffers::DeviceExt;
use super::{GpuContext, GpuError};
use crate::tree::LshTree;

/// Stride for packed node_meta buffer (6 u32s per node).
const META_STRIDE: usize = 6;

/// Flattened tree in SoA layout for GPU upload.
pub struct FlatTreeGpu {
    /// Packed per-node metadata: [num_bits, hp_offset, bucket_offset, bucket_count,
    /// children_offset, children_count] × n_nodes
    pub node_meta: Vec<u32>,
    pub hyperplanes_flat: Vec<f32>,
    pub bucket_table: Vec<u32>, // (hash, child_idx) pairs flattened
    pub centroids_flat: Vec<f32>,
    pub children_flat: Vec<u32>,
    pub root: u32,
    pub dim: u32,
    pub n_nodes: usize,
}

impl FlatTreeGpu {
    /// Flatten an LshTree into GPU-ready SoA arrays.
    pub fn from_tree(tree: &LshTree) -> Self {
        let n_nodes = tree.num_nodes();
        let dim = tree.embedding_dim();

        let mut node_meta = Vec::with_capacity(n_nodes * META_STRIDE);
        let mut hyperplanes_flat = Vec::new();
        let mut bucket_table = Vec::new();
        let mut centroids_flat = Vec::with_capacity(n_nodes * dim);
        let mut children_flat = Vec::new();

        for i in 0..n_nodes {
            let node = tree.node(i);

            let num_bits = node.num_bits as u32;
            let hp_offset = hyperplanes_flat.len() as u32;
            let bucket_offset = (bucket_table.len() / 2) as u32;
            let bucket_count = node.bucket_id_to_child.len() as u32;
            let children_offset = children_flat.len() as u32;
            let children_count = node.children.len() as u32;

            // Pack 6 metadata fields
            node_meta.push(num_bits);
            node_meta.push(hp_offset);
            node_meta.push(bucket_offset);
            node_meta.push(bucket_count);
            node_meta.push(children_offset);
            node_meta.push(children_count);

            // Hyperplanes
            if let Some(hp) = &node.hyperplanes {
                hyperplanes_flat.extend_from_slice(hp);
            }

            // Bucket table: (hash, child_idx) pairs
            for &(hash, child_idx) in &node.bucket_id_to_child {
                bucket_table.push(hash as u32);
                bucket_table.push(child_idx as u32);
            }

            // Centroids
            centroids_flat.extend_from_slice(&node.centroid);

            // Children
            for &child_idx in &node.children {
                children_flat.push(child_idx as u32);
            }
        }

        // Ensure non-empty arrays (WGSL doesn't like zero-length storage buffers)
        if hyperplanes_flat.is_empty() {
            hyperplanes_flat.push(0.0);
        }
        if bucket_table.is_empty() {
            bucket_table.push(0);
            bucket_table.push(0);
        }
        if children_flat.is_empty() {
            children_flat.push(0);
        }

        FlatTreeGpu {
            node_meta,
            hyperplanes_flat,
            bucket_table,
            centroids_flat,
            children_flat,
            root: tree.root() as u32,
            dim: dim as u32,
            n_nodes,
        }
    }
}

/// Uniform parameters passed to the shader.
#[repr(C)]
#[derive(Copy, Clone, bytemuck::Pod, bytemuck::Zeroable)]
struct Params {
    dim: u32,
    root: u32,
    n_items: u32,
    _pad: u32,
}

/// GPU tree routing pipeline.
pub struct GpuTreeRouter {
    pipeline: wgpu::ComputePipeline,
    tree_bind_group: wgpu::BindGroup,
    data_bind_group_layout: wgpu::BindGroupLayout,
    params_bind_group_layout: wgpu::BindGroupLayout,
    ctx: Arc<GpuContext>,
    flat_tree: FlatTreeGpu,
}

impl GpuTreeRouter {
    /// Create a GPU tree router, uploading the tree structure to the GPU.
    pub fn new(ctx: Arc<GpuContext>, tree: &LshTree) -> Result<Self, GpuError> {
        let flat_tree = FlatTreeGpu::from_tree(tree);

        let shader = ctx.device.create_shader_module(wgpu::ShaderModuleDescriptor {
            label: Some("tree_routing"),
            source: wgpu::ShaderSource::Wgsl(include_str!("tree_routing.wgsl").into()),
        });

        // Bind group layout for tree structure (group 0) — 5 storage buffers
        let tree_bind_group_layout =
            ctx.device
                .create_bind_group_layout(&wgpu::BindGroupLayoutDescriptor {
                    label: Some("tree_layout"),
                    entries: &(0..5)
                        .map(|i| wgpu::BindGroupLayoutEntry {
                            binding: i,
                            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,
                        })
                        .collect::<Vec<_>>(),
                });

        // Bind group layout for embeddings + results (group 1) — 2 storage buffers
        let data_bind_group_layout =
            ctx.device
                .create_bind_group_layout(&wgpu::BindGroupLayoutDescriptor {
                    label: Some("data_layout"),
                    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: false },
                                has_dynamic_offset: false,
                                min_binding_size: None,
                            },
                            count: None,
                        },
                    ],
                });

        // Bind group layout for params (group 2) — 1 uniform buffer
        let params_bind_group_layout =
            ctx.device
                .create_bind_group_layout(&wgpu::BindGroupLayoutDescriptor {
                    label: Some("params_layout"),
                    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("tree_routing_layout"),
                    bind_group_layouts: &[
                        &tree_bind_group_layout,
                        &data_bind_group_layout,
                        &params_bind_group_layout,
                    ],
                    immediate_size: 0,
                });

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

        // Upload tree buffers
        let usage = wgpu::BufferUsages::STORAGE;
        let make_buf = |label: &str, data: &[u32]| {
            ctx.device
                .create_buffer_init(&wgpu::util::BufferInitDescriptor {
                    label: Some(label),
                    contents: bytemuck::cast_slice(data),
                    usage,
                })
        };
        let make_buf_f32 = |label: &str, data: &[f32]| {
            ctx.device
                .create_buffer_init(&wgpu::util::BufferInitDescriptor {
                    label: Some(label),
                    contents: bytemuck::cast_slice(data),
                    usage,
                })
        };

        let buf_node_meta = make_buf("node_meta", &flat_tree.node_meta);
        let buf_hp_flat = make_buf_f32("hyperplanes_flat", &flat_tree.hyperplanes_flat);
        let buf_bucket_table = make_buf("bucket_table", &flat_tree.bucket_table);
        let buf_centroids = make_buf_f32("centroids_flat", &flat_tree.centroids_flat);
        let buf_children_flat = make_buf("children_flat", &flat_tree.children_flat);

        let tree_bind_group = ctx.device.create_bind_group(&wgpu::BindGroupDescriptor {
            label: Some("tree_bg"),
            layout: &tree_bind_group_layout,
            entries: &[
                wgpu::BindGroupEntry { binding: 0, resource: buf_node_meta.as_entire_binding() },
                wgpu::BindGroupEntry { binding: 1, resource: buf_hp_flat.as_entire_binding() },
                wgpu::BindGroupEntry { binding: 2, resource: buf_bucket_table.as_entire_binding() },
                wgpu::BindGroupEntry { binding: 3, resource: buf_centroids.as_entire_binding() },
                wgpu::BindGroupEntry { binding: 4, resource: buf_children_flat.as_entire_binding() },
            ],
        });

        Ok(GpuTreeRouter {
            pipeline,
            tree_bind_group,
            data_bind_group_layout,
            params_bind_group_layout,
            ctx,
            flat_tree,
        })
    }

    /// Access the flattened tree (for testing/inspection).
    pub fn flat_tree(&self) -> &FlatTreeGpu {
        &self.flat_tree
    }

    /// Route a batch of embeddings through the tree, returning leaf node indices.
    ///
    /// `embeddings`: flat f32 slice of shape (n × dim)
    /// Returns: Vec<u32> of leaf node indices, one per embedding.
    pub fn route_batch(&self, embeddings: &[f32], n: usize) -> Result<Vec<u32>, GpuError> {
        let dim = self.flat_tree.dim as usize;
        assert_eq!(embeddings.len(), n * dim, "embeddings length mismatch");

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

        let emb_buffer = device.create_buffer_init(&wgpu::util::BufferInitDescriptor {
            label: Some("embeddings"),
            contents: bytemuck::cast_slice(embeddings),
            usage: wgpu::BufferUsages::STORAGE,
        });

        let results_buffer = device.create_buffer(&wgpu::BufferDescriptor {
            label: Some("results"),
            size: (n * std::mem::size_of::<u32>()) as u64,
            usage: wgpu::BufferUsages::STORAGE | wgpu::BufferUsages::COPY_SRC,
            mapped_at_creation: false,
        });

        let staging_buffer = device.create_buffer(&wgpu::BufferDescriptor {
            label: Some("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_bind_group = device.create_bind_group(&wgpu::BindGroupDescriptor {
            label: Some("data_bg"),
            layout: &self.data_bind_group_layout,
            entries: &[
                wgpu::BindGroupEntry {
                    binding: 0,
                    resource: emb_buffer.as_entire_binding(),
                },
                wgpu::BindGroupEntry {
                    binding: 1,
                    resource: results_buffer.as_entire_binding(),
                },
            ],
        });

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

        let workgroup_size = 256;
        let num_workgroups = (n as u32 + workgroup_size - 1) / workgroup_size;

        let mut encoder = device.create_command_encoder(&wgpu::CommandEncoderDescriptor {
            label: Some("tree_routing_encoder"),
        });
        {
            let mut pass = encoder.begin_compute_pass(&wgpu::ComputePassDescriptor {
                label: Some("tree_routing_pass"),
                timestamp_writes: None,
            });
            pass.set_pipeline(&self.pipeline);
            pass.set_bind_group(0, &self.tree_bind_group, &[]);
            pass.set_bind_group(1, &data_bind_group, &[]);
            pass.set_bind_group(2, &params_bind_group, &[]);
            pass.dispatch_workgroups(num_workgroups, 1, 1);
        }

        encoder.copy_buffer_to_buffer(
            &results_buffer,
            0,
            &staging_buffer,
            0,
            (n * std::mem::size_of::<u32>()) as u64,
        );

        queue.submit(std::iter::once(encoder.finish()));

        let slice = staging_buffer.slice(..);
        let (sender, receiver) = std::sync::mpsc::channel();
        slice.map_async(wgpu::MapMode::Read, move |result| {
            sender.send(result).unwrap();
        });
        let _ = device.poll(wgpu::PollType::wait_indefinitely());
        receiver
            .recv()
            .unwrap()
            .map_err(|e| GpuError::ValidationError(format!("Buffer map failed: {e:?}")))?;

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

        Ok(results)
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::tree::LshTree;
    use rand::prelude::*;
    use rand::SeedableRng;
    use rand_chacha::ChaCha8Rng;

    /// Build a small tree and verify GPU routing matches CPU find_leaf().
    #[test]
    fn gpu_routing_matches_cpu() {
        let dim = 32;
        let n = 500;
        let seed = 42u64;

        let mut rng = ChaCha8Rng::seed_from_u64(seed);
        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;
            }
        }

        let tree = LshTree::build(&embeddings, n, dim, 3, 10, 40, seed);

        let cpu_results: Vec<usize> = (0..n)
            .map(|i| tree.find_leaf(&embeddings[i * dim..(i + 1) * dim]))
            .collect();

        let ctx = match GpuContext::new() {
            Ok(ctx) => Arc::new(ctx),
            Err(_) => {
                eprintln!("No GPU available, skipping test");
                return;
            }
        };
        let router = GpuTreeRouter::new(ctx, &tree).unwrap();
        let gpu_results = router.route_batch(&embeddings, n).unwrap();

        let mut mismatches = 0;
        for i in 0..n {
            if cpu_results[i] as u32 != gpu_results[i] {
                mismatches += 1;
            }
        }
        assert_eq!(
            mismatches, 0,
            "GPU vs CPU: {mismatches}/{n} leaf assignments differ"
        );
    }

    /// Verify FlatTreeGpu correctly flattens a small tree.
    #[test]
    fn flat_tree_roundtrip() {
        let dim = 8;
        let n = 100;
        let seed = 123u64;

        let mut rng = ChaCha8Rng::seed_from_u64(seed);
        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;
            }
        }

        let tree = LshTree::build(&embeddings, n, dim, 2, 5, 20, seed);
        let flat = FlatTreeGpu::from_tree(&tree);

        assert_eq!(flat.dim, dim as u32);
        assert_eq!(flat.node_meta.len(), tree.num_nodes() * META_STRIDE);
        assert_eq!(flat.centroids_flat.len(), tree.num_nodes() * dim);

        // Internal nodes should have num_bits > 0
        let root = tree.node(tree.root());
        if !root.is_leaf() {
            assert!(flat.node_meta[tree.root() * META_STRIDE] > 0);
        }

        // Leaves should have num_bits == 0
        for leaf_idx in tree.leaf_indices() {
            assert_eq!(flat.node_meta[leaf_idx * META_STRIDE], 0);
        }
    }
}
