// 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;
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)]
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 {}

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>, // Made this field public
    render_pipeline: wgpu::RenderPipeline,
    compute_pipeline: wgpu::ComputePipeline,
    cell_buffer: wgpu::Buffer,
    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 (device, queue) = adapter
            .request_device(
                &wgpu::DeviceDescriptor {
                    features: wgpu::Features::default(), // No need to enable additional 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,
        });

        // Bind group layout and bind group
        let bind_group_layout =
            device.create_bind_group_layout(&wgpu::BindGroupLayoutDescriptor {
                label: Some("Cell Bind Group Layout"),
                entries: &[
                    // Binding for compute shader (read-write)
                    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,
                    },
                    // Binding for vertex and fragment shaders (read-only)
                    wgpu::BindGroupLayoutEntry {
                        binding: 1,
                        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,
                    },
                ],
            });

        let bind_group = device.create_bind_group(&wgpu::BindGroupDescriptor {
            label: Some("Cell Bind Group"),
            layout: &bind_group_layout,
            entries: &[
                // Binding for compute shader
                wgpu::BindGroupEntry {
                    binding: 0,
                    resource: cell_buffer.as_entire_binding(),
                },
                // Binding for vertex and fragment shaders
                wgpu::BindGroupEntry {
                    binding: 1,
                    resource: cell_buffer.as_entire_binding(),
                },
            ],
        });

        // Load shaders
        let vertex_shader_source = include_str!("shaders/vertex.wgsl");
        // Replace placeholders with actual values (if any)
        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");
        // Replace placeholders with actual values
        let fragment_shader_source = fragment_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 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");
        // Replace placeholders with actual values
        let compute_shader_source = compute_shader_source
            .replace("$SIM_WIDTH", &SIM_WIDTH.to_string())
            .replace("$SIM_HEIGHT", &SIM_HEIGHT.to_string());

        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 pipeline_layout = device.create_pipeline_layout(&wgpu::PipelineLayoutDescriptor {
            label: Some("Render Pipeline Layout"),
            bind_group_layouts: &[&bind_group_layout],
            push_constant_ranges: &[],
        });

        // Render pipeline
        let render_pipeline = device.create_render_pipeline(&wgpu::RenderPipelineDescriptor {
            label: Some("Render Pipeline"),
            layout: Some(&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: &[&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,
            bind_group,
            cursor_position: None, // Initialize cursor_position
        }
    }

    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),
        );

        // Run compute shader to update simulation
        {
            let mut encoder = self
                .device
                .create_command_encoder(&wgpu::CommandEncoderDescriptor {
                    label: Some("Compute Encoder"),
                });

            {
                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.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);
            }

            self.queue.submit(Some(encoder.finish()));
        }

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

        // Command encoder
        let mut encoder = self
            .device
            .create_command_encoder(&wgpu::CommandEncoderDescriptor {
                label: Some("Render Encoder"),
            });

        // 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.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 }
    }

    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
    }

    // Add 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;
        }
    }
}
// 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: f32 = $WINDOW_WIDTH;
const WINDOW_HEIGHT: f32 = $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 / WINDOW_WIDTH * f32(SIM_WIDTH));
    let y = u32(frag_coord.y / 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);
}

Fix errors and warnings
warning: unused import: `material::Material`
  --> src/main.rs:15:5
   |
15 | use material::Material;
   |     ^^^^^^^^^^^^^^^^^^
   |
   = note: `#[warn(unused_imports)]` on by default

warning: method `update` is never used
  --> src\simulation.rs:19:12
   |
11 | impl Simulation {
   | --------------- method in this implementation
...
19 |     pub fn update(&mut self) {
   |            ^^^^^^
   |
   = note: `#[warn(dead_code)]` on by default

warning: variants `Water`, `Earth`, and `Lava` are never constructed
  --> src\material.rs:9:5
   |
6  | pub enum Material {
   |          -------- variants in this enum
...
9  |     Water = 2,
   |     ^^^^^
10 |     Earth = 3,
   |     ^^^^^
11 |     Lava = 4,
   |     ^^^^
   |
   = note: `Material` has derived impls for the traits `Debug` and `Clone`, but these are intentionally ignored during dead code analysis

warning: method `color` is never used
  --> src\material.rs:21:12
   |
20 | impl Material {
   | ------------- method in this implementation
21 |     pub fn color(&self) -> [f32; 3] {
   |            ^^^^^

warning: `falling_sand_simulator` (bin "falling_sand_simulator") generated 4 warnings (run `cargo fix --bin "falling_sand_simulator"` to apply 1 suggestion)
    Finished `dev` profile [unoptimized + debuginfo] target(s) in 0.22s
     Running `C:\Users\Enzo\Documents\falling_sand_simulator\target\debug\falling_sand_simulator.exe`
thread 'main' panicked at C:\Users\Enzo\.cargo\registry\src\index.crates.io-6f17d22bba15001f\wgpu-0.16.3\src\backend\direct.rs:3019:5:
wgpu error: Validation Error

Caused by:
    In Device::create_shader_module
      note: label = `Fragment Shader`

Shader 'Fragment Shader' parsing error: the type of `SIM_WIDTH` is expected to be `u32`, but got `i32`
  ┌─ wgsl:3:7
  │
3 │ const SIM_WIDTH: u32 = 500;
  │       ^^^^^^^^^ definition of `SIM_WIDTH`


    the type of `SIM_WIDTH` is expected to be `u32`, but got `i32`


note: run with `RUST_BACKTRACE=1` environment variable to display a backtrace
error: process didn't exit successfully: `C:\Users\Enzo\Documents\falling_sand_simulator\target\debug\falling_sand_simulator.exe` (exit code: 101)