use anyhow::{Result, anyhow};
use wgpu::util::DeviceExt;

pub struct Gpu {
    device: wgpu::Device,
    queue: wgpu::Queue,
}

impl Default for Gpu {
    fn default() -> Self {
        Self::new()
    }
}

impl Gpu {
    pub fn try_new() -> Result<Self> {
        pollster::block_on(Self::init())
    }

    pub fn new() -> Self {
        Self::try_new().expect("failed to initialize GPU")
    }

    async fn init() -> Result<Self> {
        let instance = wgpu::Instance::new(&wgpu::InstanceDescriptor::default());

        let adapter = instance
            .request_adapter(&wgpu::RequestAdapterOptions {
                power_preference: wgpu::PowerPreference::HighPerformance,
                ..Default::default()
            })
            .await
            .ok_or_else(|| anyhow!("no suitable GPU adapter found"))?;

        let info = adapter.get_info();
        eprintln!("gpu: {} ({:?})", info.name, info.backend);

        let (device, queue) = adapter
            .request_device(&wgpu::DeviceDescriptor::default(), None)
            .await?;

        Ok(Self { device, queue })
    }

    pub fn storage_buffer(&self, label: &str, data: &[u8]) -> wgpu::Buffer {
        self.device
            .create_buffer_init(&wgpu::util::BufferInitDescriptor {
                label: Some(label),
                contents: data,
                usage: wgpu::BufferUsages::STORAGE
                    | wgpu::BufferUsages::COPY_SRC
                    | wgpu::BufferUsages::COPY_DST,
            })
    }

    pub fn storage_buffer_empty(&self, label: &str, size: u64) -> wgpu::Buffer {
        self.device.create_buffer(&wgpu::BufferDescriptor {
            label: Some(label),
            size,
            usage: wgpu::BufferUsages::STORAGE
                | wgpu::BufferUsages::COPY_SRC
                | wgpu::BufferUsages::COPY_DST,
            mapped_at_creation: false,
        })
    }

    fn staging_buffer(&self, size: u64) -> wgpu::Buffer {
        self.device.create_buffer(&wgpu::BufferDescriptor {
            label: Some("staging"),
            size,
            usage: wgpu::BufferUsages::MAP_READ | wgpu::BufferUsages::COPY_DST,
            mapped_at_creation: false,
        })
    }

    pub fn read_buffer(&self, buffer: &wgpu::Buffer) -> Vec<u8> {
        let staging = self.staging_buffer(buffer.size());

        let mut encoder = self
            .device
            .create_command_encoder(&wgpu::CommandEncoderDescriptor { label: None });

        encoder.copy_buffer_to_buffer(buffer, 0, &staging, 0, buffer.size());
        self.queue.submit(Some(encoder.finish()));

        let slice = staging.slice(..);
        let (tx, rx) = std::sync::mpsc::channel();
        slice.map_async(wgpu::MapMode::Read, move |result| {
            tx.send(result).unwrap();
        });

        self.device.poll(wgpu::Maintain::Wait);
        rx.recv().unwrap().unwrap();

        let data = slice.get_mapped_range();
        data.to_vec()
    }

    pub fn read_buffer_as<T: bytemuck::Pod>(&self, buffer: &wgpu::Buffer) -> Vec<T> {
        let bytes = self.read_buffer(buffer);
        bytemuck::cast_slice(&bytes).to_vec()
    }

