mewgpu 3.7.3

use std::{
    cell::OnceCell,
    time::{ Duration, Instant, },
};
use mewgpu::{
    prelude::*,
    winit::{
        application::ApplicationHandler,
        dpi::PhysicalSize,
        event_loop::{ EventLoop, ActiveEventLoop, ControlFlow, },
        event::WindowEvent,
        window::{ WindowAttributes, WindowId, },
    },
};


const FRAME_TIME: Duration = Duration::new(0, 1_000_000_000 / 60);


mew! {
    // Setting up vertex type for defining a model.
    // Each point needs a physical position & texture UV.
    vertex_struct ModelVertex [
        0 pos => Float32x3,
        1 uv => Float32x2,
    ];
    buffer ModelVertexBuffer <ModelVertex> as VERTEX | COPY_DST;


    // Setting up instance data if we want to display multiple objects.
    // Each instance will have its own XYZ position, & rotation along 1 axis to keep
    // it simple.
    // Shader @location needs to be offset by 2 to compensate for vertex data.
    vertex_struct InstanceData step Instance + 2 [
        // @location(2)
        0 pos => Float32x3,
        // @location(3)
        1 rot => Float32,
    ];
    buffer InstanceDataBuffer <InstanceData> as VERTEX | COPY_DST;


    // Setting up camera uniform matrix. Usually required for all sorts of graphical
    // apps, transforms all vertices at least to compensate for window scaling.
    // IMPORTANT: If you're targetting web, WebGL only supports UNIFORM buffers in
    // bind groups (vertex buffers are special). These need to have alignment of a
    // multiple of 16 bytes. If you don't have enough data to fill in, pad it, or
    // write bits directly to matrices.
    shader_struct CameraData [
        0 matrix => Mat4x4f,
    ];
    buffer CameraDataBuffer <CameraData> as UNIFORM | COPY_DST;


    // Setting up texture & sampler types. These are pretty generic.
    texture Texture {
        usage: COPY_DST | TEXTURE_BINDING,
        dimension: D2,
        sample_type: FLOAT,
    };
    sampler Sampler as Filtering;


    // Setting up a bind group to store the texture, sampler, & camera buffer.
    bind_group BindGroup [
        0 texture @ FRAGMENT => Texture,
        1 sampler @ FRAGMENT => Sampler,
        2 camera @ VERTEX => CameraDataBuffer,
    ];


    // Setting up depth texture type for culling.
    texture DepthTexture {
        usage: COPY_DST | RENDER_ATTACHMENT,
        dimension: D2,
        sample_type: DEPTH,
    };


    // Setting up the pipeline for rendering. Specify the bind group used, vertex
    // type, generic fragment target, depth format and cull face.
    // Immediates (aka "push constants") aren't well supported and to be honest not
    // all that useful, not covered in this example but should be simple enough to
    // figure out.
    pipeline Pipeline {
        bind_groups: [
            0 => BindGroup,
        ],
        fragment_targets: [
            0 => ALPHA_BLENDING ALL as ANY,
        ],
        vertex_types: [
            0 => ModelVertex,
            1 => InstanceData,
        ],
        depth: Depth32Float,
        cull: Front,
    };
}


// All logic lives in the App, just set up window event loop.
fn main() {
    let event_loop = EventLoop::new().unwrap();
    let mut app = App::new(&event_loop);
    event_loop.run_app(&mut app).unwrap();
}


#[cfg(target_arch = "wasm32")]
fn get_webgl_canvas() -> web_sys::HtmlCanvasElement {
    use wasm_bindgen::prelude::*;
    use winit::platform::web::WindowAttributesExtWebSys;
    web_sys::window().expect("Couldn't get DOM window")
        .document().expect("Couldn't get DOM document")
        .get_element_by_id("canvas").expect("Couldn't get canvas element")
        .dyn_into::<web_sys::HtmlCanvasElement>().expect("Couldn't get GL canvas object")
}


