use std::{
    collections::{BTreeMap, BTreeSet},
    marker::PhantomData,
    mem::{MaybeUninit, size_of},
    ptr::NonNull,
    sync::Arc,
};

use ash::vk;
use ash_formats::FormatExt;

use crate::{
    commands::CommandBuffer,
    images::{ImageDesc, ImageViewDesc},
    render_graph::{
        commands::{Command, InsideRenderPass, OutsideRenderPass},
        resources::{ResourceAccess, ResourceId, ResourceRange},
    },
};

use super::graph_builder::GraphResourceRegistry;

pub struct CommandRecorder {
    cmd_buffer: CommandBuffer,
    registry: Arc<GraphResourceRegistry>,
    layouts: BTreeMap<ResourceId, vk::ImageLayout>,
}

impl CommandRecorder {
    pub(crate) fn new(cmd_buffer: CommandBuffer, registry: Arc<GraphResourceRegistry>) -> Self {
        Self {
            cmd_buffer,
            registry,
            layouts: BTreeMap::new(),
        }
    }

    pub fn cmd_buffer(&self) -> &CommandBuffer {
        &self.cmd_buffer
    }

    pub fn into_command_buffer(self) -> CommandBuffer {
        self.cmd_buffer
    }

    pub fn get_buffer_handle(&self, id: ResourceId) -> vk::Buffer {
        if let Some(buffer) = self.registry.buffers.read().get(&id) {
            unsafe { (&*buffer.get()).raw() }
        } else {
            self.registry
                .shared_buffers
                .read()
                .get(&id)
                .expect("unknown graph buffer")
                .raw()
        }
    }

    pub fn get_image_desc(&self, id: ResourceId) -> ImageDesc {
        self.registry
            .images
            .read()
            .get(&id)
            .expect("unknown graph image")
            .desc()
            .clone()
    }

    pub fn get_image_handle(&self, id: ResourceId) -> vk::Image {
        self.registry
            .images
            .read()
            .get(&id)
            .expect("unknown graph image")
            .raw()
    }

    pub fn get_image_layout(&self, id: ResourceId) -> vk::ImageLayout {
        self.layouts
            .get(&id)
            .copied()
            .unwrap_or(vk::ImageLayout::UNDEFINED)
    }

    pub fn get_image_view_handle(&self, id: ResourceId) -> vk::ImageView {
        if let Some(view) = self.registry.views.read().get(&id) {
            return view.raw();
        }

        let mut views = self.registry.views.write();
        if let Some(view) = views.get(&id) {
            return view.raw();
        }

        let image = self
            .registry
            .images
            .read()
            .get(&id)
            .expect("unknown graph image")
            .clone();
        let desc = image.desc();
        let view_desc = ImageViewDesc {
            name: None,
            flags: vk::ImageViewCreateFlags::empty(),
            kind: match desc.kind {
                vk::ImageType::TYPE_1D if desc.array_layers > 1 => vk::ImageViewType::TYPE_1D_ARRAY,
                vk::ImageType::TYPE_1D => vk::ImageViewType::TYPE_1D,
                vk::ImageType::TYPE_2D if desc.array_layers > 1 => vk::ImageViewType::TYPE_2D_ARRAY,
                vk::ImageType::TYPE_2D => vk::ImageViewType::TYPE_2D,
                vk::ImageType::TYPE_3D => vk::ImageViewType::TYPE_3D,
                _ => unreachable!("unsupported image type"),
            },
            format: desc.format,
            aspect: desc.format.aspect_flags(),
            mip_range: desc.mip_level_range(),
            layer_range: desc.layer_range(),
            components: vk::ComponentMapping::default(),
        };
        let view = image
            .create_view(view_desc)
            .expect("failed to create a full-range graph image view");
        let view = std::mem::ManuallyDrop::into_inner(view);
        let handle = view.raw();
        views.insert(id, view);
        handle
    }

    pub(crate) fn set_image_layout(&mut self, id: ResourceId, layout: vk::ImageLayout) {
        self.layouts.insert(id, layout);
    }
}

struct CommandMeta {
    apply_command: unsafe fn(&mut CommandRecorder, NonNull<()>, cursor: &mut usize),
}

