use std::{collections::BTreeMap, sync::Arc};

use ash::vk;
use gpu_allocator::MemoryLocation;
use indexmap::IndexMap;

use crate::{
    EguiState,
    buffers::{Buffer, BufferDesc},
    commands::traits::CommandBufferExt,
    device::{self, AllocationStrategy},
    images::{Image, ImageDesc, ImageViewDesc},
    render_graph::{
        Access, Barrier, GraphResourceDesc, GraphResourceId, PassDesc, RecordFn, RenderContext,
        RenderGraph, buffer_barrier, image_barrier,
    },
    texture,
    util::Rect2D,
};

pub fn egui_pre_pass(
    dev: &device::Device,
    rg: &mut RenderGraph,
    textures: &mut crate::texture::TextureManager,
    egui_state: &mut EguiState,
    output: &egui::FullOutput,
) -> crate::Result<()> {
    // allocate resource ids for textures in tessellated list (imported from texture manager)
    // define accesses for resource ids

    if output.textures_delta.set.is_empty() {
        return Ok(());
    }

    // create textures for new egui textures
    for (egui_id, delta) in output
        .textures_delta
        .set
        .iter()
        .filter(|(_, image)| image.is_whole())
    {
        tracing::trace!("creating texture image for egui image {egui_id:?}");
        let image = Image::new(
            dev.clone(),
            ImageDesc {
                name: Some(format!("egui-texture-{egui_id:?}").into()),
                format: vk::Format::R8G8B8A8_UNORM,
                extent: vk::Extent3D {
                    width: delta.image.width() as u32,
                    height: delta.image.height() as u32,
                    depth: 1,
                },
                usage: vk::ImageUsageFlags::SAMPLED | vk::ImageUsageFlags::TRANSFER_DST,
                mem_location: MemoryLocation::GpuOnly,
                alloc_scheme: AllocationStrategy::Dedicated,
                ..Default::default()
            },
        )?;

        let tid = textures.insert_image(Arc::new(image));
        if let Some(old) = egui_state.textures.insert(
            *egui_id,
            crate::EguiTextureInfo {
                id: tid,
                options: delta.options,
            },
        ) {
            textures.remove_texture(old.id);
        }
    }

    // calculate size for staging buffer.
    // calculate size for staging image.
    let (staging_size, image_size) = output.textures_delta.set.iter().fold(
        (0usize, glam::UVec2::ZERO),
        |(mut buffer, mut image), (_id, delta)| {
            let bytes = delta.image.height() * delta.image.width() * delta.image.bytes_per_pixel();
            image = image.max(glam::uvec2(
                delta.image.width() as u32,
                delta.image.height() as u32,
            ));
            buffer += bytes;

            (buffer, image)
        },
    );

    tracing::trace!(
        staging_size,
        "creating staging buffer for uploading egui textures"
    );
    let mut staging_buffer = Buffer::new(
        dev.clone(),
        BufferDesc {
            name: Some("egui-prepass-staging-buffer".into()),
            size: staging_size as u64,
            usage: vk::BufferUsageFlags::TRANSFER_SRC,
            queue_families: device::QueueFlags::empty(),
            mem_location: MemoryLocation::CpuToGpu,
            ..Default::default()
        },
    )?;

    tracing::trace!("creating staging image for uploading egui textures with dims={image_size:?}");
    let staging_image = Arc::new(Image::new(
        dev.clone(),
        ImageDesc {
            name: Some("egui-prepass-staging-buffer".into()),
            format: vk::Format::R8G8B8A8_UNORM,
            extent: vk::Extent3D {
                width: image_size.x,
                height: image_size.y,
                depth: 1,
            },
            usage: vk::ImageUsageFlags::TRANSFER_SRC | vk::ImageUsageFlags::TRANSFER_DST,
            queue_families: device::QueueFlags::empty(),
            mem_location: MemoryLocation::GpuOnly,
            ..Default::default()
        },
    )?);

    let aliased_images = {
        tracing::trace!("mmap-ing staging buffer");
        let mut staging_map = staging_buffer.map().unwrap();
        let mut offset = 0;

        output
            .textures_delta
            .set
            .iter()
            .map(|(id, delta)| {
                let bytes =
                    delta.image.height() * delta.image.width() * delta.image.bytes_per_pixel();

                let mem = &mut staging_map[offset..offset + bytes];
                match &delta.image {
                    egui::ImageData::Color(arc) => {
                        let slice = unsafe {
                            core::slice::from_raw_parts(
                                arc.pixels.as_ptr().cast::<u8>(),
                                arc.pixels.len() * size_of::<egui::Color32>(),
                            )
                        };
                        mem[..slice.len()].copy_from_slice(slice);
                    }
                }

                let old_offset = offset;
                offset += bytes;

                let pos = delta.pos.unwrap_or_default();
                let rect = Rect2D::new_from_size(
                    glam::ivec2(pos[0] as i32, pos[1] as i32),
                    glam::ivec2(delta.image.width() as i32, delta.image.height() as i32),
                );
                (*id, (old_offset, bytes, rect))
            })
            .collect::<BTreeMap<_, _>>()

        // let tessellated = egui.tessellate(output.shapes, output.pixels_per_point);
    };

    let textures = output
        .textures_delta
        .set
        .iter()
        .filter_map(|(egui_id, _)| {
            egui_state
                .lookup_texture(*egui_id)
                .and_then(|tid| textures.get_texture(tid))
                .map(|img| (*egui_id, img))
        })
        .map(|(id, img)| {
            (
                id,
                rg.import_image(
                    img,
                    Access {
                        layout: vk::ImageLayout::GENERAL,
                        ..Access::undefined()
                    },
                ),
            )
        })
        .collect::<BTreeMap<_, _>>();

    let staging_buffer = rg.import_buffer(Arc::new(staging_buffer), Access::undefined());
    let staging_image = rg.import_image(staging_image, Access::undefined());

    let record = Box::new({
        let textures = textures.clone();
        move |ctx: &RenderContext| -> crate::Result<()> {
            let staging_image = ctx.get_image(staging_image).unwrap().clone();
            let staging_buffer = ctx.get_buffer(staging_buffer).unwrap();

            for (id, (offset, _, rect)) in aliased_images {
                tracing::trace!(
                    "record-prepass: fetching alias of prepass staging image id={id:?}"
                );
                let alias = staging_image.get_alias(ImageDesc {
                    name: Some(format!("egui-prepass-staging-aliased-{id:?}v").into()),
                    format: vk::Format::R8G8B8A8_UNORM,
                    extent: vk::Extent3D {
                        width: rect.width() as u32,
                        height: rect.height() as u32,
                        depth: 1,
                    },
                    usage: vk::ImageUsageFlags::TRANSFER_SRC | vk::ImageUsageFlags::TRANSFER_DST,
                    queue_families: device::QueueFlags::empty(),
                    ..Default::default()
                })?;

                let texture = textures.get(&id).and_then(|id| ctx.get_image(*id)).unwrap();

                let image: Barrier = image_barrier(
                    alias.raw(),
                    alias.format(),
                    Access {
                        stage: vk::PipelineStageFlags2::NONE,
                        mask: vk::AccessFlags2::empty(),
                        layout: vk::ImageLayout::UNDEFINED,
                    },
                    Access {
                        stage: vk::PipelineStageFlags2::TRANSFER,
                        mask: vk::AccessFlags2::TRANSFER_WRITE,
                        layout: vk::ImageLayout::TRANSFER_DST_OPTIMAL,
                    },
                    None,
                )
                .into();

                unsafe {
                    ctx.device
                        .dev()
                        .cmd_pipeline_barrier2(ctx.cmd.buffer(), &((&image).into()));
                }

                ctx.cmd.copy_buffer_to_image(
                    staging_buffer.raw(),
                    alias.raw(),
                    vk::ImageLayout::TRANSFER_DST_OPTIMAL,
                    &[vk::BufferImageCopy {
                        buffer_offset: offset as u64,
                        buffer_row_length: alias.width(),
                        buffer_image_height: alias.height(),
                        image_subresource: vk::ImageSubresourceLayers::default()
                            .aspect_mask(vk::ImageAspectFlags::COLOR)
                            .base_array_layer(0)
                            .mip_level(0)
                            .layer_count(1),
                        image_offset: vk::Offset3D { x: 0, y: 0, z: 0 },
                        image_extent: alias.size(),
                    }],
                );

                let from_barrier = image_barrier(
                    alias.raw(),
                    alias.format(),
                    Access {
                        stage: vk::PipelineStageFlags2::TRANSFER,
                        mask: vk::AccessFlags2::TRANSFER_WRITE,
                        layout: vk::ImageLayout::TRANSFER_DST_OPTIMAL,
                    },
                    Access {
                        stage: vk::PipelineStageFlags2::TRANSFER,
                        mask: vk::AccessFlags2::TRANSFER_READ,
                        layout: vk::ImageLayout::TRANSFER_SRC_OPTIMAL,
                    },
                    None,
                );

                let to_barrier = image_barrier(
                    texture.raw(),
                    texture.format(),
                    Access {
                        stage: vk::PipelineStageFlags2::NONE,
                        mask: vk::AccessFlags2::empty(),
                        // TODO: this is somewhat sub-optimal, but I think
                        // perfectly legal to not-care about the layout of a
                        // texture? In this case, the entire texture is
                        // overwritten, so the layout doesn't matter, and else
                        // the layout will be `GENERAL` because the texture was
                        // previously written to.
                        layout: if alias.size() == texture.size() {
                            vk::ImageLayout::UNDEFINED
                        } else {
                            vk::ImageLayout::GENERAL
                        },
                    },
                    Access {
                        stage: vk::PipelineStageFlags2::TRANSFER,
                        mask: vk::AccessFlags2::TRANSFER_WRITE,
                        layout: vk::ImageLayout::TRANSFER_DST_OPTIMAL,
                    },
                    None,
                );

                unsafe {
                    ctx.device.dev().cmd_pipeline_barrier2(
                        ctx.cmd.buffer(),
                        &vk::DependencyInfo::default()
                            .image_memory_barriers(&[from_barrier, to_barrier]),
                    );
                }
                ctx.cmd.copy_images(
                    alias.raw(),
                    vk::ImageLayout::TRANSFER_SRC_OPTIMAL,
                    texture.raw(),
                    vk::ImageLayout::TRANSFER_DST_OPTIMAL,
                    &[vk::ImageCopy {
                        src_subresource: vk::ImageSubresourceLayers::default()
                            .aspect_mask(vk::ImageAspectFlags::COLOR)
                            .base_array_layer(0)
                            .mip_level(0)
                            .layer_count(1),
                        src_offset: vk::Offset3D { x: 0, y: 0, z: 0 },
                        dst_subresource: vk::ImageSubresourceLayers::default()
                            .aspect_mask(vk::ImageAspectFlags::COLOR)
                            .base_array_layer(0)
                            .mip_level(0)
                            .layer_count(1),
                        dst_offset: vk::Offset3D {
                            x: rect.top_left().x,
                            y: rect.top_left().y,
                            z: 0,
                        },
                        extent: alias.size(),
                    }],
                );

                let image: Barrier = image_barrier(
                    texture.raw(),
                    texture.format(),
                    Access {
                        stage: vk::PipelineStageFlags2::TRANSFER,
                        mask: vk::AccessFlags2::TRANSFER_WRITE,
                        layout: vk::ImageLayout::TRANSFER_DST_OPTIMAL,
                    },
                    Access {
                        stage: vk::PipelineStageFlags2::ALL_COMMANDS,
                        mask: vk::AccessFlags2::empty(),
                        layout: vk::ImageLayout::GENERAL,
                    },
                    None,
                )
                .into();

                unsafe {
                    ctx.device
                        .dev()
                        .cmd_pipeline_barrier2(ctx.cmd.buffer(), &((&image).into()));
                }
            }
            Ok(())
        }
    });

    rg.add_pass(PassDesc {
        reads: [
            (
                staging_buffer,
                Access {
                    stage: vk::PipelineStageFlags2::TRANSFER,
                    mask: vk::AccessFlags2::TRANSFER_READ,
                    ..Access::undefined()
                },
            ),
            (staging_image, Access::undefined()),
        ]
        .to_vec(),
        writes: textures
            .values()
            .map(|id| {
                (
                    *id,
                    Access {
                        layout: vk::ImageLayout::GENERAL,
                        ..Access::undefined()
                    },
                )
            })
            .collect(),
        record: Some(record),
    });
    Ok(())
}

