use super::*;
use crate::storage::GpuVec;
use crate::storage::mesh::*;
use crate::shader::{ShaderHandle, ShaderStore};
use crate::viewport::ViewportInfo;
use crate::RenderLayouts;

pub struct ShaderInfo {
    pub store: Arc<ShaderStore>,
    pub forward: ShaderHandle,
    pub skinning: ShaderHandle,
}

#[repr(C)]
#[derive(Copy, Clone, bytemuck::Pod, bytemuck::Zeroable)]
pub struct Vertex {
    pub position: [f32; 3],
    pub tan_frame: u32,
}

impl Attribute for Vertex {
    fn get_usages() -> wgpu::BufferUsages {
        wgpu::BufferUsages::STORAGE | wgpu::BufferUsages::VERTEX
    }
}

const VERTEX_ATTRS: &[wgpu::VertexAttribute] =
    &wgpu::vertex_attr_array![0 => Float32x3, 1 => Uint32];

impl Vertex {
    pub fn desc() -> wgpu::VertexBufferLayout<'static> {
        wgpu::VertexBufferLayout {
            array_stride: std::mem::size_of::<Self>() as wgpu::BufferAddress,
            step_mode: wgpu::VertexStepMode::Vertex,
            attributes: VERTEX_ATTRS,
        }
    }
}

pub type Index = u32;

#[repr(C)]
#[derive(Copy, Clone, bytemuck::Pod, bytemuck::Zeroable)]
struct SkinningUniform {
    transform: [[f32; 4]; 4],
    src_offset: u32,
    dst_offset: u32,
    count: u32,
    _padding: u32,
}

#[derive(Debug)]
struct MeshCommand {
    vertex_offset: usize,
    vertex_count: usize,
    skinned_offset: usize,
    index_offset: usize,
    index_count: usize,
    skinning_index: usize,
}

#[derive(Debug)]
struct MeshGroupCommands {
    binding_indices: MeshLayoutBindingIndices,
    bind_group: wgpu::BindGroup,
    meshes: Vec<MeshCommand>,
}

pub struct MeshInstance {
    pub transform: glam::Mat4,
    pub mesh: MeshHandle,
}

pub struct FrameData {
    skinned_vertices: GpuVec<Vertex>,
    skinning_uniforms: GpuVec<SkinningUniform>,
    groups: Vec<MeshGroupCommands>,
}

pub struct MeshPass {
    device: Arc<wgpu::Device>,
    layouts: Arc<RenderLayouts>,
    attr_store: Arc<AttrStore>,
    shader_info: ShaderInfo,
    mesh_pool: Arc<MeshPool>,
    vertex_attr_id: AttrId,
    index_attr_id: AttrId,
    mesh_layout_id: MeshLayoutId,
    example_mesh: MeshHandle,
    skinning_bind_group_layout: wgpu::BindGroupLayout,
    skinning_pipeline: wgpu::ComputePipeline,
    depth_pipeline: wgpu::RenderPipeline,
    opaque_pipeline: wgpu::RenderPipeline,
    target_info: ViewportInfo,
    instances: Vec<MeshInstance>,
}

