use std::sync::Arc;

use ash::vk;
use glam::{
    camera::rh::{proj, view},
    vec3,
};
use gpu_allocator::MemoryLocation;

pub use crate::egui_pass::{egui_pass, egui_pre_pass};

use crate::{
    Result,
    buffers::{Buffer, BufferDesc},
    commands::{self, traits::CommandBufferExt},
    device::Device,
    images::ImageViewDesc,
    pipeline,
    render_graph::{
        Access, BindPipeline, BufferRange, BufferSlice, Draw, DrawIndexedData, GraphBuffer,
        GraphBuilder, GraphResourceId, ImportResource, IndexFormat, PassDesc, PushConstants,
        RecordFn, RenderContext, RenderGraph, RenderingAttachment, SetScissor, SetViewport,
        Texture, TextureRange, TextureRegion, buffer_barrier,
    },
    sync,
    util::Rgba8,
};

pub struct Wireframe {
    positions: Arc<Buffer>,
    indices: Arc<Buffer>,
    num_indices: u32,
    colors: Arc<Buffer>,
    pipeline: Arc<pipeline::Pipeline>,
}

impl Wireframe {
    pub async fn from_triangles_indexed(
        dev: Device,
        positions: Vec<glam::Vec3>,
        indices: Vec<u32>,
        colors: Vec<Rgba8>,
    ) -> Result<Self> {
        let positions_size = positions.len() * size_of::<glam::Vec3>();
        let num_indices = indices.len() as u32;
        let indices_size = indices.len() * size_of::<u32>();
        let indices_offset = positions_size;
        let colors_size = colors.len() * size_of::<Rgba8>();
        let colors_offset = indices_offset + indices_size;

        let staging_size = positions_size + indices_size + colors_size;

        let mut staging = Buffer::new(
            dev.clone(),
            BufferDesc {
                name: Some("wireframe-staging".into()),
                size: staging_size as u64,
                usage: vk::BufferUsageFlags::TRANSFER_SRC,
                mem_location: MemoryLocation::CpuToGpu,
                ..Default::default()
            },
        )?;

        {
            let map = staging
                .map_mut()
                .expect("staging buffer should be host visible");

            map[..positions_size].copy_from_slice(bytemuck::cast_slice(&positions));
            map[indices_offset..][..indices_size].copy_from_slice(bytemuck::cast_slice(&indices));
            map[colors_offset..][..colors_size].copy_from_slice(bytemuck::cast_slice(&colors));
        }

        let positions = Buffer::new(
            dev.clone(),
            BufferDesc {
                name: Some("wireframe-positions".into()),
                size: positions_size as u64,
                usage: vk::BufferUsageFlags::TRANSFER_DST | vk::BufferUsageFlags::VERTEX_BUFFER,
                mem_location: MemoryLocation::GpuOnly,
                ..Default::default()
            },
        )?;
        let indices = Buffer::new(
            dev.clone(),
            BufferDesc {
                name: Some("wireframe-indices".into()),
                size: indices_size as u64,
                usage: vk::BufferUsageFlags::TRANSFER_DST | vk::BufferUsageFlags::INDEX_BUFFER,
                mem_location: MemoryLocation::GpuOnly,
                ..Default::default()
            },
        )?;
        let colors = Buffer::new(
            dev.clone(),
            BufferDesc {
                name: Some("wireframe-colors".into()),
                size: colors_size as u64,
                usage: vk::BufferUsageFlags::TRANSFER_DST | vk::BufferUsageFlags::VERTEX_BUFFER,
                mem_location: MemoryLocation::GpuOnly,
                ..Default::default()
            },
        )?;

        let pool = commands::SingleUseCommandPool::new(dev.clone(), dev.main_queue().clone())?;

        let cmd = pool.alloc()?;

        cmd.copy_buffers(
            staging.raw(),
            positions.raw(),
            &[vk::BufferCopy {
                src_offset: 0,
                dst_offset: 0,
                size: positions_size as u64,
            }],
        );
        cmd.copy_buffers(
            staging.raw(),
            indices.raw(),
            &[vk::BufferCopy {
                src_offset: indices_offset as u64,
                dst_offset: 0,
                size: indices_size as u64,
            }],
        );
        cmd.copy_buffers(
            staging.raw(),
            colors.raw(),
            &[vk::BufferCopy {
                src_offset: colors_offset as u64,
                dst_offset: 0,
                size: colors_size as u64,
            }],
        );

        let barriers = [
            buffer_barrier(
                positions.raw(),
                0,
                positions.len(),
                Access::transfer_write(),
                Access::vertex_read(),
                None,
            ),
            buffer_barrier(
                indices.raw(),
                0,
                indices.len(),
                Access::transfer_write(),
                Access::index_read(),
                None,
            ),
            buffer_barrier(
                colors.raw(),
                0,
                colors.len(),
                Access::transfer_write(),
                Access::vertex_read(),
                None,
            ),
        ];
        unsafe {
            dev.dev().cmd_pipeline_barrier2(
                cmd.buffer(),
                &vk::DependencyInfo::default().buffer_memory_barriers(&barriers),
            );
        }

        let pipeline = Self::create_pipeline(dev.clone())?;

        let future = cmd.submit_async(
            None,
            None,
            Arc::new(sync::Fence::from_pool(&dev.pools.fences, None)?),
        )?;

        future.await;

        Ok(Self {
            positions: Arc::new(positions),
            indices: Arc::new(indices),
            colors: Arc::new(colors),
            pipeline: Arc::new(pipeline),
            num_indices,
        })
    }