// Buffers, bind groups & pipelines rely on the context & need to be created
// in the event loop. Use a separate struct to manage all that, or if you have
// many objects with their own graphics data, they'll need to keep their own
// structs and construct them with the RenderContext.
struct GraphicsData {
    vertices: ModelVertexBuffer,
    instances: InstanceDataBuffer,
    bindgroup: BindGroup,
    depth: OnceCell<DepthTexture>,
    pipeline: Pipeline,
}

impl GraphicsData {
    const CIRNOS: usize = 9;

    fn new(context: &RenderContext, window: &ActiveSurfaceWindow) -> Self {
        // Create vertex buffer & intialize its data. For simple models or voxels you
        // could probably keep whole models in the shader, but for the sake of example
        // generate a model and write it.
        let vertices: ModelVertexBuffer = context.new_buffer(12);
        let vertex_data: [ModelVertex; 12] = std::array::from_fn(|i| {
            // Not relevant to mew, create rotating triangles to form a pyramid.
            let face = i as isize / 3;
            let face_x = (face / 2).rem_euclid(2) * 2 - 1;
            let face_z = ((face + 1) / 2).rem_euclid(2) * 2 - 1;
            let pos = match i % 3 {
                0 => [0, 1, 0],
                1 => [face_x, -1, face_z],
                2 => [-face_z, -1, face_x],
                _ => unreachable!(),
            };

            // Put above triangles in a ModelVertex, transforming UV to be in the range
            // [0, 1] from [-1, 1]. Vertex & shader structs can be into()'d from tuples.
            (
                [pos[0] as f32, pos[1] as f32, pos[2] as f32],
                [(pos[0] as f32) * 0.5 + 0.5, (pos[2] as f32) * 0.5 + 0.5],
            ).into()
        });
        // Write above ModelVertex array to the actual GPU buffer. Struct slices can be
        // cast into &[u8] with as_byte_slice().
        context.queue.write_buffer(&vertices, 0, vertex_data.as_byte_slice());

        // Create an instance buffer to draw multiple models at once. Don't worry about
        // populating with data, this will be done at run time.
        let instances: InstanceDataBuffer = context.new_buffer(Self::CIRNOS as u64);

        // Load image data first, so we can make a texture of the correct size.
        let cirnium = image::load_from_memory(include_bytes!("cirnium.webp")).unwrap().into_rgba8();
        let texture: Texture = context.new_texture((cirnium.dimensions().0, cirnium.dimensions().1, 1), wgpu::TextureFormat::Rgba8UnormSrgb);
        // Write above image into the GPU texture.
        let copy = texture.as_image_copy();
        context.queue.write_texture(copy, &cirnium, texture.texel_layout(), texture.size());

        let sampler: Sampler = context.new_sampler();
        let camera: CameraDataBuffer = context.new_buffer(1);

        // Make a bind group with all its components - these are still accessible from
        // the bind group (e.g. bindgroup.camera).
        let bindgroup: BindGroup = context.new_bind_group((texture, sampler, camera));

        // Create a shader & pipeline. Pipeline's format must match the output texture,
        // in this case the window's surface which is passed in through this function.
        let shader = context.device.create_shader_module(wgpu::include_wgsl!("shader.wgsl"));
        let pipeline: Pipeline = context.new_pipeline(window.config().format, &shader, None, None);

        // The depth buffer will need to frequently get dropped & recreated.
        let depth = OnceCell::new();

        Self {
            vertices,
            instances,
            bindgroup,
            depth,
            pipeline,
        }
    }

    // Create a depth buffer, the same size as the window and the same format as
    // defined in the pipeline. Cached but needs to be dropped & recreated when the
    // window is resized.
    fn get_depth_buffer(&self, context: &RenderContext, size: PhysicalSize<u32>) -> &DepthTexture {
        self.depth.get_or_init(|| context.new_texture((size.width, size.height, 1), wgpu::TextureFormat::Depth32Float))
    }