enum VecOrSingle<T> {
    Vec(Vec<T>),
    Single(T),
}

impl<T> VecOrSingle<T> {
    fn push(&mut self, value: T)
    where
        T: Default,
    {
        match self {
            VecOrSingle::Vec(vec) => vec.push(value),
            VecOrSingle::Single(existing) => {
                let existing = std::mem::take(existing);
                *self = VecOrSingle::Vec(vec![existing, value]);
            }
        }
    }
    fn iter(&self) -> core::slice::Iter<'_, T> {
        match self {
            VecOrSingle::Vec(items) => items.as_slice().iter(),
            VecOrSingle::Single(item) => core::slice::from_ref(item).iter(),
        }
    }
}

#[derive(Default)]
struct SideEffects {
    map: BTreeMap<(ResourceId, u32), VecOrSingle<ResourceAccess>>,
    resources: BTreeSet<ResourceId>,
}

pub struct SideEffectMap<'a>(&'a mut SideEffects, u32);

impl<'a> SideEffectMap<'a> {
    pub fn insert(&mut self, id: ResourceId, access: ResourceAccess) {
        self.0.resources.insert(id);
        self.0
            .map
            .entry((id, self.1))
            .and_modify(|entry| entry.push(access))
            .or_insert_with(|| VecOrSingle::Single(access));
    }

    pub fn reborrow(&mut self) -> SideEffectMap<'_> {
        SideEffectMap(self.0, self.1)
    }

    pub fn into_side_effect_map2(
        self,
        stage_flags: vk::PipelineStageFlags2,
        access_flags: vk::AccessFlags2,
        required_layout: Option<vk::ImageLayout>,
    ) -> SideEffectMap2<'a> {
        SideEffectMap2 {
            inner: self.0,
            command_index: self.1,
            stage_flags,
            access_flags,
            required_layout,
        }
    }
}

pub struct SideEffectMap2<'a> {
    inner: &'a mut SideEffects,
    command_index: u32,
    stage_flags: vk::PipelineStageFlags2,
    access_flags: vk::AccessFlags2,
    required_layout: Option<vk::ImageLayout>,
}

impl<'a> SideEffectMap2<'a> {
    pub fn reborrow(&mut self) -> SideEffectMap2<'_> {
        SideEffectMap2 {
            inner: self.inner,
            command_index: self.command_index,
            stage_flags: self.stage_flags,
            access_flags: self.access_flags,
            required_layout: self.required_layout,
        }
    }
    pub fn insert_range<T>(&mut self, id: ResourceId, range: T)
    where
        (T, vk::ImageLayout): Into<ResourceRange>,
    {
        self.inner.resources.insert(id);

        SideEffectMap(self.inner, self.command_index).insert(
            id,
            ResourceAccess {
                range: (range, self.required_layout.unwrap_or_default()).into(),
                stages: self.stage_flags,
                access: self.access_flags,
            },
        );
    }
}

pub struct CommandList<'cmd> {
    command_bytes: Vec<MaybeUninit<u8>>,
    cursor: usize,
    // each command that accesses a resource adds an entry to this map (id, command_index) -> access
    // during command recording, we fold accesses to the same resource as commands are recorded, and when a command reads a resource, we check for any previous writes to be made available/visible
    side_effects: SideEffects,
    num_commands: u32,

    _pd: PhantomData<&'cmd ()>,
}

#[must_use]
pub struct RenderPass<'cmd, 'a> {
    cmd_list: &'a mut CommandList<'cmd>,
}

impl<'cmd, 'a> RenderPass<'cmd, 'a> {
    pub fn finalise(self) {
        self.cmd_list.push_inner(super::commands::EndRendering);
    }

    pub fn push<C: Command + InsideRenderPass + 'cmd>(&mut self, command: C) {
        self.cmd_list.push_inner(command);
    }
}

impl<'cmd> CommandList<'cmd> {
    pub fn new() -> Self {
        Self {
            command_bytes: Vec::new(),
            cursor: 0,
            side_effects: SideEffects::default(),
            num_commands: 0,
            _pd: PhantomData,
        }
    }