    fn create_pipeline(device: Device) -> Result<pipeline::Pipeline> {
        let pipeline_layout = Arc::new(pipeline::PipelineLayout::new(
            device.clone(),
            pipeline::PipelineLayoutDesc {
                descriptor_set_layouts: &[],
                push_constant_ranges: &[vk::PushConstantRange {
                    offset: 0,
                    size: 128,
                    stage_flags: vk::ShaderStageFlags::VERTEX,
                }],
                name: Some("wireframe-pipeline-layout".into()),
            },
        )?);

        let shader = pipeline::ShaderModule::new_from_path(
            device.clone(),
            "crates/renderer/shaders/wireframe.spv",
        )?;

        let pipeline = pipeline::Pipeline::new_graphics(
            device.clone(),
            pipeline::GraphicsPipelineDesc {
                flags: Default::default(),
                name: Some("wireframe-pipeline".into()),
                shader_stages: &[
                    pipeline::ShaderStageDesc {
                        flags: vk::PipelineShaderStageCreateFlags::empty(),
                        module: &shader,
                        stage: vk::ShaderStageFlags::FRAGMENT,
                        entry: c"fragment".into(),
                    },
                    pipeline::ShaderStageDesc {
                        flags: vk::PipelineShaderStageCreateFlags::empty(),
                        module: &shader,
                        stage: vk::ShaderStageFlags::VERTEX,
                        entry: c"vertex".into(),
                    },
                ],
                render_pass: None,
                layout: pipeline_layout,
                subpass: None,
                base_pipeline: None,
                vertex_input: Some(pipeline::VertexInputState {
                    bindings: &[
                        vk::VertexInputBindingDescription {
                            binding: 0,
                            stride: size_of::<glam::Vec3>() as u32,
                            input_rate: vk::VertexInputRate::VERTEX,
                        },
                        vk::VertexInputBindingDescription {
                            binding: 1,
                            stride: size_of::<Rgba8>() as u32,
                            input_rate: vk::VertexInputRate::VERTEX,
                        },
                    ],
                    attributes: &[
                        vk::VertexInputAttributeDescription {
                            location: 0,
                            binding: 0,
                            format: vk::Format::R32G32B32_SFLOAT,
                            offset: 0,
                        },
                        vk::VertexInputAttributeDescription {
                            location: 1,
                            binding: 1,
                            format: vk::Format::R8G8B8A8_UNORM,
                            offset: 0,
                        },
                    ],
                }),
                input_assembly: Some(pipeline::InputAssemblyState {
                    topology: vk::PrimitiveTopology::TRIANGLE_LIST,
                    primitive_restart: false,
                }),
                tessellation: None,
                viewport: Some(pipeline::ViewportState {
                    num_viewports: 1,
                    num_scissors: 1,
                    ..Default::default()
                }),
                rasterization: Some(pipeline::RasterizationState {
                    cull_mode: vk::CullModeFlags::NONE,
                    polygon_mode: vk::PolygonMode::LINE,
                    ..Default::default()
                }),
                multisample: Some(pipeline::MultisampleState {
                    ..Default::default()
                }),
                depth_stencil: Some(pipeline::DepthStencilState {
                    depth: Some(pipeline::DepthState {
                        write_enable: false,
                        compare_op: Some(vk::CompareOp::LESS),
                        bounds: Some(pipeline::DepthBounds { min: 0.0, max: 1.0 }),
                    }),
                    ..Default::default()
                }),
                color_blend: Some(pipeline::ColorBlendState {
                    attachments: &[vk::PipelineColorBlendAttachmentState::default()
                        .color_write_mask(vk::ColorComponentFlags::RGBA)
                        .blend_enable(true)
                        .color_blend_op(vk::BlendOp::ADD)
                        .src_color_blend_factor(vk::BlendFactor::ONE)
                        .dst_color_blend_factor(vk::BlendFactor::ONE_MINUS_SRC_ALPHA)
                        .alpha_blend_op(vk::BlendOp::ADD)
                        .src_alpha_blend_factor(vk::BlendFactor::ONE)
                        .dst_alpha_blend_factor(vk::BlendFactor::ONE_MINUS_SRC_ALPHA)],
                    ..Default::default()
                }),
                rendering: Some(pipeline::RenderingState {
                    color_formats: &[vk::Format::R8G8B8A8_UNORM],
                    ..Default::default()
                }),
                dynamic: Some(pipeline::DynamicState {
                    dynamic_states: &[vk::DynamicState::VIEWPORT, vk::DynamicState::SCISSOR],
                    ..Default::default()
                }),
            },
        )?;

        Ok(pipeline)
    }