    fn set_camera_matrix(&self, context: &RenderContext, size: PhysicalSize<u32>) {
        let project = cgmath::perspective(cgmath::Rad(1.3), size.width as f32 / size.height as f32, 0.001, 100.0);
        let translate = cgmath::Matrix4::from_translation([0.0, -0.6, -6.0].into());
        let camera = project * translate;

        let floats: &[[f32; 4]; 4] = camera.as_ref();
        let camera_data: CameraData = (*floats,).into();
        context.queue.write_buffer(&*self.bindgroup.camera, 0, camera_data.as_ref());
    }
}


struct App {
    context: DeferredContext,
    window: OnceCell<SurfaceWindow>,
    graphics_data: Option<GraphicsData>,
    suspended: bool,
    last_frame: Instant,
    time: Duration,
}

impl App {
    fn new(_event_loop: &EventLoop<()>) -> Self {
        // Can't build the context yet, on web need to create a Window/Surface first,
        // and to create a Window we need to run the app.
        let context: RenderContextBuilder = RenderContextBuilder::new()
            .with_backends(wgpu::Backends::all());

        Self {
            context: context.deferred(),
            window: OnceCell::new(),
            graphics_data: None,
            suspended: false,
            last_frame: Instant::now(),
            time: Duration::new(0, 0),
        }
    }
}

impl ApplicationHandler for App {
    fn resumed(&mut self, event_loop: &ActiveEventLoop) {
        let _window = self.window.get_or_init(|| {
            #[allow(unused_mut)]
            let mut attributes = WindowAttributes::default();

            #[cfg(target_arch = "wasm32")]
            attributes.with_canvas(Some(get_webgl_canvas("canvas")));

            // Passing in the event loop here gives the render context access to a
            // display handle, which needs to be manually provided since wgpu 29.0.
            // On web, will also provide the surface to the context builder. The
            // surface isn't configured until it's needed.
            self.context.create_window(event_loop, attributes, Default::default())
                .expect("Couldn't create window")
        });

        // Can't do much else until we have a valid context, which may take until
        // the next function call to finish building. Move on to main loop.
        self.suspended = false;
    }

