use crate::engine::Engine;
use crate::prelude::renderer_types::{TextureResource, Vertex};
use crate::prelude::wgpu::util::DeviceExt;
use crate::prelude::wgpu::{self, Device, RenderPass};
use bytemuck;

use crate::render::BackendState;
use crate::renderer::traits::PostShader;

struct State {
    device: Device,
    pipeline: wgpu::RenderPipeline,
    bind_group_layout: wgpu::BindGroupLayout,
    vertex_buffer: wgpu::Buffer,
    index_buffer: wgpu::Buffer,
}

impl State {
    pub fn new(param: &BackendState) -> Self {
        let device = &param.device;
        let format = param.surface_view_format;
        let shader = device.create_shader_module(wgpu::ShaderModuleDescriptor {
            label: Some("screen shader"),
            source: wgpu::ShaderSource::Wgsl(
                include_str!("../../../assets/shaders/screen.wgsl").into(),
            ),
        });
        let bind_group_layout = device.create_bind_group_layout(&wgpu::BindGroupLayoutDescriptor {
            label: Some("Screen Bind Group Layout"),
            entries: &[
                wgpu::BindGroupLayoutEntry {
                    binding: 0,
                    visibility: wgpu::ShaderStages::FRAGMENT,
                    ty: wgpu::BindingType::Texture {
                        sample_type: wgpu::TextureSampleType::Float { filterable: true },
                        view_dimension: wgpu::TextureViewDimension::D2,
                        multisampled: false,
                    },
                    count: None,
                },
                wgpu::BindGroupLayoutEntry {
                    binding: 1,
                    visibility: wgpu::ShaderStages::FRAGMENT,
                    ty: wgpu::BindingType::Sampler(wgpu::SamplerBindingType::Filtering),
                    count: None,
                },
            ],
        });
        let pipeline_layout = device.create_pipeline_layout(&wgpu::PipelineLayoutDescriptor {
            label: Some("screen layout"),
            bind_group_layouts: &[&bind_group_layout],
            push_constant_ranges: &[],
        });
        let pipeline = device.create_render_pipeline(&wgpu::RenderPipelineDescriptor {
            cache: None,
            label: Some("screen pipeline"),
            layout: Some(&pipeline_layout),
            vertex: wgpu::VertexState {
                module: &shader,
                entry_point: Some("vs_main"),
                buffers: &[Vertex::desc()],
                compilation_options: Default::default(),
            },
            fragment: Some(wgpu::FragmentState {
                module: &shader,
                entry_point: Some("fs_main"),
                targets: &[Some(wgpu::ColorTargetState {
                    format,
                    blend: Some(wgpu::BlendState::ALPHA_BLENDING),
                    write_mask: wgpu::ColorWrites::ALL,
                })],
                compilation_options: Default::default(),
            }),
            primitive: wgpu::PrimitiveState {
                topology: wgpu::PrimitiveTopology::TriangleList,
                strip_index_format: None,
                front_face: wgpu::FrontFace::Ccw,
                cull_mode: Some(wgpu::Face::Back),
                polygon_mode: wgpu::PolygonMode::Fill,
                unclipped_depth: false,
                conservative: false,
            },
            depth_stencil: None,
            multisample: wgpu::MultisampleState::default(),
            multiview: None,
        });
        let vertex_buffer = device.create_buffer_init(&wgpu::util::BufferInitDescriptor {
            label: Some("screen vertex buffer"),
            contents: bytemuck::cast_slice(&[
                Vertex {
                    pos: [-1.0, -1.0, 0.0],
                    color: [1.0, 1.0, 1.0, 1.0],
                    tex_coord: [0.0, 1.0],
                },
                Vertex {
                    pos: [1.0, -1.0, 0.0],
                    color: [1.0, 1.0, 1.0, 1.0],
                    tex_coord: [1.0, 1.0],
                },
                Vertex {
                    pos: [1.0, 1.0, 0.0],
                    color: [1.0, 1.0, 1.0, 1.0],
                    tex_coord: [1.0, 0.0],
                },
                Vertex {
                    pos: [-1.0, 1.0, 0.0],
                    color: [1.0, 1.0, 1.0, 1.0],
                    tex_coord: [0.0, 0.0],
                },
            ]),
            usage: wgpu::BufferUsages::VERTEX,
        });

        let index_buffer = device.create_buffer_init(&wgpu::util::BufferInitDescriptor {
            label: Some("screen index buffer"),
            contents: bytemuck::cast_slice(&[0u16, 1, 2, 0, 2, 3]),
            usage: wgpu::BufferUsages::INDEX,
        });
        Self {
            device: device.clone(),
            pipeline,
            bind_group_layout,
            vertex_buffer,
            index_buffer,
        }
    }
}

use std::cell::RefCell;

pub struct ScreenShader {
    state: State,
    cached_bind_group: RefCell<Option<wgpu::BindGroup>>,
}

impl ScreenShader {
    pub fn setup(param: &BackendState) -> Self {
        let state = State::new(param);
        Self {
            state,
            cached_bind_group: RefCell::new(None),
        }
    }
}

impl PostShader for ScreenShader {
    fn resize(&self, _g: &Engine) {
        // Drop cached bind group so it will be recreated with the new render target view
        self.cached_bind_group.borrow_mut().take();
    }

    fn draw(&self, texture: &TextureResource, pass: &mut RenderPass) {
        if self.cached_bind_group.borrow().is_none() {
            let bg = self
                .state
                .device
                .create_bind_group(&wgpu::BindGroupDescriptor {
                    layout: &self.state.bind_group_layout,
                    entries: &[
                        wgpu::BindGroupEntry {
                            binding: 0,
                            resource: wgpu::BindingResource::TextureView(&texture.view),
                        },
                        wgpu::BindGroupEntry {
                            binding: 1,
                            resource: wgpu::BindingResource::Sampler(&texture.sampler),
                        },
                    ],
                    label: Some("Screen bind group"),
                });
            self.cached_bind_group.borrow_mut().replace(bg);
        }
        let bind_group = self.cached_bind_group.borrow();
        let bind_group = bind_group.as_ref().unwrap();
        pass.set_pipeline(&self.state.pipeline);
        pass.set_bind_group(0, bind_group, &[]);
        pass.set_vertex_buffer(0, self.state.vertex_buffer.slice(..));
        pass.set_index_buffer(self.state.index_buffer.slice(..), wgpu::IndexFormat::Uint16);
        pass.draw_indexed(0..6, 0, 0..1);
    }

    fn frame_pass(&self, _g: &Engine, _pass: &mut RenderPass) {}

    fn frame_end(&self, _pass: &mut RenderPass) {}
}