    pub fn begin_rendering(
        &mut self,
        color_attachments: Vec<super::commands::RenderingAttachment>,
        depth_attachment: Option<super::commands::RenderingAttachment>,
        stencil_attachment: Option<super::commands::RenderingAttachment>,
        area: (u32, u32),
        layers: u32,
    ) -> RenderPass<'cmd, '_> {
        self.push(super::commands::BeginRendering {
            color_attachments,
            depth_attachment,
            stencil_attachment,
            area,
            layers,
        });
        RenderPass { cmd_list: self }
    }

    pub fn push<C: Command + OutsideRenderPass + 'cmd>(&mut self, command: C) {
        self.push_inner(command);
    }

    fn push_inner<C: Command + 'cmd>(&mut self, command: C) {
        #[repr(C, packed)]
        struct Packed<C> {
            meta: CommandMeta,
            command: C,
        }

        command.side_effects(SideEffectMap(&mut self.side_effects, self.num_commands));
        self.num_commands += 1;

        let packed = Packed {
            meta: CommandMeta {
                apply_command: |recorder, command_ptr, cursor| {
                    *cursor += size_of::<CommandMeta>() + size_of::<C>();
                    let command = unsafe { command_ptr.cast::<C>().read_unaligned() };
                    C::apply(command, recorder);
                },
            },
            command,
        };

        let len = self.command_bytes.len();
        self.command_bytes.reserve(size_of::<Packed<C>>());
        unsafe {
            let bytes = self.command_bytes.as_mut_ptr().add(len);
            bytes.cast::<Packed<C>>().write_unaligned(packed);
            self.command_bytes.set_len(len + size_of::<Packed<C>>());
        }
    }

    fn fold_side_effects(
        &self,
        recorder: &mut CommandRecorder,
    ) -> Vec<(u32, super::commands::PipelineBarrier)> {
        // Buffer accesses can be granularly tracked by their offset and size;
        // Image accesses are tracked by their subresource range, which means all regions of the same mip/array layer are considered the same subresource for the purposes of ImageMemoryBarriers

        // But, for simplicity sake, and since buffers can already be created
        // with specific sizes from device memory ranges, and aliased as
        // required, we can just fold all buffer accesses together and collapse
        // the size and offsets to min([offset, ..]), max([offset + size, ..])
        // for all accesses to the same buffer.

        // Generally, we want to be as granular as possible with barriers when
        // it makes sense, but reduce the number of calls to
        // vkCmdPipelineBarrier.

        // Image accesses cannot be folded if they differ in: layout, access type, access stage/mask.

        struct SubresourceState {
            command_index: u32,
        }

        struct Barrier {
            resource_id: ResourceId,
            command_index: u32,
            from: ResourceAccess,
            to: ResourceAccess,
        }

        // map to keep track of the current state of each subresource
        let mut subresource_state: BTreeMap<ResourceId, Vec<(ResourceAccess, SubresourceState)>> =
            Default::default();

        let mut barriers = Vec::new();

        for id in &self.side_effects.resources {
            // check if this subresource has been accessed before, and if we can
            // fold this access with the previous one
            let entry = subresource_state.entry(*id).or_default();

            for (command_index, accesses) in self
                .side_effects
                .map
                .range((*id, 0)..=(*id, u32::MAX))
                .map(|((_id, command_index), accesses)| (command_index, accesses))
            {
                for access in accesses.iter() {
                    let mut access = *access;

                    // iterate over the previous accesses to this subresource:
                    // - if we find an access with the same polarity, attempt to
                    // fold; a barrier already exists if it is needed.
                    // - if we find an access with a different polarity, we need to
                    // insert a barrier before the current command index.
                    // - if we find need for a layout transition, we need to insert
                    // a barrier before the current command index.

                    let mut folded = false;
                    for (_i, (sub, _state)) in &mut entry.iter_mut().enumerate().rev() {
                        if !folded && sub.try_fold(&access) {
                            access = *sub;
                            folded = true;
                        } else if sub.conflicts(&access) {
                            // if this access conflicts with a previous one, we need
                            // to insert a barrier before the current command index.
                            // if we have managed to fold this access with a
                            // previous one, a barrier already exists.

                            // we know that `from` is correct, as the first
                            // conflicting access must be of opposite polarity (or a
                            // layout transition) and no more accesses can be folded
                            // into it.
                            barriers.push(Barrier {
                                resource_id: *id,
                                command_index: *command_index,
                                from: *sub,
                                to: access,
                            });

                            // we can stop iterating over previous accesses, as we
                            // can no longer fold into any accesses before this
                            // conflict.
                            break;
                        }
                    }

                    if !folded {
                        // if we couldn't fold this access with any previous one, add it as a new entry
                        entry.push((
                            access,
                            SubresourceState {
                                command_index: *command_index,
                            },
                        ));
                    }
                }
            }
        }

        let mut result = Vec::with_capacity(barriers.len());
        for barrier in barriers {
            let pipeline_barrier = match (&barrier.from.range, &barrier.to.range) {
                (
                    ResourceRange::Texture {
                        layout: old_layout, ..
                    },
                    ResourceRange::Texture {
                        aspect,
                        mip_level,
                        array_layers,
                        layout: new_layout,
                        ..
                    },
                ) => {
                    let format_aspects = recorder
                        .get_image_desc(barrier.resource_id)
                        .format
                        .aspect_flags();
                    let aspects = *aspect & format_aspects;
                    if aspects.is_empty() {
                        tracing::debug!(resource = ?barrier.resource_id, "skipping image barrier with no compatible aspects");
                        continue;
                    }
                    let image_barrier = vk::ImageMemoryBarrier2::default()
                        .src_stage_mask(barrier.from.stages)
                        .src_access_mask(barrier.from.access)
                        .dst_stage_mask(barrier.to.stages)
                        .dst_access_mask(barrier.to.access)
                        .old_layout(*old_layout)
                        .new_layout(*new_layout)
                        .image(recorder.get_image_handle(barrier.resource_id))
                        .subresource_range(
                            vk::ImageSubresourceRange::default()
                                .aspect_mask(aspects)
                                .base_mip_level(mip_level.start())
                                .level_count(mip_level.len())
                                .base_array_layer(array_layers.start())
                                .layer_count(array_layers.len()),
                        );
                    super::commands::PipelineBarrier {
                        image_barriers: vec![image_barrier],
                        buffer_barriers: Vec::new(),
                        image_layouts: vec![(barrier.resource_id, *new_layout)],
                    }
                }
                (ResourceRange::Buffer { offset, size }, ResourceRange::Buffer { .. }) => {
                    let buffer_barrier = vk::BufferMemoryBarrier2::default()
                        .src_stage_mask(barrier.from.stages)
                        .src_access_mask(barrier.from.access)
                        .dst_stage_mask(barrier.to.stages)
                        .dst_access_mask(barrier.to.access)
                        .buffer(recorder.get_buffer_handle(barrier.resource_id))
                        .offset(*offset)
                        .size(*size);
                    super::commands::PipelineBarrier {
                        image_barriers: Vec::new(),
                        buffer_barriers: vec![buffer_barrier],
                        image_layouts: Vec::new(),
                    }
                }
                _ => continue,
            };
            result.push((barrier.command_index, pipeline_barrier));
        }
        result.sort_by_key(|(command_index, _)| *command_index);
        result
    }

    pub(crate) fn record_commands(&self, recorder: &mut CommandRecorder) {
        let mut cursor = 0;
        let mut barriers = self.fold_side_effects(recorder);
        let mut command_index = 0;

        while cursor < self.command_bytes.len() {
            barriers
                .drain(..barriers.partition_point(|(idx, _)| *idx <= command_index))
                .for_each(|(_, barrier)| {
                    barrier.apply(recorder);
                });

            unsafe {
                let meta_ptr = self
                    .command_bytes
                    .as_ptr()
                    .add(cursor)
                    .cast::<CommandMeta>();
                let meta = meta_ptr.read_unaligned();
                let command_ptr = meta_ptr.add(1).cast::<()>().cast_mut();
                (meta.apply_command)(recorder, NonNull::new_unchecked(command_ptr), &mut cursor);
            }

            command_index += 1;
        }
    }
}