use std::collections::BTreeMap;

use ash::vk;

use crate::resources::{
    BufferRange, ResourceId, Subresource, SubresourceScope, SubresourceScopeInner,
};

/// The position of a command in its command list.
#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Hash)]
pub struct CommandIndex(pub usize);

impl CommandIndex {
    pub fn index(self) -> usize {
        self.0
    }
}

impl From<usize> for CommandIndex {
    fn from(index: usize) -> Self {
        Self(index)
    }
}

impl From<CommandIndex> for usize {
    fn from(index: CommandIndex) -> Self {
        index.0
    }
}

#[derive(Debug, Clone, Copy)]
pub struct ResourceAccess {
    pub resource: ResourceId,
    pub scope: SubresourceScope,
}

/// Stores every declared access together with the command that declared it.
#[derive(Debug, Default, Clone)]
pub struct SideEffectMap {
    accesses: Vec<Vec<ResourceAccess>>,
}

impl SideEffectMap {
    pub fn begin_command(&mut self) -> CommandIndex {
        let index = self.accesses.len();
        self.accesses.push(Vec::new());
        index.into()
    }

    /// Record an access scope for a command and resource.
    ///
    /// Each scope must describe the command's complete use for that access,
    /// not merely an enclosing bound, because `SideEffectMap::insert` may merge
    /// contained scopes for the same command and resource.
    pub fn insert(
        &mut self,
        command_index: CommandIndex,
        resource: ResourceId,
        scope: SubresourceScope,
    ) {
        if self.accesses.len() <= command_index.0 {
            self.accesses.resize_with(command_index.0 + 1, Vec::new);
        }
        let accesses = &mut self.accesses[command_index.0];
        for existing in accesses.iter_mut() {
            if existing.resource == resource {
                if existing.scope.try_merge(&scope) {
                    return;
                }
                if scope.has_same_access_polarity(&existing.scope)
                    && scope.contains_subresource(&existing.scope)
                {
                    let mut merged = scope;
                    merged.stage |= existing.scope.stage;
                    merged.access |= existing.scope.access;
                    existing.scope = merged;
                    return;
                }
            }
        }
        accesses.push(ResourceAccess { resource, scope });
    }

    pub fn accesses(&self, command_index: CommandIndex) -> &[ResourceAccess] {
        self.accesses
            .get(command_index.0)
            .map_or(&[], Vec::as_slice)
    }

    pub fn command_accesses(&self) -> &[Vec<ResourceAccess>] {
        &self.accesses
    }
}

/// A safe, type-erased temporary command buffer. Each command is owned until
/// compilation/recording, while its access declarations are stored separately.
pub struct CommandList {
    pub(crate) commands: Vec<Box<dyn RecordedCommand>>,
    pub side_effects: SideEffectMap,
}

pub trait RecordedCommand {
    fn apply(&mut self, recorder: &mut CommandRecorder<'_>) -> Result<(), RecordError>;
}

impl CommandList {
    pub fn new() -> Self {
        Self {
            commands: Vec::new(),
            side_effects: SideEffectMap::default(),
        }
    }

    pub fn push<C: crate::commands::Command + 'static>(&mut self, command: C) -> CommandIndex {
        let index = self.side_effects.begin_command();
        command.side_effects(index, &mut self.side_effects);
        self.commands.push(Box::new(command));
        index
    }

    pub fn len(&self) -> usize {
        self.commands.len()
    }

    pub fn is_empty(&self) -> bool {
        self.commands.is_empty()
    }
}

impl Default for CommandList {
    fn default() -> Self {
        Self::new()
    }
}

#[derive(Debug, Clone, PartialEq, Eq)]
pub enum RecordError {
    MissingBuffer(ResourceId),
    MissingImage(ResourceId),
    UnresolvedCreatedBuffer(ResourceId),
    UnresolvedCreatedImage(ResourceId),
}

