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_size: u64,
    uniform_buffer: Buffer,
    uniform_staging_buffer: Buffer,
}

impl Painter {
    pub fn new(app: &Arc<App>, extent: Extent3d, tex_a: &Texture, tex_b: &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::<u32>() as wgpu::BufferAddress;
        let uniform_staging_buffer = device.create_buffer(&wgpu::BufferDescriptor {
            label: None,
            size: uniform_buffer_size,
            usage: wgpu::BufferUsages::MAP_READ | wgpu::BufferUsages::COPY_DST,
            mapped_at_creation: false,
        });
        let uniform_buffer = device.create_buffer_init(&wgpu::util::BufferInitDescriptor {
            label: None,
            contents: bytemuck::cast_slice(&[0u32]),
            usage: wgpu::BufferUsages::STORAGE
                | wgpu::BufferUsages::COPY_DST
                | wgpu::BufferUsages::COPY_SRC,
        });

        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(
                        &tex_a.create_view(&wgpu::TextureViewDescriptor::default()),
                    ),
                },
                wgpu::BindGroupEntry {
                    binding: 1,
                    resource: wgpu::BindingResource::TextureView(
                        &tex_b.create_view(&wgpu::TextureViewDescriptor::default()),
                    ),
                },
                wgpu::BindGroupEntry {
                    binding: 2,
                    resource: uniform_buffer.as_entire_binding(),
                },
            ],
        });
        Self {
            app: app.clone(),
            size: extent,
            pipeline,
            uniform_buffer_size,
            uniform_buffer,
            bind_group,
            uniform_staging_buffer,
        }
    }

    pub async fn run(&mut self) -> u32 {
        let App { device, queue, .. } = self.app.as_ref();

        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);
        }

        encoder.copy_buffer_to_buffer(
            &self.uniform_buffer,
            0,
            &self.uniform_staging_buffer,
            0,
            self.uniform_buffer_size,
        );

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

        let buffer_slice = self.uniform_staging_buffer.slice(..);
        let (sender, receiver) = futures_intrusive::channel::shared::oneshot_channel();
        buffer_slice.map_async(wgpu::MapMode::Read, move |v| {
            sender.send(v.unwrap()).unwrap()
        });
        device.poll(wgpu::Maintain::Wait);
        receiver.receive().await;

        let data = buffer_slice.get_mapped_range();
        let result: u32 = bytemuck::cast_slice(&data).to_vec()[0];

        drop(data);
        self.uniform_staging_buffer.unmap();

        result
    }
}