use bevy::{
    core_pipeline::fullscreen_vertex_shader,
    prelude::*,
    render::{
        render_resource::{
            binding_types::{sampler, texture_2d},
            BindGroupLayout, BindGroupLayoutEntries, ColorTargetState, ColorWrites, FragmentState,
            RenderPipelineDescriptor, Sampler, SamplerBindingType, SamplerDescriptor, ShaderStages,
            SpecializedRenderPipeline, TextureFormat, TextureSampleType,
        },
        renderer::RenderDevice,
    },
};

#[derive(Resource)]
pub struct DitherPostProcessPipeline {
    pub layout: BindGroupLayout,
    pub screen_sampler: Sampler,
    pub threshold_map_sampler: Sampler,
    pub shader: Handle<Shader>,
}

impl FromWorld for DitherPostProcessPipeline {
    fn from_world(world: &mut World) -> Self {
        let render_device = world.resource::<RenderDevice>();

        let layout = render_device.create_bind_group_layout(
            "dither_post_process_bind_group_layout",
            &BindGroupLayoutEntries::sequential(
                ShaderStages::FRAGMENT,
                (
                    texture_2d(TextureSampleType::Float { filterable: true }),
                    sampler(SamplerBindingType::Filtering),
                    texture_2d(TextureSampleType::Float { filterable: true }),
                    sampler(SamplerBindingType::Filtering),
                ),
            ),
        );

        let screen_sampler = render_device.create_sampler(&SamplerDescriptor::default());
        let threshold_map_sampler = render_device.create_sampler(&SamplerDescriptor::default());

        let shader = world.resource::<AssetServer>().load::<Shader>(
            "embedded://bevy_dither_post_process/../assets/shaders/dither_post_process.wgsl",
        );

        Self {
            layout,
            screen_sampler,
            threshold_map_sampler,
            shader,
        }
    }
}

#[derive(Clone, Copy, PartialEq, Eq, Hash)]
pub struct DitherPostProcessingPipelineKey {
    pub texture_format: TextureFormat,
}

impl SpecializedRenderPipeline for DitherPostProcessPipeline {
    type Key = DitherPostProcessingPipelineKey;

    fn specialize(&self, key: Self::Key) -> RenderPipelineDescriptor {
        RenderPipelineDescriptor {
            label: Some("dither_post_processing".into()),
            layout: vec![self.layout.clone()],
            vertex: fullscreen_vertex_shader::fullscreen_shader_vertex_state(),
            fragment: Some(FragmentState {
                shader: self.shader.clone(),
                shader_defs: vec![],
                entry_point: "fragment".into(),
                targets: vec![Some(ColorTargetState {
                    format: key.texture_format,
                    blend: None,
                    write_mask: ColorWrites::ALL,
                })],
            }),
            primitive: default(),
            depth_stencil: None,
            multisample: default(),
            push_constant_ranges: vec![],
            zero_initialize_workgroup_memory: false,
        }
    }
}