// lib.rs
pub mod material;
pub mod render;
pub mod simulation;
// main.rs
use winit::{
    dpi::PhysicalSize,
    event::{Event, WindowEvent},
    event_loop::{ControlFlow, EventLoop},
    window::WindowBuilder,
};

mod render;
mod simulation;
mod material;

use render::Renderer;
use simulation::Simulation;
// Removed unused import:
// use material::Material;

fn main() {
    // Fixed window size (500x500)
    let window_size = PhysicalSize::new(500, 500);

    // Create the event loop and window
    let event_loop = EventLoop::new();

    let window = WindowBuilder::new()
        .with_title("Falling Sand Simulator")
        .with_inner_size(window_size)
        .with_resizable(false) // Fixed size
        .build(&event_loop)
        .unwrap();

    // Initialize the renderer and simulation
    let mut renderer = pollster::block_on(Renderer::new(&window));
    let mut simulation = Simulation::new();

    // Run the event loop
    event_loop.run(move |event, _, control_flow| {
        *control_flow = ControlFlow::Poll;

        match event {
            Event::RedrawRequested(_) => {
                // Render
                match renderer.render(&simulation) {
                    Ok(_) => {}
                    // Recreate the surface if lost
                    Err(wgpu::SurfaceError::Lost) => renderer.resize(renderer.size),
                    // The system is out of memory, exit
                    Err(wgpu::SurfaceError::OutOfMemory) => *control_flow = ControlFlow::Exit,
                    // All other errors (Outdated, Timeout) should be resolved by the next frame
                    Err(e) => eprintln!("{:?}", e),
                }
            }
            Event::MainEventsCleared => {
                // Request a redraw
                window.request_redraw();
            }
            Event::WindowEvent { ref event, .. } => {
                if !renderer.input(event, &mut simulation, &window) {
                    // Close window
                    match event {
                        WindowEvent::CloseRequested => *control_flow = ControlFlow::Exit,
                        _ => {}
                    }
                }
            }
            _ => {}
        }
    });
}
// material.rs
use bytemuck::{Pod, Zeroable};

#[repr(u32)]
#[derive(Copy, Clone, PartialEq, Eq, Debug)]
#[allow(dead_code)] // Suppress warnings for unused variants
pub enum Material {
    Empty = 0,
    Sand = 1,
    Water = 2,
    Earth = 3,
    Lava = 4,
    // Add more materials if needed
}

// SAFETY: Material is a field-less enum with a defined integer representation.
// All possible bit patterns from 0 to 4 are valid discriminants.
unsafe impl Zeroable for Material {}
unsafe impl Pod for Material {}

#[allow(dead_code)] // Suppress warning for unused method
impl Material {
    pub fn color(&self) -> [f32; 3] {
        match self {
            Material::Empty => [0.0, 0.0, 0.0],   // Black
            Material::Sand => [0.8, 0.7, 0.5],    // Sandy color
            Material::Water => [0.0, 0.5, 1.0],   // Blue
            Material::Earth => [0.4, 0.3, 0.2],   // Brown
            Material::Lava => [1.0, 0.3, 0.0],    // Orange-red
        }
    }
}
// render.rs
use winit::window::Window;

use crate::simulation::{Simulation, SIM_HEIGHT, SIM_WIDTH};
use crate::material::Material;

pub struct Renderer {
    device: wgpu::Device,
    queue: wgpu::Queue,
    surface: wgpu::Surface,
    config: wgpu::SurfaceConfiguration,
    pub size: winit::dpi::PhysicalSize<u32>,
    render_pipeline: wgpu::RenderPipeline,
    compute_pipeline: wgpu::ComputePipeline,
    cell_buffer: wgpu::Buffer,
    compute_bind_group: wgpu::BindGroup,
    render_bind_group: wgpu::BindGroup,
    cursor_position: Option<winit::dpi::PhysicalPosition<f64>>,
}

