use crate::uniform::UniformCache;

const CANVAS_FORMAT: wgpu::TextureFormat = wgpu::TextureFormat::Rgba8UnormSrgb;

const VERTEX_SHADER: &str = "\
@group(0) @binding(0) var<uniform> _wgsl_resolution: vec2<f32>;

struct VertexOutput {
    @builtin(position) clip_position: vec4<f32>,
    @location(0) uv: vec2<f32>,
    @location(1) resolution: vec2<f32>,
};

@vertex
fn vs_main(@builtin(vertex_index) vi: u32) -> VertexOutput {
    var uv = vec2<f32>(f32(vi % 2u), f32(vi / 2u));
    var out: VertexOutput;
    out.clip_position = vec4<f32>(2.0 * uv - 1.0, 0.0, 1.0);
    out.uv = vec2<f32>(uv.x, 1.0 - uv.y);
    out.resolution = _wgsl_resolution;
    return out;
}
";

const DEFAULT_FRAGMENT: &str = "\
@fragment
fn fs_main(@location(0) uv: vec2<f32>) -> @location(0) vec4<f32> {
    let check = floor(uv * 10.0);
    return vec4<f32>(vec3<f32>(fract((check.x + check.y) / 2.0)), 1.0);
}
";

pub struct Renderer {
    canvas: wgpu::Texture,
    canvas_view: wgpu::TextureView,
    vertex_module: wgpu::ShaderModule,
    vertex_bgl: wgpu::BindGroupLayout,
    vertex_bg: wgpu::BindGroup,
    render_pipeline: wgpu::RenderPipeline,
    uniforms: UniformCache,
}

impl Renderer {
    pub fn new(device: &wgpu::Device, queue: &wgpu::Queue, width: u32, height: u32) -> Self {
        let canvas = device.create_texture(&wgpu::TextureDescriptor {
            label: Some("canvas"),
            size: wgpu::Extent3d {
                width,
                height,
                depth_or_array_layers: 1,
            },
            mip_level_count: 1,
            sample_count: 1,
            dimension: wgpu::TextureDimension::D2,
            format: CANVAS_FORMAT,
            usage: wgpu::TextureUsages::RENDER_ATTACHMENT
                | wgpu::TextureUsages::TEXTURE_BINDING
                | wgpu::TextureUsages::COPY_SRC,
            view_formats: &[],
        });
        let canvas_view = canvas.create_view(&Default::default());

        let vertex_module = device.create_shader_module(wgpu::ShaderModuleDescriptor {
            label: Some("vertex_shader"),
            source: wgpu::ShaderSource::Wgsl(VERTEX_SHADER.into()),
        });

        let resolution_buffer = device.create_buffer(&wgpu::BufferDescriptor {
            label: Some("resolution"),
            size: 8,
            usage: wgpu::BufferUsages::UNIFORM | wgpu::BufferUsages::COPY_DST,
            mapped_at_creation: false,
        });

        let vertex_bgl = device.create_bind_group_layout(&wgpu::BindGroupLayoutDescriptor {
            label: None,
            entries: &[wgpu::BindGroupLayoutEntry {
                binding: 0,
                visibility: wgpu::ShaderStages::VERTEX,
                ty: wgpu::BindingType::Buffer {
                    ty: wgpu::BufferBindingType::Uniform,
                    has_dynamic_offset: false,
                    min_binding_size: None,
                },
                count: None,
            }],
        });

        let vertex_bg = device.create_bind_group(&wgpu::BindGroupDescriptor {
            label: None,
            layout: &vertex_bgl,
            entries: &[wgpu::BindGroupEntry {
                binding: 0,
                resource: resolution_buffer.as_entire_binding(),
            }],
        });

        queue.write_buffer(
            &resolution_buffer,
            0,
            bytemuck::cast_slice(&[width as f32, height as f32]),
        );

        let mut uniforms = UniformCache::new();
        let render_pipeline = build_pipeline(
            device,
            &vertex_module,
            &vertex_bgl,
            DEFAULT_FRAGMENT,
            &mut uniforms,
        )
        .expect("default shader");

        Self {
            canvas,
            canvas_view,
            vertex_module,
            vertex_bgl,
            vertex_bg,
            render_pipeline,
            uniforms,
        }
    }