    fn window_event(&mut self, event_loop: &ActiveEventLoop, _window_id: WindowId, event: WindowEvent) {
        if self.suspended { return; }

        let context = match self.context.get_context() {
            Ok(c) => c,
            // Something went wrong building the context, abort.
            Err(GetDeferredContextError::Build(e)) => panic!("Error building context: {e}"),
            // resume() didn't run for some reason, RenderContext has no display
            // handle to use to initialize with.
            Err(GetDeferredContextError::RequiresDisplayHandle) =>
                unreachable!("Display handle should have been provided in resume()"),
            // resume() didn't run for some reason, RenderContext has no surface
            // to use to initialize with (on web).
            Err(GetDeferredContextError::RequiresSurface) =>
                unreachable!("Surface should have been provided in resume()"),
            // Still building (on web), will try again next window_event().
            Err(GetDeferredContextError::StillBuilding) => return,
        };

        match event {
            WindowEvent::RedrawRequested => {
                if !self.suspended {
                    // Framerate control
                    let now = Instant::now();
                    let elapsed = now.duration_since(self.last_frame);
                    self.last_frame = now;
                    self.time += elapsed;
                    event_loop.set_control_flow(ControlFlow::WaitUntil(now + FRAME_TIME));

                    // Now that the context is ready, we can finish configuring the window surface.
                    // init_surface() will create everything when needed and make sure it's configured.
                    let active_window = self.window.get().expect("Window created in resume()")
                        .init_with_context(&context).expect("Couldn't init window");

                    // Initialize required buffers and whatnot. Needs to know the window's size to create a
                    // depth buffer & calculate camera matrix, and the surface's format for the pipeline, so
                    // pass in the active window which definitely has a configured surface that we can query.
                    let graphics_data = self.graphics_data.get_or_insert_with(|| GraphicsData::new(&context, &active_window));

                    // We also need to update depth buffer, but the resize event is actually out
                    // of sync with the actual resizes - e.g. it could come a frame late.
                    // ActiveSurfaceWindow has a flag whether it has just been resized, use that
                    // instead to keep in sync with the surface config.
                    if active_window.resized() {
                        let _ = graphics_data.depth.take();
                        graphics_data.set_camera_matrix(&context, active_window.get_config_size());
                    }
                    let depth = graphics_data.get_depth_buffer(&context, active_window.get_config_size());

                    // Populate instance data with pseudo-randomized bouncing positions. Real apps would
                    // have some sort of state preserved between frames, but for the sake of the example, eh.
                    let time = self.time.as_secs_f32();
                    let instances: [InstanceData; GraphicsData::CIRNOS] = std::array::from_fn(|i| {
                        fn random_val(i: usize, j: usize) -> f32 {
                            ((i as f32 + j as f32 - 12.34) / (4 + i + j).pow(2) as f32 + (i as f32).sin() + (j as f32).cos()).rem_euclid(2.2) - 1.1
                        }

                        fn sawtooth(time: f32, mult: f32) -> f32 {
                            (1.0 - (time * mult).rem_euclid(2.0)).abs() * 2.0 - 1.0
                        }

                        (
                            [
                                sawtooth(time, random_val(i, 0)) * 3.0,
                                (time * random_val(i, 1) * 6.0).sin() * random_val(i, 2) * 0.3 + random_val(i, 3) * 2.0,
                                sawtooth(time, random_val(i, 4)) * 3.0
                            ],
                            [random_val(i, 4) * time],
                        ).into()
                    });

                    // Slices of mew shader/vertex structs can be cast to byte slices directly.
                    // Internally, this is just a transmute.
                    context.queue.write_buffer(&*graphics_data.instances, 0, instances.as_byte_slice());

                    // Actually rendering is almost bare wgpu, with some quality-of-life sprinkled
                    // here & there. You're probably better off coming up with your own abstraction to
                    // manage this, or just rawdog it.
                    let mut encoder = context.device.create_command_encoder(&Default::default());

                    // Scoping, render pass needs to be dropped before encoder finishes
                    {
                        let mut render_pass = encoder.begin_render_pass(&wgpu::RenderPassDescriptor {
                            // The output target of the render pass. Set this to the window surface,
                            // unless you're downscaling the output for performance, in which case
                            // create a new texture and then render that to the window.
                            color_attachments: &[Some(active_window.as_colour_attachment(Some(wgpu::Color::BLUE)))],
                            // Depth buffer which figures out what to cull. Needs to be the same size
                            // as the window surface.
                            depth_stencil_attachment: Some(depth.as_depth_attachment()),
                            // Don't really care 'bout the rest.
                            ..Default::default()
                        });

                        // mew structs are deref-able into wgpu pipelines, with a couple &s and *s.
                        render_pass.set_pipeline(&*graphics_data.pipeline);

                        render_pass.set_vertex_buffer(0, graphics_data.vertices.slice(..));
                        render_pass.set_vertex_buffer(1, graphics_data.instances.slice(..));
                        render_pass.set_bind_group(0, &*graphics_data.bindgroup, &[]);

                        render_pass.draw(0..12, 0..GraphicsData::CIRNOS as u32);
                    }

                    context.queue.submit([encoder.finish()]);
                    active_window.present_texture();
                }
            },

            WindowEvent::Resized(_new_size) => {
                // Resizing is "queued" and will come into effect when applicable.
                // It also queries the window's size from, well, the Window, so not really any
                // need to pass it in & store it - just trigger the event.
                // Managing the depth buffer is kept in sync and regenerated in the frame.
               self.window.get().expect("Window should have been created")
                    .resize(None);
            },

            _ => {},
        }
    }

    fn about_to_wait(&mut self, _event_loop: &ActiveEventLoop) {
        // Need to manually call request_redraw() if it's time to render a frame.
        if self.last_frame.elapsed() >= FRAME_TIME {
            self.window.get().expect("Window should have been created")
                .window().request_redraw();
        }
    }

    fn suspended(&mut self, _event_loop: &ActiveEventLoop) {
        // Android requires surfaces be dropped on suspend. init_surface() above
        // will re-create and configure the window surface automatically when unsuspended.
        self.window.get_mut().map(|w| w.drop_surface());
        self.suspended = true;
    }
}