use bytemuck::{Pod, Zeroable};
use std::{borrow::Cow, sync::Arc};
use wgpu::{util::DeviceExt, BindGroup, Buffer, ComputePipeline, Extent3d, Texture};

use crate::app::App;

pub struct Painter {
    app: Arc<App>,
    size: Extent3d,

    pipeline: ComputePipeline,
    bind_group: BindGroup,

    uniform_buffer: Buffer,
}

#[repr(C)]
#[derive(Pod, Zeroable, Clone, Copy)]
pub struct PaintUniforms {
    pub x: f32,
    pub y: f32,
    pub rx: f32,
    pub ry: f32,
    pub r: f32,
    pub g: f32,
    pub b: f32,
}

impl Painter {
    pub fn new(app: &Arc<App>, extent: Extent3d, texture: &Texture) -> Self {
        let App { device, .. } = app.as_ref();
        let cs_module = device.create_shader_module(wgpu::ShaderModuleDescriptor {
            label: None,
            source: wgpu::ShaderSource::Wgsl(Cow::Borrowed(include_str!("paint.wgsl"))),
        });

        // let uniform_buffer_size = std::mem::size_of::<PaintUniforms>() as wgpu::BufferAddress;
        let uniform_buffer = device.create_buffer_init(&wgpu::util::BufferInitDescriptor {
            label: None,
            contents: bytemuck::cast_slice(&[PaintUniforms::zeroed()]),
            usage: wgpu::BufferUsages::UNIFORM | wgpu::BufferUsages::COPY_DST,
        });

        let pipeline = device.create_compute_pipeline(&wgpu::ComputePipelineDescriptor {
            label: None,
            layout: None,
            module: &cs_module,
            entry_point: "main",
        });

        let bind_group = device.create_bind_group(&wgpu::BindGroupDescriptor {
            label: None,
            layout: &pipeline.get_bind_group_layout(0),
            entries: &[
                wgpu::BindGroupEntry {
                    binding: 0,
                    resource: wgpu::BindingResource::TextureView(
                        &texture.create_view(&wgpu::TextureViewDescriptor::default()),
                    ),
                },
                wgpu::BindGroupEntry {
                    binding: 1,
                    resource: uniform_buffer.as_entire_binding(),
                },
            ],
        });
        Self {
            app: app.clone(),
            size: extent,
            pipeline,
            uniform_buffer,
            bind_group,
        }
    }

    pub fn run(&self, params: PaintUniforms) {
        let App { device, queue, .. } = self.app.as_ref();

        queue.write_buffer(&self.uniform_buffer, 0, bytemuck::cast_slice(&[params]));

        let mut encoder =
            device.create_command_encoder(&wgpu::CommandEncoderDescriptor { label: None });
        {
            let mut cpass =
                encoder.begin_compute_pass(&wgpu::ComputePassDescriptor { label: None });
            cpass.set_pipeline(&self.pipeline);
            cpass.set_bind_group(0, &self.bind_group, &[]);
            cpass.dispatch_workgroups(self.size.width, self.size.height, 1);
        }
        queue.submit(Some(encoder.finish()));
    }
}