use std::collections::HashMap;
use tracing::debug;
use wgpu::{BindGroupDescriptor, BindGroupEntry, BindingResource, Color, ColorTargetState, ColorWrites, CommandEncoderDescriptor, Device, FilterMode, FragmentState, include_wgsl, LoadOp, Operations, Queue, RenderPassColorAttachment, RenderPassDescriptor, RenderPipeline, RenderPipelineDescriptor, Sampler, SamplerDescriptor, ShaderModule, StoreOp, TextureFormat, TextureViewDescriptor, VertexState};
use crate::rendering::texture::Texture;

pub struct MipGenerator {
    shader: ShaderModule,
    sampler: Sampler,
    pipelines: HashMap<TextureFormat, RenderPipeline>,
}
impl MipGenerator {
    pub fn new(device: &Device) -> Self {
        debug!("initializing MipGenerator, compiling shader module");
        let shader =
            device.create_shader_module(include_wgsl!("../shaders/text_quad_mips.wgsl"));
        Self {
            shader,
            sampler: device.create_sampler(&SamplerDescriptor {
                min_filter: FilterMode::Linear,
                label: Some("MipGenerator sampler"),
                ..Default::default()
            }),
            pipelines: HashMap::new(),
        }
    }

    pub fn generate_mips(&mut self, texture: &Texture, device: &Device, queue: &Queue) {
        let pipeline = self.pipelines.entry(texture.texture.format()).or_insert_with(|| {
            device.create_render_pipeline(&RenderPipelineDescriptor {
                label: Some("MipGenerator format pipeline"),
                layout: None,
                vertex: VertexState {
                    module: &self.shader,
                    entry_point: Some("vs"),
                    compilation_options: Default::default(),
                    buffers: &[],
                },
                primitive: Default::default(),
                depth_stencil: None,
                multisample: Default::default(),
                fragment: Some(FragmentState {
                    module: &self.shader,
                    entry_point: Some("fs"),
                    compilation_options: Default::default(),
                    targets: &[Some(ColorTargetState {
                        format: texture.texture.format(),
                        blend: None,
                        write_mask: ColorWrites::default(),
                    })],
                }),
                multiview: None,
                cache: None,
            })
        });

        let mut encoder = device.create_command_encoder(&CommandEncoderDescriptor {
            label: Some("MipGenerator command encoder"),
        });

        let mut width = texture.texture.width();
        let mut height = texture.texture.height();
        let mut base_mip_level = 0;

        while width > 1 || height > 1 {
            width = 1.max(width / 2);
            height = 1.max(height / 2);

            let bind_group = device.create_bind_group(&BindGroupDescriptor {
                label: Some("MipGenerator bind group"),
                layout: &pipeline.get_bind_group_layout(0),
                entries: &[
                    BindGroupEntry {
                        binding: 0,
                        resource: BindingResource::Sampler(&self.sampler),
                    },
                    BindGroupEntry {
                        binding: 1,
                        resource: BindingResource::TextureView(&texture.texture.create_view(
                            &TextureViewDescriptor {
                                base_mip_level,
                                mip_level_count: Some(1),
                                ..Default::default()
                            },
                        )),
                    },
                ],
            });

            base_mip_level += 1;

            let texture_view = texture.texture.create_view(&TextureViewDescriptor {
                base_mip_level,
                mip_level_count: Some(1),
                ..Default::default()
            });

            let render_pass_descriptor = RenderPassDescriptor {
                label: Some("MipGenerator render pass"),
                color_attachments: &[Some(RenderPassColorAttachment {
                    view: &texture_view,
                    resolve_target: None,
                    ops: Operations {
                        load: LoadOp::Clear(Color {
                            r: 0.0,
                            g: 0.0,
                            b: 0.0,
                            a: 1.0,
                        }),
                        store: StoreOp::Store,
                    },
                })],
                depth_stencil_attachment: None,
                timestamp_writes: None,
                occlusion_query_set: None,
            };

            let mut pass = encoder.begin_render_pass(&render_pass_descriptor);
            pass.set_pipeline(pipeline);
            pass.set_bind_group(0, Some(&bind_group), &[]);
            pass.draw(0..6, 0..1);
        }

        let command_buffer = encoder.finish();
        queue.submit(std::iter::once(command_buffer));
    }
}