/// Vulkan recording context. Handle maps are populated by the compiled graph.
pub struct CommandRecorder<'a> {
    pub device: &'a ash::Device,
    pub command_buffer: vk::CommandBuffer,
    pub buffers: BTreeMap<ResourceId, vk::Buffer>,
    pub images: BTreeMap<ResourceId, vk::Image>,
}

impl<'a> CommandRecorder<'a> {
    pub fn new(device: &'a ash::Device, command_buffer: vk::CommandBuffer) -> Self {
        Self {
            device,
            command_buffer,
            buffers: BTreeMap::new(),
            images: BTreeMap::new(),
        }
    }

    pub fn buffer(&self, id: ResourceId) -> Result<vk::Buffer, RecordError> {
        self.buffers
            .get(&id)
            .copied()
            .ok_or(RecordError::MissingBuffer(id))
    }

    pub fn image(&self, id: ResourceId) -> Result<vk::Image, RecordError> {
        self.images
            .get(&id)
            .copied()
            .ok_or(RecordError::MissingImage(id))
    }

    pub fn pipeline_barrier(&self, barrier: &super::graph::PipelineBarrier) {
        let buffer_barriers: Vec<_> = barrier
            .resources
            .iter()
            .filter_map(|resource| match resource.range {
                Subresource::Buffer(BufferRange { offset, size }) => {
                    self.buffers.get(&resource.id).map(|&buffer| {
                        let (src, dst) = if let Some(queue_transfer) = resource.queue_transfer {
                            (queue_transfer.src.index(), queue_transfer.dst.index())
                        } else {
                            (vk::QUEUE_FAMILY_IGNORED, vk::QUEUE_FAMILY_IGNORED)
                        };

                        vk::BufferMemoryBarrier2::default()
                            .src_stage_mask(resource.src_stage)
                            .src_access_mask(resource.src_access)
                            .dst_stage_mask(resource.dst_stage)
                            .dst_access_mask(resource.dst_access)
                            .src_queue_family_index(src)
                            .dst_queue_family_index(dst)
                            .buffer(buffer)
                            .offset(offset)
                            .size(size)
                    })
                }
                Subresource::Texture(_) => None,
            })
            .collect();

        let image_barriers: Vec<_> = barrier
            .resources
            .iter()
            .filter_map(|resource| match resource.range {
                Subresource::Texture(range) => self.images.get(&resource.id).map(|&image| {
                    vk::ImageMemoryBarrier2::default()
                        .src_stage_mask(resource.src_stage)
                        .src_access_mask(resource.src_access)
                        .dst_stage_mask(resource.dst_stage)
                        .dst_access_mask(resource.dst_access)
                        .old_layout(resource.old_layout.unwrap_or(vk::ImageLayout::UNDEFINED))
                        .new_layout(resource.new_layout.unwrap_or(vk::ImageLayout::GENERAL))
                        .image(image)
                        .subresource_range(range)
                }),
                Subresource::Buffer(_) => None,
            })
            .collect();

        let dependency = vk::DependencyInfo::default()
            .buffer_memory_barriers(&buffer_barriers)
            .image_memory_barriers(&image_barriers);
        unsafe {
            self.device
                .cmd_pipeline_barrier2(self.command_buffer, &dependency);
        }
    }
}

pub(crate) fn scope_range(scope: SubresourceScope) -> Subresource {
    match scope.subresource {
        SubresourceScopeInner::Buffer(range) => Subresource::Buffer(range),
        SubresourceScopeInner::Texture { range, .. } => Subresource::Texture(range),
    }
}

#[cfg(test)]
mod tests {
    use super::*;

    fn read_scope(
        offset: u64,
        size: u64,
        stage: vk::PipelineStageFlags2,
        access: vk::AccessFlags2,
    ) -> SubresourceScope {
        SubresourceScope {
            access,
            stage,
            subresource: SubresourceScopeInner::Buffer(BufferRange { offset, size }),
        }
    }

