use crate::ecs::sdf::brick_map::{BRICK_CORNER_COUNT, BRICK_SIZE, BRICKS_PER_ATLAS_DIM};
use crate::ecs::sdf::primitives::{CsgOperation, SdfEdit, SdfPrimitive};
use crate::ecs::sdf::resources::TerrainConfig;

const MAX_EDITS: usize = 1024;
const MAX_DIRTY_BRICKS_PER_FRAME: usize = 4096;

#[repr(C)]
#[derive(Copy, Clone, Debug, bytemuck::Pod, bytemuck::Zeroable)]
struct ComputeUniforms {
    brick_count: u32,
    terrain_enabled: u32,
    terrain_base_height: f32,
    terrain_material_id: u32,
    edit_count: u32,
    smoothness_scale: f32,
    terrain_seed: u32,
    terrain_octaves: u32,
    terrain_frequency: f32,
    terrain_amplitude: f32,
    terrain_gain: f32,
    _pad0: u32,
    _pad1: u32,
    _pad2: u32,
    _pad3: u32,
    _pad4: u32,
}

#[repr(C)]
#[derive(Copy, Clone, Debug, bytemuck::Pod, bytemuck::Zeroable)]
struct GpuSdfEdit {
    inverse_transform: [[f32; 4]; 4],
    primitive_type: u32,
    operation_type: u32,
    material_id: u32,
    smoothness: f32,
    uniform_scale: f32,
    param0: f32,
    param1: f32,
    param2: f32,
    param3: f32,
    _pad0: u32,
    _pad1: u32,
    _pad2: u32,
}

#[repr(C)]
#[derive(Copy, Clone, Debug, bytemuck::Pod, bytemuck::Zeroable)]
pub struct GpuDirtyBrick {
    pub world_origin: [f32; 3],
    pub voxel_size: f32,
    pub atlas_slot: u32,
    pub material_atlas_x: u32,
    pub material_atlas_y: u32,
    pub material_atlas_z: u32,
}

pub struct SdfComputePipelines {
    corners_pipeline: wgpu::ComputePipeline,
    materials_pipeline: wgpu::ComputePipeline,
    bind_group_layout: wgpu::BindGroupLayout,
    uniforms_buffer: wgpu::Buffer,
    edits_buffer: wgpu::Buffer,
    dirty_bricks_buffer: wgpu::Buffer,
    surface_flags_buffer: wgpu::Buffer,
    bind_group: Option<wgpu::BindGroup>,
    last_brick_count: u32,
}

