use eframe::{
    egui_wgpu::wgpu::util::DeviceExt as _,
    egui_wgpu::{self, wgpu},
};
use std::num::NonZeroU64;

pub struct Preview {
    ready: bool,
}

impl Preview {
    pub fn new<'a>(cc: &'a eframe::CreationContext<'a>) -> Option<Self> {
        let state = cc.wgpu_render_state.as_ref()?;
        let resources = &mut state.renderer.write().callback_resources;
        resources.insert(TriangleRenderResources::init_with_shader(
            &state.device,
            state.target_format,
            "#version 450
in vec2 in_uv;
out vec4 out_color;

void main() {
  out_color = vec4(1, 0, 1, 1.0);
}",
        ));

        Some(Self { ready: false })
    }

    pub fn update(&mut self, ui: &mut egui::Ui, frag_shader: Option<String>) -> bool {
        ui.set_min_width(150.0);
        let size = ui.available_size().min_elem();
        let (rect, _) = ui.allocate_exact_size(egui::Vec2::splat(size), egui::Sense::empty());

        let have_shader = frag_shader.is_some();

        ui.painter().add(egui_wgpu::Callback::new_paint_callback(
            rect,
            CustomTriangleCallback {
                frag_shader,
                time: ui.ctx().input(|i| i.time) as f32,
            },
        ));

        ui.ctx().request_repaint();

        if !self.ready {
            self.ready = true;
            false
        } else {
            have_shader && !ui.ctx().will_discard()
        }
    }

    pub fn reset(&mut self) {
        self.ready = false;
    }
}

struct CustomTriangleCallback {
    frag_shader: Option<String>,
    time: f32,
}

impl egui_wgpu::CallbackTrait for CustomTriangleCallback {
    fn prepare(
        &self,
        device: &wgpu::Device,
        queue: &wgpu::Queue,
        screen: &egui_wgpu::ScreenDescriptor,
        _egui_encoder: &mut wgpu::CommandEncoder,
        resources: &mut egui_wgpu::CallbackResources,
    ) -> Vec<wgpu::CommandBuffer> {
        let triangle_resources: &mut TriangleRenderResources = resources.get_mut().unwrap();
        if let Some(code) = &self.frag_shader {
            *triangle_resources = TriangleRenderResources::init_with_shader(
                device,
                triangle_resources.target_format,
                code,
            );
        }
        triangle_resources.prepare(device, queue, screen, self.time);

        Vec::new()
    }

    fn paint(
        &self,
        _info: egui::PaintCallbackInfo,
        render_pass: &mut wgpu::RenderPass<'static>,
        resources: &egui_wgpu::CallbackResources,
    ) {
        let resources: &TriangleRenderResources = resources.get().unwrap();
        resources.paint(render_pass);
    }
}

struct TriangleRenderResources {
    pipeline: wgpu::RenderPipeline,
    bind_group: wgpu::BindGroup,
    target_format: wgpu::TextureFormat,
    uniform_buffer: wgpu::Buffer,
}

impl TriangleRenderResources {
    pub fn init_with_shader(
        device: &wgpu::Device,
        target_format: wgpu::TextureFormat,
        fragment: &str,
    ) -> Self {
        let vertex_shader = device.create_shader_module(wgpu::ShaderModuleDescriptor {
            label: Some("custom3d"),
            source: wgpu::ShaderSource::Glsl {
                shader: "#version 450
out vec2 out_uv;

const vec2 v_positions[3] = vec2[3](
    vec2( 0.0,  3.0),
    vec2( 3.0, -3.0),
    vec2(-3.0, -3.0)
);

const vec4 v_colors[3] = vec4[3](
    vec4(1.0, 0.0, 0.0, 1.0),
    vec4(0.0, 1.0, 0.0, 1.0),
    vec4(0.0, 0.0, 1.0, 1.0)
);

void main() {
    vec2 pos = v_positions[gl_VertexIndex];
    gl_Position = vec4(pos, 0.0, 1.0);
    out_uv = gl_Position.xy / 2.0 + 0.5;
}"
                .into(),
                stage: wgpu::naga::ShaderStage::Vertex,
                defines: &[],
            },
        });

        let fragment_shader = device.create_shader_module(wgpu::ShaderModuleDescriptor {
            label: Some("custom3d"),
            source: wgpu::ShaderSource::Glsl {
                shader: fragment.into(),
                stage: wgpu::naga::ShaderStage::Fragment,
                defines: &[],
            },
        });
        let bind_group_layout = device.create_bind_group_layout(&wgpu::BindGroupLayoutDescriptor {
            label: Some("custom3d"),
            entries: &[wgpu::BindGroupLayoutEntry {
                binding: 0,
                visibility: wgpu::ShaderStages::FRAGMENT,
                ty: wgpu::BindingType::Buffer {
                    ty: wgpu::BufferBindingType::Uniform,
                    has_dynamic_offset: false,
                    min_binding_size: NonZeroU64::new(16),
                },
                count: None,
            }],
        });

        let pipeline_layout = device.create_pipeline_layout(&wgpu::PipelineLayoutDescriptor {
            label: Some("custom3d"),
            bind_group_layouts: &[&bind_group_layout],
            push_constant_ranges: &[],
        });

        let pipeline = device.create_render_pipeline(&wgpu::RenderPipelineDescriptor {
            label: Some("custom3d"),
            layout: Some(&pipeline_layout),
            vertex: wgpu::VertexState {
                module: &vertex_shader,
                entry_point: None,
                buffers: &[],
                compilation_options: wgpu::PipelineCompilationOptions::default(),
            },
            fragment: Some(wgpu::FragmentState {
                module: &fragment_shader,
                entry_point: None,
                targets: &[Some(target_format.into())],
                compilation_options: wgpu::PipelineCompilationOptions::default(),
            }),
            primitive: wgpu::PrimitiveState::default(),
            depth_stencil: None,
            multisample: wgpu::MultisampleState::default(),
            multiview: None,
            cache: None,
        });

        let uniform_buffer = device.create_buffer_init(&wgpu::util::BufferInitDescriptor {
            label: Some("custom3d"),
            contents: bytemuck::cast_slice(&[0.0_f32; 4]), // 16 bytes aligned!
            // Mapping at creation (as done by the create_buffer_init utility) doesn't require us to to add the MAP_WRITE usage
            // (this *happens* to workaround this bug )
            usage: wgpu::BufferUsages::COPY_DST | wgpu::BufferUsages::UNIFORM,
        });

        let bind_group = device.create_bind_group(&wgpu::BindGroupDescriptor {
            label: Some("custom3d"),
            layout: &bind_group_layout,
            entries: &[wgpu::BindGroupEntry {
                binding: 0,
                resource: uniform_buffer.as_entire_binding(),
            }],
        });

        Self {
            pipeline,
            bind_group,
            target_format,
            uniform_buffer,
        }
    }

    fn prepare(
        &mut self,
        _device: &wgpu::Device,
        queue: &wgpu::Queue,
        screen: &egui_wgpu::ScreenDescriptor,
        time: f32,
    ) {
        queue.write_buffer(
            &self.uniform_buffer,
            0,
            bytemuck::cast_slice(&[
                screen.size_in_pixels[0] as f32,
                screen.size_in_pixels[1] as f32,
                time,
                0.0,
            ]),
        );
    }

    fn paint(&self, render_pass: &mut wgpu::RenderPass<'_>) {
        // Draw our triangle!
        render_pass.set_pipeline(&self.pipeline);
        render_pass.set_bind_group(0, &self.bind_group, &[]);
        render_pass.draw(0..3, 0..1);
    }
}