impl MeshPass {
    pub fn new(
        device: Arc<wgpu::Device>,
        layouts: Arc<RenderLayouts>,
        target_info: ViewportInfo,
        shader_info: ShaderInfo,
    ) -> Self {
        let attr_store = AttrStore::new();
        let mesh_pool = MeshPool::new(device.clone(), attr_store.to_owned());

        let vertex_attr_id = attr_store.get_type::<Vertex>();
        let index_attr_id = attr_store.add(AttrInfo {
            layout: AttrLayout {
                size: std::mem::size_of::<Index>(),
            },
            usages: wgpu::BufferUsages::INDEX,
            default_pool_size: 1_000_000,
        });

        let mut mesh_layout = MeshLayoutDesc::new();
        mesh_layout.insert(vertex_attr_id, ());
        mesh_layout.insert(index_attr_id, ());
        let mesh_layout_id = mesh_pool.add_layout(mesh_layout).unwrap();

        let example_vertices = vec![
            Vertex {
                position: [-0.5, 0.5, 0.0],
                tan_frame: 0,
            },
            Vertex {
                position: [0.5, 0.5, 0.0],
                tan_frame: 0,
            },
            Vertex {
                position: [0.0, -0.5, 0.0],
                tan_frame: 0,
            },
        ];

        let example_vertices = AttrBuffer {
            id: vertex_attr_id,
            count: example_vertices.len(),
            data: bytemuck::cast_slice(&example_vertices).to_vec(),
        };

        let example_indices = vec![0u32, 1u32, 2u32];
        let example_indices = AttrBuffer {
            id: index_attr_id,
            count: example_indices.len(),
            data: bytemuck::cast_slice(&example_indices).to_vec(),
        };

        let mut example_mesh = MeshBuffer::default();
        example_mesh.attributes.push(example_vertices);
        example_mesh.attributes.push(example_indices);
        let example_mesh = mesh_pool.load(example_mesh).unwrap();

        let mut instances = Vec::new();
        let r = 4;
        for x in -r..r {
            for y in -r..r {
                for z in -r..r {
                    let translation = glam::Vec3::new(x as f32, y as f32, z as f32);
                    let transform = glam::Mat4::from_translation(translation);
                    instances.push(MeshInstance {
                        transform,
                        mesh: example_mesh.clone(),
                    });
                }
            }
        }

        let render_pipeline_layout =
            device.create_pipeline_layout(&wgpu::PipelineLayoutDescriptor {
                label: Some("MeshPass Pipeline Layout"),
                bind_group_layouts: &[&layouts.bind_viewport],
                push_constant_ranges: &[],
            });

        let shader = shader_info.store.get(&shader_info.forward).unwrap();

        let targets = &[wgpu::ColorTargetState {
            format: target_info.output_format,
            blend: Some(wgpu::BlendState::REPLACE),
            write_mask: wgpu::ColorWrites::ALL,
        }];

        let mut pipeline_desc = wgpu::RenderPipelineDescriptor {
            label: Some("Opaque MeshPass Pipeline"),
            layout: Some(&render_pipeline_layout),
            vertex: wgpu::VertexState {
                module: shader.as_ref(),
                entry_point: "vs_main",
                buffers: &[Vertex::desc()],
            },
            fragment: Some(wgpu::FragmentState {
                module: shader.as_ref(),
                entry_point: "fs_main",
                targets,
            }),
            primitive: wgpu::PrimitiveState {
                topology: wgpu::PrimitiveTopology::TriangleList,
                strip_index_format: None,
                front_face: wgpu::FrontFace::Ccw,
                cull_mode: None, // Some(wgpu::Face::Back),
                polygon_mode: wgpu::PolygonMode::Fill,
                unclipped_depth: false,
                conservative: false,
            },
            depth_stencil: Some(wgpu::DepthStencilState {
                format: target_info.depth_format,
                depth_write_enabled: true,
                depth_compare: wgpu::CompareFunction::Less,
                stencil: Default::default(),
                bias: Default::default(),
            }),
            multisample: wgpu::MultisampleState {
                count: 1,
                mask: !0,
                alpha_to_coverage_enabled: false,
            },
            multiview: None,
        };

        let depth_pipeline = device.create_render_pipeline(&pipeline_desc);

        pipeline_desc.depth_stencil = Some(wgpu::DepthStencilState {
            format: target_info.depth_format,
            depth_write_enabled: false,
            depth_compare: wgpu::CompareFunction::Equal,
            stencil: Default::default(),
            bias: Default::default(),
        });

        let opaque_pipeline = device.create_render_pipeline(&pipeline_desc);

        drop(shader);

        let skinning_bind_group_layout =
            device.create_bind_group_layout(&wgpu::BindGroupLayoutDescriptor {
                label: Some("Skinning Bind Group Layout"),
                entries: &[
                    wgpu::BindGroupLayoutEntry {
                        binding: 0,
                        visibility: wgpu::ShaderStages::COMPUTE,
                        ty: wgpu::BindingType::Buffer {
                            ty: wgpu::BufferBindingType::Storage { read_only: true },
                            has_dynamic_offset: true,
                            min_binding_size: None, // TODO ???
                        },
                        count: None,
                    },
                    wgpu::BindGroupLayoutEntry {
                        binding: 1,
                        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,
                    },
                    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,
                    },
                ],
            });

