use std::io::{stdin, stdout, Read, Write};

use pollster::FutureExt;
use wgpu::{
    include_wgsl, Backends, BindGroupDescriptor, BindGroupEntry, BindGroupLayoutDescriptor,
    BindGroupLayoutEntry, BindingType, CommandEncoder, ComputePipelineDescriptor, DeviceDescriptor,
    Extent3d, Features, ImageCopyTexture, Instance, InstanceDescriptor, InstanceFlags, Limits,
    MaintainBase, Origin3d, PipelineLayoutDescriptor, PowerPreference, RequestAdapterOptions,
    ShaderStages, TextureDescriptor, TextureDimension, TextureFormat, TextureSampleType,
    TextureUsages, TextureViewDimension,
};

fn main() {
    let instance = Instance::new(InstanceDescriptor {
        backends: Backends::all(),
        ..Default::default()
    });

    let adapter = instance
        .request_adapter(&RequestAdapterOptions {
            power_preference: PowerPreference::HighPerformance,
            force_fallback_adapter: false,
            compatible_surface: None,
        })
        .block_on()
        .unwrap();

    let (device, queue) = adapter
        .request_device(
            &DeviceDescriptor {
                label: None,
                features: Features::BGRA8UNORM_STORAGE,
                limits: Limits::default(),
            },
            None,
        )
        .block_on()
        .unwrap();

    let (width, height) = (1920, 1080);
    let extent = Extent3d {
        width: width as u32,
        height: height as u32,
        depth_or_array_layers: 1,
    };

    let [texa, texb] = [(), ()].map(|_| {
        device.create_texture(&TextureDescriptor {
            label: None,
            size: extent,
            mip_level_count: 1,
            sample_count: 1,
            dimension: TextureDimension::D2,
            format: TextureFormat::Bgra8Unorm,
            usage: TextureUsages::TEXTURE_BINDING | TextureUsages::COPY_DST,
            view_formats: &[],
        })
    });
    let out_tex = device.create_texture(&TextureDescriptor {
        label: None,
        size: extent,
        mip_level_count: 1,
        sample_count: 1,
        dimension: TextureDimension::D2,
        format: TextureFormat::Bgra8Unorm,
        usage: TextureUsages::STORAGE_BINDING | TextureUsages::COPY_SRC,
        view_formats: &[],
    });

    let mut buffer = vec![0u8; width * height * 4];

    let download_buffer = device.create_buffer(&wgpu::BufferDescriptor {
        label: None,
        size: buffer.len() as u64,
        usage: wgpu::BufferUsages::COPY_DST | wgpu::BufferUsages::MAP_READ,
        mapped_at_creation: false,
    });

    let bind_group_layout = device.create_bind_group_layout(&BindGroupLayoutDescriptor {
        label: None,
        entries: &[
            BindGroupLayoutEntry {
                binding: 0,
                count: None,
                ty: wgpu::BindingType::StorageTexture {
                    access: wgpu::StorageTextureAccess::WriteOnly,
                    format: TextureFormat::Bgra8Unorm,
                    view_dimension: wgpu::TextureViewDimension::D2,
                },
                visibility: ShaderStages::COMPUTE,
            },
            BindGroupLayoutEntry {
                binding: 1,
                count: None,
                ty: BindingType::Texture {
                    sample_type: TextureSampleType::Float { filterable: false },
                    view_dimension: TextureViewDimension::D2,
                    multisampled: false,
                },
                visibility: ShaderStages::COMPUTE,
            },
        ],
    });

    let bind_goup = device.create_bind_group(&BindGroupDescriptor {
        label: None,
        layout: &bind_group_layout,
        entries: &[
            BindGroupEntry {
                binding: 0,
                resource: wgpu::BindingResource::TextureView(
                    &out_tex.create_view(&Default::default()),
                ),
            },
            BindGroupEntry {
                binding: 1,
                resource: wgpu::BindingResource::TextureView(
                    &texa.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!("shader.wgsl"));
    let pipeline = device.create_compute_pipeline(&ComputePipelineDescriptor {
        label: None,
        layout: Some(&pipeline_layout),
        module: &module,
        entry_point: "main",
    });

    loop {
        eprintln!("read");
        stdin().read_exact(&mut buffer).unwrap();

        eprintln!("upload");
        queue.write_texture(
            ImageCopyTexture {
                aspect: wgpu::TextureAspect::All,
                mip_level: 0,
                origin: Origin3d::ZERO,
                texture: &texa,
            },
            &buffer,
            wgpu::ImageDataLayout {
                offset: 0,
                bytes_per_row: Some(extent.width * 4),
                rows_per_image: Some(extent.height),
            },
            extent,
        );

        eprintln!("compute");
        let mut encoder = device.create_command_encoder(&Default::default());

        {
            let mut cpass = encoder.begin_compute_pass(&wgpu::ComputePassDescriptor {
                label: None,
                timestamp_writes: None,
            });
            cpass.set_pipeline(&pipeline);
            cpass.set_bind_group(0, &bind_goup, &[]);
            cpass.dispatch_workgroups(width as u32, height as u32, 1);
        }

        encoder.copy_texture_to_buffer(
            wgpu::ImageCopyTexture {
                texture: &out_tex,
                mip_level: 0,
                origin: wgpu::Origin3d::ZERO,
                aspect: wgpu::TextureAspect::All,
            },
            wgpu::ImageCopyBuffer {
                buffer: &download_buffer,
                layout: wgpu::ImageDataLayout {
                    offset: 0,
                    bytes_per_row: Some(extent.width * 4),
                    rows_per_image: Some(extent.height),
                },
            },
            extent,
        );

        queue.submit(Some(encoder.finish()));
        device.poll(MaintainBase::Wait);

        eprintln!("download");
        let buffer_slice = download_buffer.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[..]);
        }
        download_buffer.unmap();
        eprintln!("write");
        stdout().write_all(&buffer).unwrap();
    }
}