impl Renderer {
    pub async fn new(window: &Window) -> Self {
        let size = window.inner_size();

        // Instance
        let instance = wgpu::Instance::default();

        // Surface
        let surface = unsafe { instance.create_surface(window) }.unwrap();

        // Adapter
        let adapter = instance
            .request_adapter(&wgpu::RequestAdapterOptions {
                power_preference: wgpu::PowerPreference::HighPerformance,
                compatible_surface: Some(&surface),
                ..Default::default()
            })
            .await
            .unwrap();

        // Device and Queue
        let required_features = wgpu::Features::VERTEX_WRITABLE_STORAGE;

        // Get the features supported by the adapter
        let available_features = adapter.features();

        if !available_features.contains(required_features) {
            panic!("Required features are not supported by the adapter");
        }

        let (device, queue) = adapter.request_device(
            &wgpu::DeviceDescriptor {
                features: required_features & available_features, // Ensure we're only requesting supported features
                limits: wgpu::Limits::default(),
                label: None,
            },
            None,
        ).await.unwrap();

        // Surface configuration
        let config = wgpu::SurfaceConfiguration {
            usage: wgpu::TextureUsages::RENDER_ATTACHMENT,
            format: surface.get_capabilities(&adapter).formats[0],
            width: size.width,
            height: size.height,
            present_mode: wgpu::PresentMode::Fifo,
            alpha_mode: surface.get_capabilities(&adapter).alpha_modes[0],
            view_formats: vec![],
        };
        surface.configure(&device, &config);

        // Create a buffer for the simulation data
        let cell_buffer_size =
            (std::mem::size_of::<u32>() * SIM_WIDTH as usize * SIM_HEIGHT as usize) as u64;

        let cell_buffer = device.create_buffer(&wgpu::BufferDescriptor {
            label: Some("Cell Buffer"),
            size: cell_buffer_size,
            usage: wgpu::BufferUsages::STORAGE
                | wgpu::BufferUsages::COPY_DST
                | wgpu::BufferUsages::COPY_SRC,
            mapped_at_creation: false,
        });

        // Create separate bind group layouts
        // Compute Bind Group Layout (read-write access)
        let compute_bind_group_layout = device.create_bind_group_layout(&wgpu::BindGroupLayoutDescriptor {
            label: Some("Compute Bind Group Layout"),
            entries: &[
                wgpu::BindGroupLayoutEntry {
                    binding: 0,
                    visibility: wgpu::ShaderStages::COMPUTE,
                    ty: wgpu::BindingType::Buffer {
                        ty: wgpu::BufferBindingType::Storage { read_only: false },
                        has_dynamic_offset: false,
                        min_binding_size: None,
                    },
                    count: None,
                },
            ],
        });

        // Render Bind Group Layout (read-only access)
        let render_bind_group_layout = device.create_bind_group_layout(&wgpu::BindGroupLayoutDescriptor {
            label: Some("Render Bind Group Layout"),
            entries: &[
                wgpu::BindGroupLayoutEntry {
                    binding: 0,
                    visibility: wgpu::ShaderStages::VERTEX | wgpu::ShaderStages::FRAGMENT,
                    ty: wgpu::BindingType::Buffer {
                        ty: wgpu::BufferBindingType::Storage { read_only: true },
                        has_dynamic_offset: false,
                        min_binding_size: None,
                    },
                    count: None,
                },
            ],
        });

        // Create bind groups
        let compute_bind_group = device.create_bind_group(&wgpu::BindGroupDescriptor {
            label: Some("Compute Bind Group"),
            layout: &compute_bind_group_layout,
            entries: &[
                wgpu::BindGroupEntry {
                    binding: 0,
                    resource: cell_buffer.as_entire_binding(),
                },
            ],
        });

        let render_bind_group = device.create_bind_group(&wgpu::BindGroupDescriptor {
            label: Some("Render Bind Group"),
            layout: &render_bind_group_layout,
            entries: &[
                wgpu::BindGroupEntry {
                    binding: 0,
                    resource: cell_buffer.as_entire_binding(),
                },
            ],
        });

        // Load shaders
        let vertex_shader_source = include_str!("shaders/vertex.wgsl");
        let vertex_shader_source = vertex_shader_source
            .replace("$SIM_WIDTH", &SIM_WIDTH.to_string())
            .replace("$SIM_HEIGHT", &SIM_HEIGHT.to_string())
            .replace("$WINDOW_WIDTH", &size.width.to_string())
            .replace("$WINDOW_HEIGHT", &size.height.to_string());

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

        let fragment_shader_source = include_str!("shaders/fragment.wgsl");
        let fragment_shader_source = fragment_shader_source
        .replace("$SIM_WIDTH", &format!("{}u", SIM_WIDTH))
        .replace("$SIM_HEIGHT", &format!("{}u", SIM_HEIGHT))
        .replace("$WINDOW_WIDTH", &format!("{}u", size.width))
        .replace("$WINDOW_HEIGHT", &format!("{}u", size.height));

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

        let compute_shader_source = include_str!("shaders/compute.wgsl");
        let compute_shader_source = compute_shader_source
            .replace("$SIM_WIDTH", &format!("{}u", SIM_WIDTH))
            .replace("$SIM_HEIGHT", &format!("{}u", SIM_HEIGHT));

        let compute_shader = device.create_shader_module(wgpu::ShaderModuleDescriptor {
            label: Some("Compute Shader"),
            source: wgpu::ShaderSource::Wgsl(compute_shader_source.into()),
        });

        // Render pipeline layout
        let render_pipeline_layout = device.create_pipeline_layout(&wgpu::PipelineLayoutDescriptor {
            label: Some("Render Pipeline Layout"),
            bind_group_layouts: &[&render_bind_group_layout],
            push_constant_ranges: &[],
        });

        // Render pipeline
        let render_pipeline = device.create_render_pipeline(&wgpu::RenderPipelineDescriptor {
            label: Some("Render Pipeline"),
            layout: Some(&render_pipeline_layout),
            vertex: wgpu::VertexState {
                module: &vertex_shader,
                entry_point: "main",
                buffers: &[],
            },
            fragment: Some(wgpu::FragmentState {
                module: &fragment_shader,
                entry_point: "main",
                targets: &[Some(wgpu::ColorTargetState {
                    format: config.format,
                    blend: None,
                    write_mask: wgpu::ColorWrites::ALL,
                })],
            }),
            primitive: wgpu::PrimitiveState {
                topology: wgpu::PrimitiveTopology::TriangleList,
                ..Default::default()
            },
            depth_stencil: None,
            multisample: wgpu::MultisampleState::default(),
            multiview: None,
        });

        // Compute pipeline layout
        let compute_pipeline_layout =
            device.create_pipeline_layout(&wgpu::PipelineLayoutDescriptor {
                label: Some("Compute Pipeline Layout"),
                bind_group_layouts: &[&compute_bind_group_layout],
                push_constant_ranges: &[],
            });

        // Compute pipeline
        let compute_pipeline = device.create_compute_pipeline(&wgpu::ComputePipelineDescriptor {
            label: Some("Compute Pipeline"),
            layout: Some(&compute_pipeline_layout),
            module: &compute_shader,
            entry_point: "main",
        });

        Self {
            device,
            queue,
            surface,
            config,
            size,
            render_pipeline,
            compute_pipeline,
            cell_buffer,
            compute_bind_group,
            render_bind_group,
            cursor_position: None,
        }
    }