impl SdfComputePipelines {
    pub fn new(device: &wgpu::Device) -> Self {
        let shader = device.create_shader_module(wgpu::ShaderModuleDescriptor {
            label: Some("SDF Evaluate Compute Shader"),
            source: wgpu::ShaderSource::Wgsl(
                include_str!("../../shaders/sdf_evaluate.wgsl").into(),
            ),
        });

        let bind_group_layout = device.create_bind_group_layout(&wgpu::BindGroupLayoutDescriptor {
            label: Some("SDF Compute Bind Group Layout"),
            entries: &[
                wgpu::BindGroupLayoutEntry {
                    binding: 0,
                    visibility: wgpu::ShaderStages::COMPUTE,
                    ty: wgpu::BindingType::Buffer {
                        ty: wgpu::BufferBindingType::Uniform,
                        has_dynamic_offset: false,
                        min_binding_size: None,
                    },
                    count: None,
                },
                wgpu::BindGroupLayoutEntry {
                    binding: 1,
                    visibility: wgpu::ShaderStages::COMPUTE,
                    ty: wgpu::BindingType::Buffer {
                        ty: wgpu::BufferBindingType::Storage { read_only: true },
                        has_dynamic_offset: false,
                        min_binding_size: None,
                    },
                    count: None,
                },
                wgpu::BindGroupLayoutEntry {
                    binding: 2,
                    visibility: wgpu::ShaderStages::COMPUTE,
                    ty: wgpu::BindingType::Buffer {
                        ty: wgpu::BufferBindingType::Storage { read_only: true },
                        has_dynamic_offset: false,
                        min_binding_size: None,
                    },
                    count: None,
                },
                wgpu::BindGroupLayoutEntry {
                    binding: 3,
                    visibility: wgpu::ShaderStages::COMPUTE,
                    ty: wgpu::BindingType::StorageTexture {
                        access: wgpu::StorageTextureAccess::WriteOnly,
                        format: wgpu::TextureFormat::R32Float,
                        view_dimension: wgpu::TextureViewDimension::D3,
                    },
                    count: None,
                },
                wgpu::BindGroupLayoutEntry {
                    binding: 4,
                    visibility: wgpu::ShaderStages::COMPUTE,
                    ty: wgpu::BindingType::StorageTexture {
                        access: wgpu::StorageTextureAccess::WriteOnly,
                        format: wgpu::TextureFormat::R32Uint,
                        view_dimension: wgpu::TextureViewDimension::D3,
                    },
                    count: None,
                },
                wgpu::BindGroupLayoutEntry {
                    binding: 5,
                    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,
                },
            ],
        });

        let pipeline_layout = device.create_pipeline_layout(&wgpu::PipelineLayoutDescriptor {
            label: Some("SDF Compute Pipeline Layout"),
            bind_group_layouts: &[&bind_group_layout],
            push_constant_ranges: &[],
        });

        let corners_pipeline = device.create_compute_pipeline(&wgpu::ComputePipelineDescriptor {
            label: Some("SDF Evaluate Corners Pipeline"),
            layout: Some(&pipeline_layout),
            module: &shader,
            entry_point: Some("evaluate_corners"),
            compilation_options: Default::default(),
            cache: None,
        });

        let materials_pipeline = device.create_compute_pipeline(&wgpu::ComputePipelineDescriptor {
            label: Some("SDF Evaluate Materials Pipeline"),
            layout: Some(&pipeline_layout),
            module: &shader,
            entry_point: Some("evaluate_materials"),
            compilation_options: Default::default(),
            cache: None,
        });

        let uniforms_buffer = device.create_buffer(&wgpu::BufferDescriptor {
            label: Some("SDF Compute Uniforms"),
            size: std::mem::size_of::<ComputeUniforms>() as u64,
            usage: wgpu::BufferUsages::UNIFORM | wgpu::BufferUsages::COPY_DST,
            mapped_at_creation: false,
        });

        let edits_buffer = device.create_buffer(&wgpu::BufferDescriptor {
            label: Some("SDF Compute Edits"),
            size: (std::mem::size_of::<GpuSdfEdit>() * MAX_EDITS) as u64,
            usage: wgpu::BufferUsages::STORAGE | wgpu::BufferUsages::COPY_DST,
            mapped_at_creation: false,
        });

        let dirty_bricks_buffer = device.create_buffer(&wgpu::BufferDescriptor {
            label: Some("SDF Compute Dirty Bricks"),
            size: (std::mem::size_of::<GpuDirtyBrick>() * MAX_DIRTY_BRICKS_PER_FRAME) as u64,
            usage: wgpu::BufferUsages::STORAGE | wgpu::BufferUsages::COPY_DST,
            mapped_at_creation: false,
        });

        let surface_flags_buffer = device.create_buffer(&wgpu::BufferDescriptor {
            label: Some("SDF Compute Surface Flags"),
            size: (std::mem::size_of::<u32>() * MAX_DIRTY_BRICKS_PER_FRAME) as u64,
            usage: wgpu::BufferUsages::STORAGE | wgpu::BufferUsages::COPY_DST,
            mapped_at_creation: false,
        });

        Self {
            corners_pipeline,
            materials_pipeline,
            bind_group_layout,
            uniforms_buffer,
            edits_buffer,
            dirty_bricks_buffer,
            surface_flags_buffer,
            bind_group: None,
            last_brick_count: 0,
        }
    }