// fn egui_pass()
#[allow(clippy::too_many_arguments)]
pub fn egui_pass(
    dev: &device::Device,
    rg: &mut RenderGraph,
    texture_handler: &mut crate::texture::TextureManager,
    samplers: &mut crate::SamplerCache,
    egui_state: &mut EguiState,
    egui: &egui::Context,
    output: egui::FullOutput,
    target: GraphResourceId,
) -> crate::Result<Vec<texture::TextureId>> {
    let draw_data = egui.tessellate(output.shapes, output.pixels_per_point);

    #[repr(C)]
    #[derive(Debug, Clone, Copy, bytemuck::Pod, bytemuck::Zeroable)]
    struct Vertex {
        pos: glam::Vec2,
        uv: glam::Vec2,
        color: egui::epaint::Color32,
    }

    #[repr(transparent)]
    #[derive(Debug, Clone, Copy)]
    struct DrawCall(vk::DrawIndexedIndirectCommand);
    unsafe impl bytemuck::Zeroable for DrawCall {}
    unsafe impl bytemuck::Pod for DrawCall {}

    let mut vertices = Vec::new();
    let mut indices = Vec::new();
    let mut draw_calls = Vec::new();
    let mut textures = IndexMap::new();
    let mut textures_indices = Vec::new();
    for draw in draw_data {
        let egui::epaint::Primitive::Mesh(mesh) = draw.primitive else {
            continue;
        };

        draw_calls.push(DrawCall(vk::DrawIndexedIndirectCommand {
            index_count: mesh.indices.len() as u32,
            instance_count: 1,
            first_index: indices.len() as u32,
            vertex_offset: vertices.len() as i32,
            first_instance: 0,
        }));

        vertices.extend(mesh.vertices.iter().map(|v| Vertex {
            pos: glam::vec2(v.pos.x, v.pos.y),
            uv: glam::vec2(v.uv.x, v.uv.y),
            color: v.color,
        }));

        indices.extend(mesh.indices);
        let texture = egui_state.textures.get(&mesh.texture_id).cloned().unwrap();
        if !textures.contains_key(&texture.id) {
            textures.insert(texture.id, texture);
        }
        let idx = textures.get_index_of(&texture.id).unwrap();
        textures_indices.push(idx as u32);
    }

    let num_draw_calls = draw_calls.len();
    let vertices_size = vertices.len() * size_of::<Vertex>();
    let indices_size = indices.len() * size_of::<u32>();
    let draw_calls_size = draw_calls.len() * size_of::<vk::DrawIndexedIndirectCommand>();

    let (draw_staging, vertices, indices, draw_calls, texture_ids) = {
        let staging_size = vertices_size + indices_size + draw_calls_size;

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

        {
            let map = staging.map_mut().unwrap();

            let (st_vertices, rest) = map.split_at_mut(vertices_size);
            let (st_indices, st_drawcalls) = rest.split_at_mut(indices_size);
            st_vertices.copy_from_slice(bytemuck::cast_slice(&vertices));
            st_indices.copy_from_slice(bytemuck::cast_slice(&indices));
            st_drawcalls.copy_from_slice(bytemuck::cast_slice(&draw_calls));
        }
        let staging = rg.import_buffer(Arc::new(staging), Access::undefined());

        let vertices = rg.add_resource(GraphResourceDesc::Buffer(BufferDesc {
            name: Some("egui-draw-vertices".into()),
            size: vertices_size as u64,
            usage: vk::BufferUsageFlags::TRANSFER_DST | vk::BufferUsageFlags::VERTEX_BUFFER,
            mem_location: MemoryLocation::GpuOnly,
            ..Default::default()
        }));
        let indices = rg.add_resource(GraphResourceDesc::Buffer(BufferDesc {
            name: Some("egui-draw-indices".into()),
            size: indices_size as u64,
            usage: vk::BufferUsageFlags::TRANSFER_DST | vk::BufferUsageFlags::INDEX_BUFFER,
            mem_location: MemoryLocation::GpuOnly,
            ..Default::default()
        }));
        let draw_calls = rg.add_resource(GraphResourceDesc::Buffer(BufferDesc {
            name: Some("egui-draw-draw_calls".into()),
            size: draw_calls_size as u64,
            usage: vk::BufferUsageFlags::TRANSFER_DST | vk::BufferUsageFlags::INDIRECT_BUFFER,
            mem_location: MemoryLocation::GpuOnly,
            ..Default::default()
        }));

        let mut texture_ids = Buffer::new(
            dev.clone(),
            BufferDesc {
                name: Some("egui-draw-texture_ids".into()),
                size: (textures_indices.len() * size_of::<u32>()) as u64,
                usage: vk::BufferUsageFlags::STORAGE_BUFFER,
                mem_location: MemoryLocation::CpuToGpu,
                ..Default::default()
            },
        )?;
        {
            let map = texture_ids
                .map_mut()
                .expect("texture id buffer should be device visible");
            map.copy_from_slice(bytemuck::cast_slice(&textures_indices));
        }

        (staging, vertices, indices, draw_calls, texture_ids)
    };

    let descriptor_infos = textures
        .values()
        .map(|entry| {
            let texture = texture_handler.get_texture(entry.id).unwrap();
            vk::DescriptorImageInfo {
                sampler: samplers.get_sampler(entry.as_sampler_desc()).unwrap(),
                image_view: texture
                    .create_view(ImageViewDesc {
                        kind: vk::ImageViewType::TYPE_2D,
                        format: texture.format(),
                        aspect: vk::ImageAspectFlags::COLOR,
                        mip_range: (0..1).into(),
                        layer_range: (0..1).into(),
                        ..Default::default()
                    })
                    .unwrap()
                    .raw(),
                image_layout: vk::ImageLayout::SHADER_READ_ONLY_OPTIMAL,
            }
        })
        .collect::<Vec<_>>();

    let uniform_info = vk::DescriptorBufferInfo {
        buffer: texture_ids.raw(),
        offset: 0,
        range: texture_ids.len(),
    };

    let descriptor_writes = descriptor_infos
        .iter()
        .enumerate()
        .map(|(i, info)| {
            vk::WriteDescriptorSet::default()
                .image_info(core::slice::from_ref(info))
                .descriptor_count(1)
                .descriptor_type(vk::DescriptorType::COMBINED_IMAGE_SAMPLER)
                .dst_set(egui_state.descriptor_set)
                .dst_binding(EguiState::TEXTURE_BINDING)
                .dst_array_element(i as u32)
        })
        .chain(core::iter::once({
            vk::WriteDescriptorSet::default()
                .buffer_info(core::slice::from_ref(&uniform_info))
                .descriptor_count(1)
                .dst_binding(EguiState::UNIFORM_BINDING)
                .descriptor_type(vk::DescriptorType::STORAGE_BUFFER)
                .dst_array_element(0)
                .dst_set(egui_state.descriptor_set)
        }))
        .collect::<Vec<_>>();

    unsafe {
        dev.dev().update_descriptor_sets(&descriptor_writes, &[]);
    }

    let to_remove_tex_ids = output
        .textures_delta
        .free
        .iter()
        .filter_map(|id| egui_state.textures.get(id).cloned())
        .map(|entry| entry.id)
        .collect::<Vec<_>>();

    let reads = textures
        .keys()
        .filter_map(|id| {
            texture_handler
                .get_texture(*id)
                .map(|img| (rg.import_image(img, Access::general()), Access::general()))
        })
        .chain([(target, Access::color_attachment_read_write())])
        .collect::<Vec<_>>();
    let writes = [(target, Access::color_attachment_write_only())].to_vec();

    let record: Box<RecordFn> = Box::new({
        let pipeline = egui_state.pipeline.clone();
        let pipeline_layout = egui_state.pipeline_layout.clone();
        let descriptor_set = egui_state.descriptor_set;
        let screen_rect = egui.screen_rect();

        move |ctx: &RenderContext| -> crate::Result<()> {
            let cmd = &ctx.cmd;
            let staging = ctx.get_buffer(draw_staging).unwrap();
            let vertices = ctx.get_buffer(vertices).unwrap();
            let indices = ctx.get_buffer(indices).unwrap();
            let draw_calls = ctx.get_buffer(draw_calls).unwrap();
            let target = ctx.get_image(target).unwrap();

            cmd.copy_buffers(
                staging.raw(),
                vertices.raw(),
                &[vk::BufferCopy {
                    src_offset: 0,
                    dst_offset: 0,
                    size: vertices_size as u64,
                }],
            );
            cmd.copy_buffers(
                staging.raw(),
                indices.raw(),
                &[vk::BufferCopy {
                    src_offset: vertices_size as u64,
                    dst_offset: 0,
                    size: indices_size as u64,
                }],
            );
            cmd.copy_buffers(
                staging.raw(),
                draw_calls.raw(),
                &[vk::BufferCopy {
                    src_offset: (vertices_size + indices_size) as u64,
                    dst_offset: 0,
                    size: draw_calls_size as u64,
                }],
            );

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

            let color_attachment = &vk::RenderingAttachmentInfo::default()
                .image_layout(vk::ImageLayout::COLOR_ATTACHMENT_OPTIMAL)
                .image_view(
                    target
                        .create_view(ImageViewDesc {
                            kind: vk::ImageViewType::TYPE_2D,
                            format: target.format(),
                            aspect: vk::ImageAspectFlags::COLOR,
                            ..Default::default()
                        })?
                        .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(&[vertices.raw()], &[0]);
            cmd.push_constants(
                &pipeline_layout,
                vk::ShaderStageFlags::VERTEX,
                0,
                bytemuck::cast_slice(
                    &[screen_rect.width(), screen_rect.height()].map(|f| f.to_bits()),
                ),
            );
            cmd.bind_descriptor_sets(
                &pipeline_layout,
                vk::PipelineBindPoint::GRAPHICS,
                &[descriptor_set],
            );
            cmd.draw_indexed_indirect(
                draw_calls.raw(),
                0,
                num_draw_calls as u32,
                size_of::<vk::DrawIndexedIndirectCommand>() as u32,
            );

            cmd.end_rendering();

            Ok(())
        }
    });

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

    Ok(to_remove_tex_ids)
}