    pub fn input(
        &mut self,
        event: &winit::event::WindowEvent,
        simulation: &mut Simulation,
        _window: &Window,
    ) -> bool {
        match event {
            winit::event::WindowEvent::CursorMoved { position, .. } => {
                // Store the current cursor position
                self.cursor_position = Some(*position);
                false
            }
            winit::event::WindowEvent::MouseInput {
                state,
                button,
                ..
            } => {
                if *button == winit::event::MouseButton::Left
                    && *state == winit::event::ElementState::Pressed
                {
                    if let Some(pos) = self.cursor_position {
                        let x = (pos.x / self.size.width as f64 * SIM_WIDTH as f64) as u32;
                        let y = (pos.y / self.size.height as f64 * SIM_HEIGHT as f64) as u32;
                        simulation.add_material(x, y, Material::Sand);
                    }
                    true
                } else {
                    false
                }
            }
            _ => false,
        }
    }

    pub fn resize(&mut self, new_size: winit::dpi::PhysicalSize<u32>) {
        if new_size.width > 0 && new_size.height > 0 {
            self.size = new_size;
            self.config.width = new_size.width;
            self.config.height = new_size.height;
            self.surface.configure(&self.device, &self.config);
        }
    }

    pub fn render(&mut self, simulation: &Simulation) -> Result<(), wgpu::SurfaceError> {
        // Update the cell buffer with the simulation data
        self.queue.write_buffer(
            &self.cell_buffer,
            0,
            bytemuck::cast_slice(&simulation.cells),
        );

        // Get current frame
        let output = self.surface.get_current_texture()?;
        let view = output
            .texture
            .create_view(&wgpu::TextureViewDescriptor::default());

        // Create a single command encoder
        let mut encoder = self
            .device
            .create_command_encoder(&wgpu::CommandEncoderDescriptor {
                label: Some("Main Encoder"),
            });

        // Run compute shader to update simulation
        {
            let mut compute_pass = encoder.begin_compute_pass(&wgpu::ComputePassDescriptor {
                label: Some("Compute Pass"),
            });
            compute_pass.set_pipeline(&self.compute_pipeline);
            compute_pass.set_bind_group(0, &self.compute_bind_group, &[]);
            // Adjust workgroup counts
            let workgroup_size = 8u32;
            let dispatch_x = (SIM_WIDTH + workgroup_size - 1) / workgroup_size;
            let dispatch_y = (SIM_HEIGHT + workgroup_size - 1) / workgroup_size;
            compute_pass.dispatch_workgroups(dispatch_x, dispatch_y, 1);
        }

        // Render pass
        {
            let mut render_pass = encoder.begin_render_pass(&wgpu::RenderPassDescriptor {
                label: Some("Render Pass"),
                color_attachments: &[Some(wgpu::RenderPassColorAttachment {
                    view: &view,
                    resolve_target: None,
                    ops: wgpu::Operations {
                        load: wgpu::LoadOp::Load,
                        store: true,
                    },
                })],
                depth_stencil_attachment: None,
            });

            render_pass.set_pipeline(&self.render_pipeline);
            render_pass.set_bind_group(0, &self.render_bind_group, &[]);
            render_pass.draw(0..6, 0..1); // Drawing a quad
        }

        // Submit command buffer
        self.queue.submit(std::iter::once(encoder.finish()));
        output.present();

        Ok(())
    }
}
// simulation.rs
use crate::material::Material;