    pub fn prepare(
        &mut self,
        queue: &wgpu::Queue,
        edits: &[SdfEdit],
        dirty_bricks: &[GpuDirtyBrick],
        terrain: &TerrainConfig,
        smoothness_scale: f32,
    ) {
        let brick_count = dirty_bricks.len().min(MAX_DIRTY_BRICKS_PER_FRAME);
        self.last_brick_count = brick_count as u32;

        if brick_count == 0 {
            return;
        }

        let edit_count = edits.len().min(MAX_EDITS);

        let uniforms = ComputeUniforms {
            brick_count: brick_count as u32,
            terrain_enabled: if terrain.enabled { 1 } else { 0 },
            terrain_base_height: terrain.base_height,
            terrain_material_id: terrain.material_id,
            edit_count: edit_count as u32,
            smoothness_scale,
            terrain_seed: terrain.seed,
            terrain_octaves: terrain.octaves,
            terrain_frequency: terrain.frequency,
            terrain_amplitude: terrain.amplitude,
            terrain_gain: terrain.gain,
            _pad0: 0,
            _pad1: 0,
            _pad2: 0,
            _pad3: 0,
            _pad4: 0,
        };

        queue.write_buffer(&self.uniforms_buffer, 0, bytemuck::cast_slice(&[uniforms]));

        let gpu_edits: Vec<GpuSdfEdit> =
            edits.iter().take(edit_count).map(serialize_edit).collect();

        if !gpu_edits.is_empty() {
            queue.write_buffer(&self.edits_buffer, 0, bytemuck::cast_slice(&gpu_edits));
        }

        queue.write_buffer(
            &self.dirty_bricks_buffer,
            0,
            bytemuck::cast_slice(&dirty_bricks[..brick_count]),
        );

        let zero_flags = vec![0u32; brick_count];
        queue.write_buffer(
            &self.surface_flags_buffer,
            0,
            bytemuck::cast_slice(&zero_flags),
        );
    }

    pub fn create_bind_group(
        &mut self,
        device: &wgpu::Device,
        brick_atlas_view: &wgpu::TextureView,
        material_atlas_view: &wgpu::TextureView,
    ) {
        self.bind_group = Some(device.create_bind_group(&wgpu::BindGroupDescriptor {
            label: Some("SDF Compute Bind Group"),
            layout: &self.bind_group_layout,
            entries: &[
                wgpu::BindGroupEntry {
                    binding: 0,
                    resource: self.uniforms_buffer.as_entire_binding(),
                },
                wgpu::BindGroupEntry {
                    binding: 1,
                    resource: self.edits_buffer.as_entire_binding(),
                },
                wgpu::BindGroupEntry {
                    binding: 2,
                    resource: self.dirty_bricks_buffer.as_entire_binding(),
                },
                wgpu::BindGroupEntry {
                    binding: 3,
                    resource: wgpu::BindingResource::TextureView(brick_atlas_view),
                },
                wgpu::BindGroupEntry {
                    binding: 4,
                    resource: wgpu::BindingResource::TextureView(material_atlas_view),
                },
                wgpu::BindGroupEntry {
                    binding: 5,
                    resource: self.surface_flags_buffer.as_entire_binding(),
                },
            ],
        }));
    }