    pub fn canvas_texture(&self) -> &wgpu::Texture {
        &self.canvas
    }

    pub fn canvas_view(&self) -> &wgpu::TextureView {
        &self.canvas_view
    }

    pub fn uniforms(&mut self) -> &mut UniformCache {
        &mut self.uniforms
    }

    pub fn load_shader(&mut self, device: &wgpu::Device, source: &str) -> Result<(), String> {
        let pipeline = build_pipeline(
            device,
            &self.vertex_module,
            &self.vertex_bgl,
            source,
            &mut self.uniforms,
        )?;
        self.render_pipeline = pipeline;
        log::info!("shader loaded");
        Ok(())
    }

    pub fn render(&mut self, device: &wgpu::Device, queue: &wgpu::Queue) {
        self.uniforms.flush(queue);

        let mut encoder =
            device.create_command_encoder(&wgpu::CommandEncoderDescriptor { label: None });

        {
            let mut pass = encoder.begin_render_pass(&wgpu::RenderPassDescriptor {
                label: None,
                color_attachments: &[Some(wgpu::RenderPassColorAttachment {
                    view: &self.canvas_view,
                    depth_slice: None,
                    resolve_target: None,
                    ops: wgpu::Operations {
                        load: wgpu::LoadOp::Clear(wgpu::Color::BLACK),
                        store: wgpu::StoreOp::Store,
                    },
                })],
                depth_stencil_attachment: None,
                occlusion_query_set: None,
                timestamp_writes: None,
                multiview_mask: None,
            });

            pass.set_pipeline(&self.render_pipeline);
            pass.set_bind_group(0, &self.vertex_bg, &[]);
            if let Some(ref bg) = self.uniforms.bind_group {
                pass.set_bind_group(1, bg, &[]);
            }
            pass.draw(0..4, 0..1);
        }

        queue.submit([encoder.finish()]);
    }
}

fn build_pipeline(
    device: &wgpu::Device,
    vertex_module: &wgpu::ShaderModule,
    vertex_bgl: &wgpu::BindGroupLayout,
    fragment_source: &str,
    cache: &mut UniformCache,
) -> Result<wgpu::RenderPipeline, String> {
    let module = naga::front::wgsl::parse_str(fragment_source)
        .map_err(|e| format!("WGSL parse error: {e}"))?;

    naga::valid::Validator::new(
        naga::valid::ValidationFlags::all(),
        naga::valid::Capabilities::all(),
    )
    .validate(&module)
    .map_err(|e| format!("WGSL validation error: {e}"))?;

    cache.refresh(module, device);
    cache.rebuild_bind_group(device);

    let fragment_module = device.create_shader_module(wgpu::ShaderModuleDescriptor {
        label: Some("fragment_shader"),
        source: wgpu::ShaderSource::Wgsl(fragment_source.into()),
    });

    let mut bind_group_layouts: Vec<Option<&wgpu::BindGroupLayout>> = vec![Some(vertex_bgl)];
    if let Some(ref bgl) = cache.bind_group_layout {
        bind_group_layouts.push(Some(bgl));
    }

    let pipeline_layout = device.create_pipeline_layout(&wgpu::PipelineLayoutDescriptor {
        label: None,
        bind_group_layouts: &bind_group_layouts,
        immediate_size: 0,
    });

    Ok(
        device.create_render_pipeline(&wgpu::RenderPipelineDescriptor {
            label: None,
            layout: Some(&pipeline_layout),
            vertex: wgpu::VertexState {
                module: vertex_module,
                entry_point: Some("vs_main"),
                buffers: &[],
                compilation_options: Default::default(),
            },
            fragment: Some(wgpu::FragmentState {
                module: &fragment_module,
                entry_point: Some("fs_main"),
                targets: &[Some(wgpu::ColorTargetState {
                    format: CANVAS_FORMAT,
                    blend: Some(wgpu::BlendState::REPLACE),
                    write_mask: wgpu::ColorWrites::ALL,
                })],
                compilation_options: Default::default(),
            }),
            primitive: wgpu::PrimitiveState {
                topology: wgpu::PrimitiveTopology::TriangleStrip,
                ..Default::default()
            },
            depth_stencil: None,
            multisample: Default::default(),
            multiview_mask: None,
            cache: None,
        }),
    )
}