        let skinning_pipeline_layout =
            device.create_pipeline_layout(&wgpu::PipelineLayoutDescriptor {
                label: Some("Skinning Pipeline Layout"),
                bind_group_layouts: &[&skinning_bind_group_layout],
                push_constant_ranges: &[],
            });

        let skinning_pipeline = device.create_compute_pipeline(&wgpu::ComputePipelineDescriptor {
            label: Some("Skinning Pipeline"),
            layout: Some(&skinning_pipeline_layout),
            module: shader_info
                .store
                .get(&shader_info.skinning)
                .unwrap()
                .as_ref(),
            entry_point: "cs_main",
        });

        Self {
            device,
            layouts,
            attr_store,
            shader_info,
            mesh_pool,
            index_attr_id,
            vertex_attr_id,
            mesh_layout_id,
            example_mesh,
            skinning_bind_group_layout,
            skinning_pipeline,
            depth_pipeline,
            opaque_pipeline,
            target_info,
            instances,
        }
    }
}

impl RenderPass for MeshPass {
    type FrameData = FrameData;

    fn create_frame_data(&self) -> FrameData {
        FrameData {
            skinned_vertices: GpuVec::new(
                self.device.clone(),
                Vertex::get_usages(),
                1024 * 128,
                Some("Skinned Vertices".to_string()),
                false,
            ),
            skinning_uniforms: GpuVec::new(
                self.device.clone(),
                wgpu::BufferUsages::STORAGE,
                1024 * 128,
                Some("Skinning Uniforms".to_string()),
                true,
            ),
            groups: Default::default(),
        }
    }

    fn begin_frame(&self, data: &mut FrameData, phases: &mut Vec<Phase>, queue: &wgpu::Queue) {
        println!("MeshPass::begin_frame()");

        phases.push(Phase::Upload);
        phases.push(Phase::Skinning);
        phases.push(Phase::Depth);
        phases.push(Phase::Opaque);
        phases.push(Phase::Transparent);

        data.groups.clear();
        data.skinning_uniforms.clear();

        let mesh_bindings = self
            .mesh_pool
            .iter_meshes(self.mesh_layout_id, self.instances.iter(), |i| &i.mesh)
            .unwrap();

        let mut skinned_cursor = 0;

        for MeshLayoutInstances {
            bindings,
            instances,
        } in mesh_bindings.iter()
        {
            let pools = self.mesh_pool.get_bindings(bindings.clone());
            let vertices_pool = pools.get(self.vertex_attr_id).unwrap();

            // TODO defer bind group creation into separate Vec after GpuVecs have been written
            let bind_group = self.device.create_bind_group(&wgpu::BindGroupDescriptor {
                label: None,
                layout: &self.skinning_bind_group_layout,
                entries: &[
                    wgpu::BindGroupEntry {
                        binding: 0,
                        resource: wgpu::BindingResource::Buffer(wgpu::BufferBinding {
                            buffer: data.skinning_uniforms.as_ref(),
                            offset: 0,
                            // TODO ugly!
                            size: Some(
                                std::num::NonZeroU64::new(
                                    std::mem::size_of::<SkinningUniform>() as u64
                                )
                                .unwrap(),
                            ),
                        }),
                    },
                    wgpu::BindGroupEntry {
                        binding: 1,
                        resource: wgpu::BindingResource::Buffer(wgpu::BufferBinding {
                            buffer: data.skinned_vertices.as_ref(),
                            offset: 0,
                            size: None,
                        }),
                    },
                    wgpu::BindGroupEntry {
                        binding: 2,
                        resource: wgpu::BindingResource::Buffer(wgpu::BufferBinding {
                            buffer: vertices_pool.get_buffer(),
                            offset: 0,
                            size: None,
                        }),
                    },
                ],
            });

            let mut group = MeshGroupCommands {
                binding_indices: bindings.clone(),
                meshes: Default::default(),
                bind_group,
            };

            for (mesh, infos) in instances {
                let vertices = infos.iter().find(|i| i.0 == self.vertex_attr_id).unwrap().1;
                let indices = infos.iter().find(|i| i.0 == self.index_attr_id).unwrap().1;

                group.meshes.push(MeshCommand {
                    vertex_offset: vertices.offset,
                    vertex_count: vertices.count,
                    skinned_offset: skinned_cursor,
                    index_offset: indices.offset,
                    index_count: indices.count,
                    skinning_index: data.skinning_uniforms.len(),
                });

                data.skinning_uniforms.push(SkinningUniform {
                    transform: mesh.transform.to_cols_array_2d(),
                    src_offset: vertices.offset as u32,
                    dst_offset: skinned_cursor as u32,
                    count: vertices.count as u32,
                    _padding: 0,
                });

                skinned_cursor += vertices.count;
            }

            data.groups.push(group);
        }

        data.skinned_vertices.reserve(skinned_cursor);
        data.skinned_vertices.write(&queue);
        data.skinning_uniforms.write(&queue);
    }

