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

        // Bind group layout with a single binding
        let bind_group_layout =
            device.create_bind_group_layout(&wgpu::BindGroupLayoutDescriptor {
                label: Some("Cell Bind Group Layout"),
                entries: &[
                    wgpu::BindGroupLayoutEntry {
                        binding: 0,
                        visibility: wgpu::ShaderStages::VERTEX
                            | wgpu::ShaderStages::FRAGMENT
                            | wgpu::ShaderStages::COMPUTE,
                        ty: wgpu::BindingType::Buffer {
                            ty: wgpu::BufferBindingType::Storage { read_only: false },
                            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: &[
                wgpu::BindGroupEntry {
                    binding: 0,
                    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 and type suffixes
        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!("{}.", size.width))
            .replace("$WINDOW_HEIGHT", &format!("{}.", 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");
        // Replace placeholders with actual values and type suffixes
        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 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,
        }
    }

    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.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.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(())
    }
}
// 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) // Changed to 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);
}

 ➡ cargo run
   Compiling falling_sand_simulator v0.1.0 (C:\Users\Enzo\Documents\falling_sand_simulator)
    Finished `dev` profile [unoptimized + debuginfo] target(s) in 8.40s
     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_render_pipeline
      note: label = `Render Pipeline`
    Error matching ShaderStages(FRAGMENT) shader requirements against the pipeline
    Shader global ResourceBinding { group: 0, binding: 0 } is not available in the layout pipeline layout
    Storage class Storage { access: LOAD | STORE } doesn't match the shader Storage { access: LOAD }


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)