    pub fn dispatch(&self, encoder: &mut wgpu::CommandEncoder) {
        if self.last_brick_count == 0 {
            return;
        }

        let bind_group = match &self.bind_group {
            Some(bind_group) => bind_group,
            None => return,
        };

        let total_corners = self.last_brick_count
            * (BRICK_CORNER_COUNT * BRICK_CORNER_COUNT * BRICK_CORNER_COUNT) as u32;
        let corner_workgroups = total_corners.div_ceil(128);

        {
            let mut compute_pass = encoder.begin_compute_pass(&wgpu::ComputePassDescriptor {
                label: Some("SDF Evaluate Corners"),
                timestamp_writes: None,
            });
            compute_pass.set_pipeline(&self.corners_pipeline);
            compute_pass.set_bind_group(0, bind_group, &[]);
            compute_pass.dispatch_workgroups(corner_workgroups, 1, 1);
        }

        let total_voxels = self.last_brick_count * (BRICK_SIZE * BRICK_SIZE * BRICK_SIZE) as u32;
        let material_workgroups = total_voxels.div_ceil(128);

        {
            let mut compute_pass = encoder.begin_compute_pass(&wgpu::ComputePassDescriptor {
                label: Some("SDF Evaluate Materials"),
                timestamp_writes: None,
            });
            compute_pass.set_pipeline(&self.materials_pipeline);
            compute_pass.set_bind_group(0, bind_group, &[]);
            compute_pass.dispatch_workgroups(material_workgroups, 1, 1);
        }
    }

    pub fn reset_frame(&mut self) {
        self.last_brick_count = 0;
    }

    pub fn has_pending_work(&self) -> bool {
        self.last_brick_count > 0
    }
}

fn serialize_edit(edit: &SdfEdit) -> GpuSdfEdit {
    let inverse_transform: [[f32; 4]; 4] = (*edit.inverse_transform()).into();

    let (primitive_type, param0, param1, param2, param3) = match *edit.primitive() {
        SdfPrimitive::Sphere { radius } => (0u32, radius, 0.0, 0.0, 0.0),
        SdfPrimitive::Box { half_extents } => {
            (1u32, half_extents.x, half_extents.y, half_extents.z, 0.0)
        }
        SdfPrimitive::Cylinder {
            radius,
            half_height,
        } => (2u32, radius, half_height, 0.0, 0.0),
        SdfPrimitive::Torus {
            major_radius,
            minor_radius,
        } => (3u32, major_radius, minor_radius, 0.0, 0.0),
        SdfPrimitive::Capsule {
            radius,
            half_height,
        } => (4u32, radius, half_height, 0.0, 0.0),
        SdfPrimitive::Plane { normal, offset } => (5u32, normal.x, normal.y, normal.z, offset),
    };

    let operation_type = match edit.operation() {
        CsgOperation::Union => 0u32,
        CsgOperation::Subtraction => 1u32,
        CsgOperation::Intersection => 2u32,
        CsgOperation::SmoothUnion { .. } => 3u32,
        CsgOperation::SmoothSubtraction { .. } => 4u32,
        CsgOperation::SmoothIntersection { .. } => 5u32,
    };

    GpuSdfEdit {
        inverse_transform,
        primitive_type,
        operation_type,
        material_id: edit.material_id(),
        smoothness: edit.smoothness(),
        uniform_scale: edit.uniform_scale(),
        param0,
        param1,
        param2,
        param3,
        _pad0: 0,
        _pad1: 0,
        _pad2: 0,
    }
}

pub fn prepare_dirty_brick_metadata(
    atlas_slot: u32,
    world_origin: nalgebra_glm::Vec3,
    voxel_size: f32,
) -> GpuDirtyBrick {
    let bricks_per_slice = BRICKS_PER_ATLAS_DIM * BRICKS_PER_ATLAS_DIM;
    let brick_z = atlas_slot as usize / bricks_per_slice;
    let remainder = atlas_slot as usize % bricks_per_slice;
    let brick_y = remainder / BRICKS_PER_ATLAS_DIM;
    let brick_x = remainder % BRICKS_PER_ATLAS_DIM;

    GpuDirtyBrick {
        world_origin: [world_origin.x, world_origin.y, world_origin.z],
        voxel_size,
        atlas_slot,
        material_atlas_x: (brick_x * BRICK_SIZE) as u32,
        material_atlas_y: (brick_y * BRICK_SIZE) as u32,
        material_atlas_z: (brick_z * BRICK_SIZE) as u32,
    }
}