    fn build_pipeline(
        &self,
        shader_src: &str,
        entry_point: &str,
        buffers: &[(&wgpu::Buffer, bool)],
    ) -> (wgpu::ComputePipeline, wgpu::BindGroup) {
        let module = self
            .device
            .create_shader_module(wgpu::ShaderModuleDescriptor {
                label: None,
                source: wgpu::ShaderSource::Wgsl(shader_src.into()),
            });

        let bind_group_layout_entries = buffers
            .iter()
            .enumerate()
            .map(|(i, &(_, read_only))| wgpu::BindGroupLayoutEntry {
                binding: i.try_into().unwrap(),
                visibility: wgpu::ShaderStages::COMPUTE,
                ty: wgpu::BindingType::Buffer {
                    ty: wgpu::BufferBindingType::Storage { read_only },
                    has_dynamic_offset: false,
                    min_binding_size: None,
                },
                count: None,
            })
            .collect::<Vec<_>>();

        let bind_group_layout =
            self.device
                .create_bind_group_layout(&wgpu::BindGroupLayoutDescriptor {
                    label: None,
                    entries: &bind_group_layout_entries,
                });

        let pipeline_layout = self
            .device
            .create_pipeline_layout(&wgpu::PipelineLayoutDescriptor {
                label: None,
                bind_group_layouts: &[&bind_group_layout],
                push_constant_ranges: &[],
            });

        let pipeline = self
            .device
            .create_compute_pipeline(&wgpu::ComputePipelineDescriptor {
                label: None,
                layout: Some(&pipeline_layout),
                module: &module,
                entry_point: Some(entry_point),
                cache: None,
                compilation_options: wgpu::PipelineCompilationOptions::default(),
            });

        let bind_group_entries = buffers
            .iter()
            .enumerate()
            .map(|(i, (buf, _))| wgpu::BindGroupEntry {
                binding: i.try_into().unwrap(),
                resource: buf.as_entire_binding(),
            })
            .collect::<Vec<_>>();

        let bind_group = self.device.create_bind_group(&wgpu::BindGroupDescriptor {
            label: None,
            layout: &bind_group_layout,
            entries: &bind_group_entries,
        });

        (pipeline, bind_group)
    }