    fn record_commands(&self, data: PhaseData<&FrameData>, cmds: &mut wgpu::CommandEncoder) {
        match data.phase {
            Phase::Upload => self.mesh_pool.flush(cmds),
            _ => {}
        }
    }

    fn record_compute<'a>(
        &'a self,
        data: PhaseData<&'a FrameData>,
        cmds: &mut wgpu::ComputePass<'a>,
    ) {
        cmds.set_pipeline(&self.skinning_pipeline);
        for group in data.frame_data.groups.iter() {
            for mesh in group.meshes.iter() {
                let ubo_offset = data
                    .frame_data
                    .skinning_uniforms
                    .buf_offset(mesh.skinning_index);
                cmds.set_bind_group(0, &group.bind_group, &[ubo_offset as u32]);

                // TODO use div_ceil instead
                let workgroup_num = if mesh.vertex_count % 64 == 0 {
                    mesh.vertex_count / 64
                } else {
                    mesh.vertex_count / 64 + 1
                };

                cmds.dispatch(workgroup_num as u32, 1, 1);
            }
        }
    }

    fn record_render(&self, data: PhaseData<&FrameData>) -> Option<wgpu::RenderBundle> {
        println!("MeshPass::record_render(phase: {:?})", data.phase);

        let pipeline = match data.phase {
            Phase::Depth => &self.depth_pipeline,
            Phase::Opaque => &self.opaque_pipeline,
            _ => return None,
        };

        let mut cmds =
            self.device
                .create_render_bundle_encoder(&wgpu::RenderBundleEncoderDescriptor {
                    label: Some("Opaque Pass Render Bundle"),
                    color_formats: &[self.target_info.output_format],
                    depth_stencil: Some(wgpu::RenderBundleDepthStencil {
                        format: self.target_info.depth_format,
                        depth_read_only: false, // TODO optimize?
                        stencil_read_only: true,
                    }),
                    sample_count: 1,
                    multiview: None,
                });

        let meshes = &[&self.example_mesh, &self.example_mesh];

        // yikes
        let mesh_bindings: Vec<(MeshLayoutBindings, Vec<(&&MeshHandle, MeshAllocInfos)>)> = self
            .mesh_pool
            .iter_meshes(self.mesh_layout_id, meshes.iter(), |v| v)
            .unwrap()
            .into_iter()
            .map(
                |MeshLayoutInstances {
                     bindings,
                     instances,
                 }| (self.mesh_pool.get_bindings(bindings), instances),
            )
            .collect();

        cmds.set_pipeline(pipeline);
        cmds.set_bind_group(0, data.bind_viewport, &[]);

        for (bindings, instances) in mesh_bindings.iter() {
            let vertices_pool = bindings.get(self.vertex_attr_id).unwrap();
            let indices_pool = bindings.get(self.index_attr_id).unwrap();

            cmds.set_vertex_buffer(0, vertices_pool.get_buffer().slice(..));
            cmds.set_index_buffer(
                indices_pool.get_buffer().slice(..),
                wgpu::IndexFormat::Uint32,
            );

            for (mesh, infos) in instances {
                let vertices = infos.iter().find(|i| i.0 == self.vertex_attr_id).unwrap().1;
                let indices = infos.iter().find(|i| i.0 == self.index_attr_id).unwrap().1;

                let is_start = indices.offset as u32;
                let is_end = is_start + indices.count as u32;

                cmds.draw_indexed(is_start..is_end, vertices.offset as i32, 0..1);
                println!(
                    "drew a mesh! {}..{} + {}",
                    is_start, is_end, vertices.offset
                );
            }
        }

        Some(cmds.finish(&wgpu::RenderBundleDescriptor::default()))
    }
}