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

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));
    }
}