    pub fn dispatch(
        &self,
        shader_src: &str,
        entry_point: &str,
        buffers: &[(&wgpu::Buffer, bool)],
        workgroups: u32,
    ) {
        let (pipeline, bind_group) = self.build_pipeline(shader_src, entry_point, buffers);

        let mut encoder = self
            .device
            .create_command_encoder(&wgpu::CommandEncoderDescriptor { label: None });

        {
            let mut cpass = encoder.begin_compute_pass(&Default::default());
            cpass.set_pipeline(&pipeline);
            cpass.set_bind_group(0, &bind_group, &[]);
            cpass.dispatch_workgroups(workgroups, 1, 1);
        }

        self.queue.submit(Some(encoder.finish()));
    }
}

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

    #[test]
    fn test_scan_u32() {
        let gpu = match Gpu::try_new() {
            Ok(g) => g,
            Err(_) => return,
        };
        let input = vec![1u32; 256];
        let buf = gpu.storage_buffer("scan", bytemuck::cast_slice(&input));

        gpu.dispatch(
            include_str!("shaders/scan_u32.wgsl"),
            "main",
            &[(&buf, false)],
            1,
        );

        let out = gpu.read_buffer_as::<u32>(&buf);
        assert_eq!(out, (1..=256).collect::<Vec<_>>());
    }

    #[test]
    fn test_fsm_string_detect() {
        let gpu = match Gpu::try_new() {
            Ok(g) => g,
            Err(_) => return,
        };

        let input_str = r#"hello "world" end"#;

        let mut input = input_str.bytes().collect::<Vec<u8>>();
        input.resize(256, 0);

        let input_buf = gpu.storage_buffer("fsm_in", &input);
        let output_buf = gpu.storage_buffer_empty(
            "fsm_out",
            // vec3 is padded to 16 bytes in storage
            (256 * std::mem::size_of::<[u32; 4]>()) as u64,
        );

        gpu.dispatch(
            include_str!("shaders/scan_fsm.wgsl"),
            "main",
            &[(&input_buf, true), (&output_buf, false)],
            1,
        );

        let out = gpu.read_buffer_as::<[u32; 4]>(&output_buf);

        let states = out[..input_str.len()]
            .iter()
            .map(|f| f[0])
            .collect::<Vec<_>>();

        //                  h  e  l  l  o     "  w  o  r  l  d  "     e  n  d
        assert_eq!(states, [0, 0, 0, 0, 0, 0, 1, 1, 1, 1, 1, 1, 0, 0, 0, 0, 0]);
    }

    #[test]
    fn test_fsm_escape() {
        let gpu = match Gpu::try_new() {
            Ok(g) => g,
            Err(_) => return,
        };

        // the \" inside the string should NOT close it
        let input_str = r#"hello "wo\"rld" end"#;

        let mut input = input_str.bytes().collect::<Vec<u8>>();
        input.resize(256, 0);

        let input_buf = gpu.storage_buffer("fsm_in", &input);
        let output_buf =
            gpu.storage_buffer_empty("fsm_out", (256 * std::mem::size_of::<[u32; 4]>()) as u64);

        gpu.dispatch(
            include_str!("shaders/scan_fsm.wgsl"),
            "main",
            &[(&input_buf, true), (&output_buf, false)],
            1,
        );

        let out = gpu.read_buffer_as::<[u32; 4]>(&output_buf);

        let states = out[..input_str.len()]
            .iter()
            .map(|f| f[0])
            .collect::<Vec<_>>();

        assert_eq!(
            states,
            // h   e  l  l  o     "  w  o  \  "  r  l  d  "     e  n  d
            [0u32, 0, 0, 0, 0, 0, 1, 1, 1, 2, 1, 1, 1, 1, 0, 0, 0, 0, 0]
        );
    }

    #[test]
    fn test_depth() {
        let gpu = match Gpu::try_new() {
            Ok(g) => g,
            Err(_) => return,
        };

        // {"a":{"b":[1,2]}}
        let input_str = r#"{"a":{"b":[1,2]}}"#;

        let mut input = input_str.bytes().collect::<Vec<u8>>();
        input.resize(256, 0);

        let input_buf = gpu.storage_buffer("bytes", &input);
        let fsm_buf =
            gpu.storage_buffer_empty("fsm", (256 * std::mem::size_of::<[u32; 4]>()) as u64);

        gpu.dispatch(
            include_str!("shaders/scan_fsm.wgsl"),
            "main",
            &[(&input_buf, true), (&fsm_buf, false)],
            1,
        );

        let compact_buf =
            gpu.storage_buffer_empty("compact", (256 * std::mem::size_of::<u32>()) as u64);
        let count_buf = gpu.storage_buffer_empty("count", 4);

        gpu.dispatch(
            include_str!("shaders/scan_compact.wgsl"),
            "main",
            &[
                (&input_buf, true),
                (&fsm_buf, true),
                (&compact_buf, false),
                (&count_buf, false),
            ],
            1,
        );

        let depth_buf =
            gpu.storage_buffer_empty("depth", (256 * std::mem::size_of::<i32>()) as u64);

        gpu.dispatch(
            include_str!("shaders/scan_depth.wgsl"),
            "main",
            &[
                (&input_buf, true),
                (&compact_buf, true),
                (&depth_buf, false),
            ],
            1,
        );

        let depths = gpu.read_buffer_as::<i32>(&depth_buf);
        let count =
            usize::try_from(gpu.read_buffer_as::<u32>(&count_buf)[0]).expect("count fits usize");

        assert_eq!(&depths[..count], &[1, 1, 1, 2, 2, 2, 3, 3, 3, 3, 2, 1, 0]);
    }

    #[test]
    fn test_sort() {
        let gpu = match Gpu::try_new() {
            Ok(g) => g,
            Err(_) => return,
        };

        // compacted structural chars from {"a":{"b":[1,2]}}
        // positions: {  :  {  :  [  ,  ]  }  }
        //            0  4  5  9  10 12 14 15 17
        // depths:    1  1  2  2  3  3  2  1  0
        let mut positions: Vec<u32> = vec![0, 4, 5, 9, 10, 12, 14, 15, 17];
        let mut depths: Vec<i32> = vec![1, 1, 2, 2, 3, 3, 2, 1, 0];

        positions.resize(256, 0);
        depths.resize(256, 0);

        let depth_buf = gpu.storage_buffer("depths", bytemuck::cast_slice(&depths));
        let pos_buf = gpu.storage_buffer("positions", bytemuck::cast_slice(&positions));

        gpu.dispatch(
            include_str!("shaders/sort.wgsl"),
            "main",
            &[(&depth_buf, false), (&pos_buf, false)],
            1,
        );

        let sorted_depths = gpu.read_buffer_as::<i32>(&depth_buf);
        let sorted_positions = gpu.read_buffer_as::<u32>(&pos_buf);

        let mut pairs = depths[..256]
            .iter()
            .zip(positions[..256].iter())
            .map(|(&d, &p)| (d, p))
            .collect::<Vec<_>>();

        pairs.sort_by(|a, b| a.0.cmp(&b.0).then(a.1.cmp(&b.1)));

        let expected_depths = pairs.iter().map(|p| p.0).collect::<Vec<_>>();
        let expected_positions = pairs.iter().map(|p| p.1).collect::<Vec<_>>();

        assert_eq!(sorted_depths, expected_depths);
        assert_eq!(sorted_positions, expected_positions);
    }

    #[test]
    fn test_parent_link() {
        let gpu = match Gpu::try_new() {
            Ok(g) => g,
            Err(_) => return,
        };

        let input_str = r#"{"a":{"b":[1,2]}}"#;

        let mut input = input_str.bytes().collect::<Vec<u8>>();
        input.resize(256, 0);

        let input_buf = gpu.storage_buffer("bytes", &input);
        let fsm_buf =
            gpu.storage_buffer_empty("fsm", (256 * std::mem::size_of::<[u32; 4]>()) as u64);

        gpu.dispatch(
            include_str!("shaders/scan_fsm.wgsl"),
            "main",
            &[(&input_buf, true), (&fsm_buf, false)],
            1,
        );

        let compact_buf =
            gpu.storage_buffer_empty("compact", (256 * std::mem::size_of::<u32>()) as u64);
        let count_buf = gpu.storage_buffer_empty("count", 4);

        gpu.dispatch(
            include_str!("shaders/scan_compact.wgsl"),
            "main",
            &[
                (&input_buf, true),
                (&fsm_buf, true),
                (&compact_buf, false),
                (&count_buf, false),
            ],
            1,
        );

        let depth_buf =
            gpu.storage_buffer_empty("depth", (256 * std::mem::size_of::<i32>()) as u64);

        gpu.dispatch(
            include_str!("shaders/scan_depth.wgsl"),
            "main",
            &[
                (&input_buf, true),
                (&compact_buf, true),
                (&depth_buf, false),
            ],
            1,
        );

        let parent_buf =
            gpu.storage_buffer_empty("parents", (256 * std::mem::size_of::<i32>()) as u64);

        gpu.dispatch(
            include_str!("shaders/parent_link.wgsl"),
            "main",
            &[
                (&input_buf, true),
                (&compact_buf, true),
                (&depth_buf, true),
                (&parent_buf, false),
            ],
            1,
        );

        let parents = gpu.read_buffer_as::<i32>(&parent_buf);
        let depths = gpu.read_buffer_as::<i32>(&depth_buf);
        let positions = gpu.read_buffer_as::<u32>(&compact_buf);
        let count =
            usize::try_from(gpu.read_buffer_as::<u32>(&count_buf)[0]).expect("count fits usize");

        assert_eq!(parents[0], -1);
        assert_eq!(
            &positions[..count],
            &[0, 1, 4, 5, 6, 9, 10, 11, 12, 13, 14, 15, 16]
        );
        assert_eq!(&depths[..count], &[1, 1, 1, 2, 2, 2, 3, 3, 3, 3, 2, 1, 0]);
        assert_eq!(&parents[..count], &[-1, 0, 0, 0, 3, 3, 3, 6, 6, 6, 6, 3, 0]);
    }
}
