tirbofish/dropbear

use egui::{Context, FullOutput};
use egui_wgpu::wgpu::{CommandEncoder, Device, Queue, StoreOp, TextureFormat, TextureView};
use egui_wgpu::{Renderer, RendererOptions, ScreenDescriptor, wgpu};
use egui_winit::State;
use winit::event::WindowEvent;
use winit::window::Window;

pub struct EguiRenderer {
    state: State,
    renderer: Renderer,
    frame_started: bool,
}

impl EguiRenderer {
    pub fn context(&self) -> &Context {
        self.state.egui_ctx()
    }

    pub fn renderer(&mut self) -> &mut Renderer {
        &mut self.renderer
    }

    pub fn new(
        device: &Device,
        output_color_format: TextureFormat,
        output_depth_format: Option<TextureFormat>,
        msaa_samples: u32,
        window: &Window,
    ) -> EguiRenderer {
        puffin::profile_function!();
        let egui_context = Context::default();

        let egui_state = egui_winit::State::new(
            egui_context,
            egui::viewport::ViewportId::ROOT,
            &window,
            Some(window.scale_factor() as f32),
            None,
            Some(2 * 1024), // default dimension is 2048
        );
        let options = RendererOptions {
            msaa_samples,
            depth_stencil_format: output_depth_format,
            dithering: true,
            predictable_texture_filtering: false,
        };
        let egui_renderer = Renderer::new(device, output_color_format, options);

        EguiRenderer {
            state: egui_state,
            renderer: egui_renderer,
            frame_started: false,
        }
    }

    pub fn handle_input(&mut self, window: &Window, event: &WindowEvent) {
        let _ = self.state.on_window_event(window, event);
    }

    pub fn ppp(&mut self, v: f32) {
        self.context().set_pixels_per_point(v);
    }

    pub fn begin_frame(&mut self, window: &Window) {
        let raw_input = self.state.take_egui_input(window);
        self.state.egui_ctx().begin_pass(raw_input);
        self.frame_started = true;
    }

    pub fn take_input(&mut self, window: &Window) -> (egui::RawInput, egui::Context) {
        egui_extras::install_image_loaders(self.state.egui_ctx());
        let raw_input = self.state.take_egui_input(window);
        let ctx = self.state.egui_ctx().clone();
        (raw_input, ctx)
    }

    pub fn process_output(
        &mut self,
        full_output: egui::FullOutput,
        device: &Device,
        queue: &Queue,
        window: &Window,
        window_surface_view: &TextureView,
        screen_descriptor: ScreenDescriptor,
    ) -> CommandEncoder {
        let mut encoder = device
            .create_command_encoder(&wgpu::CommandEncoderDescriptor {
                label: Some("egui render encoder"),
            });

        self.handle_full_output(
            full_output,
            device,
            queue,
            &mut encoder,
            window,
            window_surface_view,
            screen_descriptor,
        );

        encoder
    }

    pub fn run_ui(
        &mut self,
        run_ui: impl FnMut(&mut egui::Ui),

        device: &Device,
        queue: &Queue,
        window: &Window,
        window_surface_view: &TextureView,
        screen_descriptor: ScreenDescriptor,
    ) -> CommandEncoder {
        let (raw_input, ctx) = self.take_input(window);
        self.frame_started = true;
        let full_output = ctx.run_ui(raw_input, run_ui);
        self.frame_started = false;

        self.process_output(full_output, device, queue, window, window_surface_view, screen_descriptor)
    }

    pub fn end_frame_and_draw(
        &mut self,
        device: &Device,
        queue: &Queue,
        encoder: &mut CommandEncoder,
        window: &Window,
        window_surface_view: &TextureView,
        screen_descriptor: ScreenDescriptor,
    ) {
        puffin::profile_function!();

        if !self.frame_started {
            panic!("begin_frame must be called before end_frame_and_draw can be called!");
        }

        let full_output = self.state.egui_ctx().end_pass();

        self.handle_full_output(
            full_output,
            device,
            queue,
            encoder,
            window,
            window_surface_view,
            screen_descriptor,
        );
    }

    fn handle_full_output(
        &mut self,
        full_output: FullOutput,
        device: &Device,
        queue: &Queue,
        encoder: &mut CommandEncoder,
        window: &Window,
        window_surface_view: &TextureView,
        screen_descriptor: ScreenDescriptor,
    ) {
        self.ppp(screen_descriptor.pixels_per_point);

        self.state
            .handle_platform_output(window, full_output.platform_output);

        let tris = self
            .state
            .egui_ctx()
            .tessellate(full_output.shapes, self.state.egui_ctx().pixels_per_point());
        for (id, image_delta) in &full_output.textures_delta.set {
            self.renderer
                .update_texture(device, queue, *id, image_delta);
        }
        self.renderer
            .update_buffers(device, queue, encoder, &tris, &screen_descriptor);
        let rpass = encoder.begin_render_pass(&wgpu::RenderPassDescriptor {
            color_attachments: &[Some(wgpu::RenderPassColorAttachment {
                view: window_surface_view,
                depth_slice: None,
                resolve_target: None,
                ops: egui_wgpu::wgpu::Operations {
                    load: egui_wgpu::wgpu::LoadOp::Load,
                    store: StoreOp::Store,
                },
            })],
            depth_stencil_attachment: None,
            timestamp_writes: None,
            label: Some("egui main render pass"),
            occlusion_query_set: None,
            multiview_mask: None,
        });

        self.renderer
            .render(&mut rpass.forget_lifetime(), &tris, &screen_descriptor);
        for x in &full_output.textures_delta.free {
            self.renderer.free_texture(x)
        }

        self.frame_started = false;
    }
}