pub const SIM_WIDTH: u32 = 500;
pub const SIM_HEIGHT: u32 = 500;

pub struct Simulation {
    pub cells: Vec<Material>, // Simulation grid of materials
}

impl Simulation {
    pub fn new() -> Self {
        // Initialize the grid with empty materials
        let cells = vec![Material::Empty; (SIM_WIDTH * SIM_HEIGHT) as usize];

        Self { cells }
    }

    #[allow(dead_code)] // Suppress warning for unused method
    pub fn update(&mut self) {
        // Update the simulation state on CPU if needed
        // In this case, we are handling the simulation on GPU, so this may be minimal
    }

    // Methods to interact with the simulation (e.g., add material at position)
    pub fn add_material(&mut self, x: u32, y: u32, material: Material) {
        if x < SIM_WIDTH && y < SIM_HEIGHT {
            let index = (y * SIM_WIDTH + x) as usize;
            self.cells[index] = material;
        }
    }
}
 ➡ cd .\shaders\
 ➡ cat *
// compute.wgsl

const SIM_WIDTH: u32 = $SIM_WIDTH;
const SIM_HEIGHT: u32 = $SIM_HEIGHT;

@group(0) @binding(0)
var<storage, read_write> cells: array<u32>;

@compute @workgroup_size(8, 8, 1)
fn main(@builtin(global_invocation_id) global_id: vec3<u32>) {
    let x = global_id.x;
    let y = global_id.y;
    let index = y * SIM_WIDTH + x;

    // Check bounds
    if x >= SIM_WIDTH || y >= SIM_HEIGHT {
        return;
    }

    // Simple simulation logic: Sand falls down
    if cells[index] == 1u {
        var moved = false;

        // Try moving down
        if y + 1u < SIM_HEIGHT {
            let below_index = (y + 1u) * SIM_WIDTH + x;
            if cells[below_index] == 0u {
                cells[below_index] = 1u;
                cells[index] = 0u;
                moved = true;
            }
        }

        // Try moving down-left
        if !moved && x > 0u && y + 1u < SIM_HEIGHT {
            let below_left_index = (y + 1u) * SIM_WIDTH + (x - 1u);
            if cells[below_left_index] == 0u {
                cells[below_left_index] = 1u;
                cells[index] = 0u;
                moved = true;
            }
        }

        // Try moving down-right
        if !moved && x + 1u < SIM_WIDTH && y + 1u < SIM_HEIGHT {
            let below_right_index = (y + 1u) * SIM_WIDTH + (x + 1u);
            if cells[below_right_index] == 0u {
                cells[below_right_index] = 1u;
                cells[index] = 0u;
                moved = true;
            }
        }
    }

    // Add more simulation logic for other materials if needed
}
// fragment.wgsl

