use super::{CommonBuffers, Params, RoundParams};
use bytemuck::{Pod, Zeroable};
use std::mem::size_of;
use wgpu::{
    include_wgsl, BindGroup, BindGroupDescriptor, BindGroupEntry, BindGroupLayoutDescriptor,
    BindGroupLayoutEntry, BindingType, Buffer, BufferDescriptor, BufferUsages, CommandEncoder,
    ComputePipeline, ComputePipelineDescriptor, Device, PipelineLayoutDescriptor, Queue,
    ShaderStages, TextureSampleType, TextureViewDimension,
};

pub struct MotionEncoder {
    pipeline: ComputePipeline,
    bind_groups: [BindGroup; 2],

    uniform_buffer: Buffer,
    uniform: EncoderUniform,
}

#[repr(C)]
#[derive(Debug, Clone, Copy, Pod, Zeroable, Default)]
pub struct EncoderUniform {
    block_size: [i32; 2],
    output_stride: u32,
    search_radius: i32,
    skip_threshold: f32,
    _pad: u32,
}

impl MotionEncoder {
    pub fn create(device: &Device, params: &Params, bufs: &CommonBuffers) -> Self {
        let uniform_buffer = device.create_buffer(&BufferDescriptor {
            label: Some("encoder uniforms"),
            size: size_of::<EncoderUniform>() as u64,
            usage: BufferUsages::UNIFORM | BufferUsages::COPY_DST,
            mapped_at_creation: false,
        });
        let uniform = EncoderUniform {
            block_size: [params.block_width as i32, params.block_height as i32],
            output_stride: (params.width / params.block_width) as u32,
            search_radius: 16,
            skip_threshold: 0.1,
            ..Default::default()
        };

        let bind_group_layout = device.create_bind_group_layout(&BindGroupLayoutDescriptor {
            label: None,
            entries: &[
                BindGroupLayoutEntry {
                    binding: 0,
                    count: None,
                    ty: BindingType::Buffer {
                        ty: wgpu::BufferBindingType::Uniform,
                        has_dynamic_offset: false,
                        min_binding_size: None,
                    },
                    visibility: ShaderStages::COMPUTE,
                },
                BindGroupLayoutEntry {
                    binding: 1,
                    count: None,
                    ty: BindingType::Buffer {
                        has_dynamic_offset: false,
                        min_binding_size: None,
                        ty: wgpu::BufferBindingType::Storage { read_only: false },
                    },
                    visibility: ShaderStages::COMPUTE,
                },
                BindGroupLayoutEntry {
                    binding: 2,
                    count: None,
                    ty: BindingType::Texture {
                        sample_type: TextureSampleType::Float { filterable: false },
                        view_dimension: TextureViewDimension::D2,
                        multisampled: false,
                    },
                    visibility: ShaderStages::COMPUTE,
                },
                BindGroupLayoutEntry {
                    binding: 3,
                    count: None,
                    ty: BindingType::Texture {
                        sample_type: TextureSampleType::Float { filterable: false },
                        view_dimension: TextureViewDimension::D2,
                        multisampled: false,
                    },
                    visibility: ShaderStages::COMPUTE,
                },
            ],
        });

        let bind_groups = [0, 1].map(|i| {
            device.create_bind_group(&BindGroupDescriptor {
                label: None,
                layout: &bind_group_layout,
                entries: &[
                    BindGroupEntry {
                        binding: 0,
                        resource: uniform_buffer.as_entire_binding(),
                    },
                    BindGroupEntry {
                        binding: 1,
                        resource: bufs.offsets.as_entire_binding(),
                    },
                    BindGroupEntry {
                        binding: 2,
                        resource: wgpu::BindingResource::TextureView(
                            &bufs.textures[i].create_view(&Default::default()),
                        ),
                    },
                    BindGroupEntry {
                        binding: 3,
                        resource: wgpu::BindingResource::TextureView(
                            &bufs.textures[1 - i].create_view(&Default::default()),
                        ),
                    },
                ],
            })
        });

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

        let module = device.create_shader_module(include_wgsl!("enc.wgsl"));
        let pipeline = device.create_compute_pipeline(&ComputePipelineDescriptor {
            label: None,
            layout: Some(&pipeline_layout),
            module: &module,
            entry_point: "main",
        });

        Self {
            bind_groups,
            uniform,
            uniform_buffer,
            pipeline,
        }
    }

    pub fn write_uniforms(&self, queue: &Queue) {
        queue.write_buffer(
            &self.uniform_buffer,
            0,
            bytemuck::cast_slice(&[self.uniform]),
        )
    }

    pub fn pass(&self, encoder: &mut CommandEncoder, params: &Params, rp: &RoundParams) {
        let mut cpass = encoder.begin_compute_pass(&wgpu::ComputePassDescriptor {
            label: None,
            timestamp_writes: None,
        });
        cpass.set_pipeline(&self.pipeline);
        cpass.set_bind_group(0, &self.bind_groups[rp.swap], &[]);
        cpass.dispatch_workgroups(params.blocks_x as u32, params.blocks_y as u32, 1);
    }
}