    pub fn pass(&self, rg: &mut RenderGraph, target: GraphResourceId) -> Result<()> {
        let reads = [(target, Access::color_attachment_read_write())].to_vec();
        let writes = [(target, Access::color_attachment_write_only())].to_vec();

        let record: Box<RecordFn> = Box::new({
            let positions = self.positions.clone();
            let indices = self.indices.clone();
            let colors = self.colors.clone();
            let num_indices = self.num_indices;

            let pipeline = self.pipeline.clone();

            move |ctx: &RenderContext| -> Result<()> {
                let target = ctx.get_image(target).unwrap();
                let cmd = &ctx.cmd;

                //let model = Mat4::from_scale(vec3(0.01, 0.01, 0.01));
                let proj = proj::vulkan::perspective(
                    core::f32::consts::FRAC_PI_4,
                    target.width() as f32 / target.height() as f32,
                    0.1,
                    100.0,
                );
                let pos = vec3(0.0, 0.0, -5.0);
                let view = view::look_at_mat4(pos, pos + vec3(0.0, 0.0, 1.0), vec3(0.0, 1.0, 0.0));
                let to_clip = proj * view; // * model;

                let color_attachment = &vk::RenderingAttachmentInfo::default()
                    .image_layout(vk::ImageLayout::COLOR_ATTACHMENT_OPTIMAL)
                    .image_view(target.create_view(ImageViewDesc::color_2d())?.raw())
                    .load_op(vk::AttachmentLoadOp::LOAD)
                    .store_op(vk::AttachmentStoreOp::STORE);

                cmd.begin_rendering(
                    vk::RenderingInfo::default()
                        .color_attachments(core::slice::from_ref(color_attachment))
                        .layer_count(1)
                        .render_area(vk::Rect2D::default().extent(target.extent_2d())),
                );

                cmd.set_scissors(&[vk::Rect2D::default()
                    .offset(vk::Offset2D::default())
                    .extent(target.extent_2d())]);

                cmd.set_viewport(&[vk::Viewport::default()
                    .x(0.0)
                    .y(0.0)
                    .min_depth(0.0)
                    .max_depth(1.0)
                    .width(target.width() as f32)
                    .height(target.height() as f32)]);

                cmd.bind_pipeline(&pipeline);
                cmd.bind_indices(indices.raw(), 0, vk::IndexType::UINT32);
                cmd.bind_vertex_buffers(&[positions.raw(), colors.raw()], &[0, 0]);
                cmd.push_constants(
                    &pipeline.layout(),
                    vk::ShaderStageFlags::VERTEX,
                    0,
                    bytemuck::cast_slice(&[to_clip]),
                );
                cmd.draw_indexed(num_indices, 1, 0, 0, 0);
                cmd.end_rendering();

                Ok(())
            }
        });

        rg.add_pass(PassDesc {
            reads,
            writes,
            record: Some(record),
        });
        Ok(())
    }

