use crate::app::App;
use image::RgbaImage;
use std::{num::NonZeroU32, sync::Arc};
use wgpu::{
    Buffer, Extent3d, ImageCopyBuffer, ImageCopyTexture, ImageDataLayout, Origin3d, Texture,
};

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

    padded_bytes_per_row: u32,
    export_buffer: Buffer,
}

impl Exporter {
    pub fn new(app: &Arc<App>, size: Extent3d) -> Self {
        let App { device, .. } = app.as_ref();

        let bytes_per_pixel = std::mem::size_of::<u32>() as u32;
        let unpadded_bytes_per_row = size.width * bytes_per_pixel;
        let align = wgpu::COPY_BYTES_PER_ROW_ALIGNMENT;
        let padded_bytes_per_row_padding = (align - unpadded_bytes_per_row % align) % align;
        let padded_bytes_per_row = unpadded_bytes_per_row + padded_bytes_per_row_padding;

        let export_buffer_size = (padded_bytes_per_row * size.height) as u64;
        let export_buffer = device.create_buffer(&wgpu::BufferDescriptor {
            label: None,
            size: export_buffer_size,
            usage: wgpu::BufferUsages::MAP_READ | wgpu::BufferUsages::COPY_DST,
            mapped_at_creation: false,
        });

        Self {
            padded_bytes_per_row,
            app: app.clone(),
            size,
            export_buffer,
        }
    }
    pub async fn run(&self, texture: &Texture, save_path: &str) {
        let App { device, queue, .. } = self.app.as_ref();

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

        encoder.copy_texture_to_buffer(
            ImageCopyTexture {
                texture,
                aspect: wgpu::TextureAspect::All,
                mip_level: 0,
                origin: Origin3d::ZERO,
            },
            ImageCopyBuffer {
                buffer: &self.export_buffer,
                layout: ImageDataLayout {
                    offset: 0,
                    bytes_per_row: Some(NonZeroU32::new(self.padded_bytes_per_row).unwrap()),
                    rows_per_image: None,
                },
            },
            self.size,
        );

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

        let buffer_slice = self.export_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: Vec<u8> = bytemuck::cast_slice(&data).to_vec();
        let image = RgbaImage::from_raw(self.size.width, self.size.height, result).unwrap();
        image.save(save_path).unwrap();

        drop(data);
        self.export_buffer.unmap();
    }
}