const SIM_WIDTH: u32 = $SIM_WIDTH;
const SIM_HEIGHT: u32 = $SIM_HEIGHT;
const WINDOW_WIDTH: u32 = $WINDOW_WIDTH;
const WINDOW_HEIGHT: u32 = $WINDOW_HEIGHT;

@group(0) @binding(0)
var<storage, read> cells: array<u32>;

@fragment
fn main(@builtin(position) frag_coord: vec4<f32>) -> @location(0) vec4<f32> {
    let x = u32(frag_coord.x / f32(WINDOW_WIDTH) * f32(SIM_WIDTH));
    let y = u32(frag_coord.y / f32(WINDOW_HEIGHT) * f32(SIM_HEIGHT));

    // Check bounds
    if x >= SIM_WIDTH || y >= SIM_HEIGHT {
        discard;
    }

    let index = y * SIM_WIDTH + x;
    let material = cells[index];

    var color: vec3<f32>;
    if material == 0u {
        color = vec3(0.0, 0.0, 0.0); // Empty
    } else if material == 1u {
        color = vec3(0.8, 0.7, 0.5); // Sand
    } else if material == 2u {
        color = vec3(0.0, 0.5, 1.0); // Water
    } else if material == 3u {
        color = vec3(0.4, 0.3, 0.2); // Earth
    } else if material == 4u {
        color = vec3(1.0, 0.3, 0.0); // Lava
    } else {
        color = vec3(1.0, 1.0, 1.0); // Default
    }

    return vec4(color, 1.0);
}
// vertex.wgsl

@vertex
fn main(@builtin(vertex_index) vertex_index: u32) -> @builtin(position) vec4<f32> {
    var positions = array<vec2<f32>, 6>(
        vec2(-1.0, -1.0),
        vec2(1.0, -1.0),
        vec2(-1.0, 1.0),
        vec2(-1.0, 1.0),
        vec2(1.0, -1.0),
        vec2(1.0, 1.0),
    );

    let pos = positions[vertex_index];
    return vec4(pos, 0.0, 1.0);
}

Can you make it so that the pixels are actually 4 pixels, like make the resolution smaller.. secondly, make it so that the element selection are buttons on the screen, thirdly, make the elements actually are affected by gravity and does like the cellular automata where it acts like sand (sand & dirt) and liquid (water & lava) 
Output every full scripts, at its fullest