use std::sync::Arc;

use ash::vk;
use glam::{f32::Mat4, vec3};

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

use crate::{
    Result,
    buffers::{Buffer, BufferDesc},
    commands::{self, traits::CommandBufferExt},
    device::{Device, DeviceOwned},
    images::ImageViewDesc,
    pipeline,
    render_graph::{
        Access, GraphResourceId, PassDesc, RecordFn, RenderContext, RenderGraph, buffer_barrier,
    },
    sync,
    util::Rgba8,
};

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

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_usage: vk_mem::MemoryUsage::AutoPreferHost,
                alloc_flags: vk_mem::AllocationCreateFlags::MAPPED
                    | vk_mem::AllocationCreateFlags::HOST_ACCESS_SEQUENTIAL_WRITE
                    | vk_mem::AllocationCreateFlags::STRATEGY_FIRST_FIT,
                ..Default::default()
            },
        )?;

        {
            let mut map = staging.map()?;

            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_usage: vk_mem::MemoryUsage::AutoPreferDevice,
                ..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_usage: vk_mem::MemoryUsage::AutoPreferDevice,
                ..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_usage: vk_mem::MemoryUsage::AutoPreferDevice,
                ..Default::default()
            },
        )?;

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

        let cmd = pool.alloc()?;

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

        let barriers = [
            buffer_barrier(
                positions.handle(),
                0,
                positions.len(),
                Access::transfer_write(),
                Access::vertex_read(),
                None,
            ),
            buffer_barrier(
                indices.handle(),
                0,
                indices.len(),
                Access::transfer_write(),
                Access::index_read(),
                None,
            ),
            buffer_barrier(
                colors.handle(),
                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 future = cmd.submit_async(
            None,
            None,
            Arc::new(sync::Fence::from_pool(&dev.pools.fences, None)?),
        )?;

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

        future.await;

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

    fn create_pipeline(device: Device) -> Result<(pipeline::Pipeline, pipeline::PipelineLayout)> {
        let pipeline_layout = 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(
            device.clone(),
            pipeline::PipelineDesc::Graphics(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"main".into(),
                    },
                    pipeline::ShaderStageDesc {
                        flags: vk::PipelineShaderStageCreateFlags::empty(),
                        module: &shader,
                        stage: vk::ShaderStageFlags::VERTEX,
                        entry: c"main".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, pipeline_layout))
    }

    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();
            let layout = self.layout.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 = Mat4::perspective_lh(
                    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 = Mat4::look_at_rh(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.get_view(ImageViewDesc {
                        kind: vk::ImageViewType::TYPE_2D,
                        format: target.format(),
                        aspect: vk::ImageAspectFlags::COLOR,
                        ..Default::default()
                    })?)
                    .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.buffer(), 0, vk::IndexType::UINT32);
                cmd.bind_vertex_buffers(&[positions.handle(), colors.handle()], &[0, 0]);
                cmd.push_constants(
                    &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(())
    }
}