    #[test]
    fn side_effect_map_folds_same_command_scopes_in_both_containment_directions() {
        let mut map = SideEffectMap::default();
        let command = map.begin_command();
        map.insert(
            command,
            ResourceId(0),
            read_scope(
                0,
                64,
                vk::PipelineStageFlags2::VERTEX_SHADER,
                vk::AccessFlags2::SHADER_READ,
            ),
        );
        map.insert(
            command,
            ResourceId(0),
            read_scope(
                16,
                16,
                vk::PipelineStageFlags2::FRAGMENT_SHADER,
                vk::AccessFlags2::UNIFORM_READ,
            ),
        );

        let accesses = map.accesses(command);
        assert_eq!(accesses.len(), 1);
        assert_eq!(
            accesses[0].scope.stage,
            vk::PipelineStageFlags2::VERTEX_SHADER | vk::PipelineStageFlags2::FRAGMENT_SHADER
        );
        assert_eq!(
            accesses[0].scope.access,
            vk::AccessFlags2::SHADER_READ | vk::AccessFlags2::UNIFORM_READ
        );
        assert!(matches!(
            accesses[0].scope.subresource,
            SubresourceScopeInner::Buffer(BufferRange {
                offset: 0,
                size: 64
            })
        ));

        let mut reverse = SideEffectMap::default();
        let command = reverse.begin_command();
        reverse.insert(
            command,
            ResourceId(1),
            read_scope(
                16,
                16,
                vk::PipelineStageFlags2::VERTEX_SHADER,
                vk::AccessFlags2::SHADER_READ,
            ),
        );
        reverse.insert(
            command,
            ResourceId(1),
            read_scope(
                0,
                64,
                vk::PipelineStageFlags2::FRAGMENT_SHADER,
                vk::AccessFlags2::UNIFORM_READ,
            ),
        );

        let accesses = reverse.accesses(command);
        assert_eq!(accesses.len(), 1);
        assert_eq!(
            accesses[0].scope.stage,
            vk::PipelineStageFlags2::VERTEX_SHADER | vk::PipelineStageFlags2::FRAGMENT_SHADER
        );
        assert_eq!(
            accesses[0].scope.access,
            vk::AccessFlags2::SHADER_READ | vk::AccessFlags2::UNIFORM_READ
        );
        assert!(matches!(
            accesses[0].scope.subresource,
            SubresourceScopeInner::Buffer(BufferRange {
                offset: 0,
                size: 64
            })
        ));
    }

    #[test]
    fn side_effect_map_keeps_overlapping_non_subset_scopes_separate() {
        let mut map = SideEffectMap::default();
        let command = map.begin_command();
        map.insert(
            command,
            ResourceId(0),
            read_scope(
                0,
                16,
                vk::PipelineStageFlags2::VERTEX_SHADER,
                vk::AccessFlags2::SHADER_READ,
            ),
        );
        map.insert(
            command,
            ResourceId(0),
            read_scope(
                8,
                16,
                vk::PipelineStageFlags2::FRAGMENT_SHADER,
                vk::AccessFlags2::UNIFORM_READ,
            ),
        );

        assert_eq!(map.accesses(command).len(), 2);
    }

    #[test]
    fn side_effect_map_does_not_fold_opposite_polarity_scopes() {
        let mut map = SideEffectMap::default();
        let command = map.begin_command();
        map.insert(
            command,
            ResourceId(0),
            read_scope(
                0,
                64,
                vk::PipelineStageFlags2::COMPUTE_SHADER,
                vk::AccessFlags2::SHADER_READ,
            ),
        );
        map.insert(
            command,
            ResourceId(0),
            read_scope(
                16,
                16,
                vk::PipelineStageFlags2::COMPUTE_SHADER,
                vk::AccessFlags2::SHADER_WRITE,
            ),
        );

        let accesses = map.accesses(command);
        assert_eq!(accesses.len(), 2);
        assert!(matches!(
            accesses[0].scope.subresource,
            SubresourceScopeInner::Buffer(BufferRange {
                offset: 0,
                size: 64
            })
        ));
        assert!(matches!(
            accesses[1].scope.subresource,
            SubresourceScopeInner::Buffer(BufferRange {
                offset: 16,
                size: 16
            })
        ));
    }
}