pub mod debug;
pub mod dec;
pub mod enc;

use std::mem::size_of;
use wgpu::{
    Buffer, BufferUsages, CommandEncoder, Device, Extent3d, ImageCopyTexture, Origin3d, Queue,
    Texture, TextureDescriptor, TextureDimension, TextureFormat, TextureUsages,
};

pub struct Params {
    pub width: usize,
    pub height: usize,
    pub extent: Extent3d,
    pub block_width: usize,
    pub block_height: usize,
    pub blocks_x: usize,
    pub blocks_y: usize,
    pub blocks: usize,
}

pub struct RoundParams {
    pub swap: usize,
}

pub struct CommonBuffers {
    textures: [Texture; 2],
    offsets: Buffer,
    offsets_download: Option<Buffer>,
    debug_output: Option<Texture>,
    texture_download: Option<Buffer>,
}

#[repr(C)]
pub struct BlockOffset {
    score: f32,
    _pad: u32,
    offset: [f32; 2],
}

impl CommonBuffers {
    pub fn create(device: &Device, params: &Params) -> Self {
        let textures = [(), ()].map(|_| {
            device.create_texture(&TextureDescriptor {
                label: None,
                size: params.extent,
                mip_level_count: 1,
                sample_count: 1,
                dimension: TextureDimension::D2,
                format: TextureFormat::Bgra8Unorm,
                usage: TextureUsages::TEXTURE_BINDING
                    | TextureUsages::STORAGE_BINDING
                    | TextureUsages::COPY_DST
                    | TextureUsages::COPY_SRC,
                view_formats: &[],
            })
        });

        let debug_output = Some(device.create_texture(&TextureDescriptor {
            label: None,
            size: params.extent,
            mip_level_count: 1,
            sample_count: 1,
            dimension: TextureDimension::D2,
            format: TextureFormat::Bgra8Unorm,
            usage: TextureUsages::STORAGE_BINDING | TextureUsages::COPY_SRC,
            view_formats: &[],
        }));

        let texture_download = Some(device.create_buffer(&wgpu::BufferDescriptor {
            label: None,
            size: (params.width * params.height * 4) as u64,
            usage: BufferUsages::COPY_DST | BufferUsages::MAP_READ,
            mapped_at_creation: false,
        }));

        let offsets = device.create_buffer(&wgpu::BufferDescriptor {
            label: None,
            size: (params.blocks * size_of::<BlockOffset>()) as u64,
            usage: BufferUsages::COPY_DST | BufferUsages::STORAGE,
            mapped_at_creation: false,
        });
        let offsets_download = Some(device.create_buffer(&wgpu::BufferDescriptor {
            label: None,
            size: (params.blocks * size_of::<BlockOffset>()) as u64,
            usage: BufferUsages::COPY_DST | BufferUsages::MAP_READ,
            mapped_at_creation: false,
        }));

        Self {
            debug_output,
            textures,
            offsets_download,
            offsets,
            texture_download,
        }
    }

    pub fn upload(&self, queue: &Queue, params: &Params, rp: &RoundParams, buffer: &[u8]) {
        queue.write_texture(
            ImageCopyTexture {
                aspect: wgpu::TextureAspect::All,
                mip_level: 0,
                origin: Origin3d::ZERO,
                texture: &self.textures[rp.swap],
            },
            buffer,
            wgpu::ImageDataLayout {
                offset: 0,
                bytes_per_row: Some(params.extent.width * 4),
                rows_per_image: Some(params.extent.height),
            },
            params.extent,
        );
    }

    pub fn prepare_texture_download(
        &self,
        encoder: &mut CommandEncoder,
        params: &Params,
        rp: &RoundParams,
    ) {
        encoder.copy_texture_to_buffer(
            wgpu::ImageCopyTexture {
                texture: &self.textures[rp.swap],
                mip_level: 0,
                origin: wgpu::Origin3d::ZERO,
                aspect: wgpu::TextureAspect::All,
            },
            wgpu::ImageCopyBuffer {
                buffer: self.texture_download.as_ref().unwrap(),
                layout: wgpu::ImageDataLayout {
                    offset: 0,
                    bytes_per_row: Some(params.extent.width * 4),
                    rows_per_image: Some(params.extent.height),
                },
            },
            params.extent,
        );
    }

    pub fn download(&self, device: &Device, buffer: &mut [u8]) {
        let buffer_slice = self.texture_download.as_ref().unwrap().slice(..);
        let (sender, receiver) = oneshot::channel();
        buffer_slice.map_async(wgpu::MapMode::Read, move |r| sender.send(r).unwrap());
        device.poll(wgpu::Maintain::Wait);
        receiver.recv().unwrap().unwrap();
        {
            let view = buffer_slice.get_mapped_range();
            buffer.copy_from_slice(&view[..]);
        }
        self.texture_download.as_ref().unwrap().unmap();
    }
}