    pub fn typed_pass(
        &self,
        builder: &mut GraphBuilder<'_>,
        target: Texture,
        extent: (u32, u32),
        list: &mut crate::render_graph::CommandList<'_>,
    ) {
        // Buffers are shared with this long-lived scene object, so the typed
        // graph retains Arc references rather than taking ownership of them.
        let positions_size = self.positions.len();
        let indices_size = self.indices.len();
        let colors_size = self.colors.len();
        let positions = builder.import_shared_buffer(self.positions.clone(), None);
        let indices = builder.import_shared_buffer(self.indices.clone(), None);
        let colors = builder.import_shared_buffer(self.colors.clone(), None);

        list.push(ImportResource {
            resource: GraphBuffer(positions),
            access: vk::AccessFlags2::VERTEX_ATTRIBUTE_READ,
            stage: vk::PipelineStageFlags2::VERTEX_ATTRIBUTE_INPUT,
            layout: None,
        });
        list.push(ImportResource {
            resource: GraphBuffer(indices),
            access: vk::AccessFlags2::INDEX_READ,
            stage: vk::PipelineStageFlags2::INDEX_INPUT,
            layout: None,
        });
        list.push(ImportResource {
            resource: GraphBuffer(colors),
            access: vk::AccessFlags2::VERTEX_ATTRIBUTE_READ,
            stage: vk::PipelineStageFlags2::VERTEX_ATTRIBUTE_INPUT,
            layout: None,
        });

        let proj = proj::vulkan::perspective(
            core::f32::consts::FRAC_PI_4,
            extent.0 as f32 / extent.1 as f32,
            0.1,
            100.0,
        );
        let pos = vec3(0.0, 0.0, -5.0);
        let view = view::look_at_mat4(pos, pos + vec3(0.0, 0.0, 1.0), vec3(0.0, 1.0, 0.0));
        let to_clip = proj * view;

        let mut pass = list.begin_rendering(
            vec![RenderingAttachment {
                region: TextureRegion {
                    texture: target,
                    range: TextureRange {
                        aspect: vk::ImageAspectFlags::COLOR,
                        mip_levels: (0..1).into(),
                        array_layers: (0..1).into(),
                        origin: (0, 0, 0),
                        extent: (extent.0, extent.1, 1),
                    },
                },
                load_op: vk::AttachmentLoadOp::CLEAR,
                store_op: vk::AttachmentStoreOp::STORE,
                clear_value: vk::ClearValue {
                    color: vk::ClearColorValue {
                        float32: [0.1, 0.1, 0.12, 1.0],
                    },
                },
            }],
            None,
            None,
            extent,
            1,
        );
        pass.push(SetScissor {
            x: 0,
            y: 0,
            width: extent.0,
            height: extent.1,
        });
        pass.push(SetViewport {
            x: 0.0,
            y: 0.0,
            width: extent.0 as f32,
            height: extent.1 as f32,
            min_depth: 0.0,
            max_depth: 1.0,
        });
        pass.push(BindPipeline(
            self.pipeline.raw(),
            vk::PipelineBindPoint::GRAPHICS,
        ));
        pass.push(PushConstants {
            pipeline_layout: self.pipeline.layout().raw(),
            stage_flags: vk::ShaderStageFlags::VERTEX,
            offset: 0,
            data: bytemuck::cast_slice(&[to_clip]).to_vec(),
        });
        pass.push(Draw {
            data: DrawIndexedData {
                index_count: self.num_indices,
                instance_count: 1,
                first_index: 0,
                vertex_offset: 0,
                first_instance: 0,
            },
            vertex_buffers: vec![
                BufferSlice {
                    buffer: GraphBuffer(positions),
                    range: BufferRange {
                        offset: 0,
                        size: positions_size,
                    },
                },
                BufferSlice {
                    buffer: GraphBuffer(colors),
                    range: BufferRange {
                        offset: 0,
                        size: colors_size,
                    },
                },
            ],
            index_buffer: Some((
                BufferSlice {
                    buffer: GraphBuffer(indices),
                    range: BufferRange {
                        offset: 0,
                        size: indices_size,
                    },
                },
                IndexFormat::Uint32,
            )),
        });
        pass.finalise();
    }
}