// Copyright (c) 2016 The vulkano developers
// Licensed under the Apache License, Version 2.0
// <LICENSE-APACHE or
// https://www.apache.org/licenses/LICENSE-2.0> or the MIT
// license <LICENSE-MIT or https://opensource.org/licenses/MIT>,
// at your option. All files in the project carrying such
// notice may not be copied, modified, or distributed except
// according to those terms.

//! A program that is run on the device.
//!
//! In Vulkan, shaders are grouped in *shader modules*. Each shader module is built from SPIR-V
//! code and can contain one or more entry points. Note that for the moment the official
//! GLSL-to-SPIR-V compiler does not support multiple entry points.
//!
//! The vulkano library can parse and introspect SPIR-V code, but it does not fully validate the
//! code. You are encouraged to use the `vulkano-shaders` crate that will generate Rust code that
//! wraps around vulkano's shaders API.

use crate::{
    descriptor_set::layout::DescriptorType,
    device::{Device, DeviceOwned},
    format::{Format, NumericType},
    image::view::ImageViewType,
    macros::{vulkan_bitflags, vulkan_enum},
    pipeline::{graphics::input_assembly::PrimitiveTopology, layout::PushConstantRange},
    shader::spirv::{Capability, Spirv, SpirvError},
    sync::PipelineStages,
    DeviceSize, OomError, Version, VulkanError, VulkanObject,
};
use ahash::{HashMap, HashSet};
use std::{
    borrow::Cow,
    error::Error,
    ffi::{CStr, CString},
    fmt::{Display, Error as FmtError, Formatter},
    mem,
    mem::MaybeUninit,
    num::NonZeroU64,
    ptr,
    sync::Arc,
};

pub mod reflect;
pub mod spirv;

use spirv::ExecutionModel;

// Generated by build.rs
include!(concat!(env!("OUT_DIR"), "/spirv_reqs.rs"));

/// Contains SPIR-V code with one or more entry points.
#[derive(Debug)]
pub struct ShaderModule {
    handle: ash::vk::ShaderModule,
    device: Arc<Device>,
    id: NonZeroU64,
    entry_points: HashMap<String, HashMap<ExecutionModel, EntryPointInfo>>,
}

impl ShaderModule {
    /// Builds a new shader module from SPIR-V 32-bit words. The shader code is parsed and the
    /// necessary information is extracted from it.
    ///
    /// # Safety
    ///
    /// - The SPIR-V code is not validated beyond the minimum needed to extract the information.
    #[inline]
    pub unsafe fn from_words(
        device: Arc<Device>,
        words: &[u32],
    ) -> Result<Arc<ShaderModule>, ShaderCreationError> {
        let spirv = Spirv::new(words)?;

        Self::from_words_with_data(
            device,
            words,
            spirv.version(),
            reflect::spirv_capabilities(&spirv),
            reflect::spirv_extensions(&spirv),
            reflect::entry_points(&spirv),
        )
    }

    /// As `from_words`, but takes a slice of bytes.
    ///
    /// # Panics
    ///
    /// - Panics if the length of `bytes` is not a multiple of 4.
    #[inline]
    pub unsafe fn from_bytes(
        device: Arc<Device>,
        bytes: &[u8],
    ) -> Result<Arc<ShaderModule>, ShaderCreationError> {
        assert!((bytes.len() % 4) == 0);

        Self::from_words(
            device,
            std::slice::from_raw_parts(
                bytes.as_ptr() as *const _,
                bytes.len() / mem::size_of::<u32>(),
            ),
        )
    }

    /// As `from_words`, but does not parse the code. Instead, you must provide the needed
    /// information yourself. This can be useful if you've already done parsing yourself and
    /// want to prevent Vulkano from doing it a second time.
    ///
    /// # Safety
    ///
    /// - The SPIR-V code is not validated at all.
    /// - The provided information must match what the SPIR-V code contains.
    pub unsafe fn from_words_with_data<'a>(
        device: Arc<Device>,
        words: &[u32],
        spirv_version: Version,
        spirv_capabilities: impl IntoIterator<Item = &'a Capability>,
        spirv_extensions: impl IntoIterator<Item = &'a str>,
        entry_points: impl IntoIterator<Item = (String, ExecutionModel, EntryPointInfo)>,
    ) -> Result<Arc<ShaderModule>, ShaderCreationError> {
        if let Err(reason) = check_spirv_version(&device, spirv_version) {
            return Err(ShaderCreationError::SpirvVersionNotSupported {
                version: spirv_version,
                reason,
            });
        }

        for &capability in spirv_capabilities {
            if let Err(reason) = check_spirv_capability(&device, capability) {
                return Err(ShaderCreationError::SpirvCapabilityNotSupported {
                    capability,
                    reason,
                });
            }
        }

        for extension in spirv_extensions {
            if let Err(reason) = check_spirv_extension(&device, extension) {
                return Err(ShaderCreationError::SpirvExtensionNotSupported {
                    extension: extension.to_owned(),
                    reason,
                });
            }
        }

        let handle = {
            let infos = ash::vk::ShaderModuleCreateInfo {
                flags: ash::vk::ShaderModuleCreateFlags::empty(),
                code_size: words.len() * mem::size_of::<u32>(),
                p_code: words.as_ptr(),
                ..Default::default()
            };

            let fns = device.fns();
            let mut output = MaybeUninit::uninit();
            (fns.v1_0.create_shader_module)(
                device.handle(),
                &infos,
                ptr::null(),
                output.as_mut_ptr(),
            )
            .result()
            .map_err(VulkanError::from)?;
            output.assume_init()
        };

        let entries = entry_points.into_iter().collect::<Vec<_>>();
        let entry_points = entries
            .iter()
            .map(|(name, _, _)| name)
            .collect::<HashSet<_>>()
            .iter()
            .map(|name| {
                (
                    (*name).clone(),
                    entries
                        .iter()
                        .filter_map(|(entry_name, entry_model, info)| {
                            if &entry_name == name {
                                Some((*entry_model, info.clone()))
                            } else {
                                None
                            }
                        })
                        .collect::<HashMap<_, _>>(),
                )
            })
            .collect();

        Ok(Arc::new(ShaderModule {
            handle,
            device,
            id: Self::next_id(),
            entry_points,
        }))
    }

    /// As `from_words_with_data`, but takes a slice of bytes.
    ///
    /// # Panics
    ///
    /// - Panics if the length of `bytes` is not a multiple of 4.
    pub unsafe fn from_bytes_with_data<'a>(
        device: Arc<Device>,
        bytes: &[u8],
        spirv_version: Version,
        spirv_capabilities: impl IntoIterator<Item = &'a Capability>,
        spirv_extensions: impl IntoIterator<Item = &'a str>,
        entry_points: impl IntoIterator<Item = (String, ExecutionModel, EntryPointInfo)>,
    ) -> Result<Arc<ShaderModule>, ShaderCreationError> {
        assert!((bytes.len() % 4) == 0);

        Self::from_words_with_data(
            device,
            std::slice::from_raw_parts(
                bytes.as_ptr() as *const _,
                bytes.len() / mem::size_of::<u32>(),
            ),
            spirv_version,
            spirv_capabilities,
            spirv_extensions,
            entry_points,
        )
    }

    /// Returns information about the entry point with the provided name. Returns `None` if no entry
    /// point with that name exists in the shader module or if multiple entry points with the same
    /// name exist.
    #[inline]
    pub fn entry_point<'a>(&'a self, name: &str) -> Option<EntryPoint<'a>> {
        self.entry_points.get(name).and_then(|infos| {
            if infos.len() == 1 {
                infos.iter().next().map(|(_, info)| EntryPoint {
                    module: self,
                    name: CString::new(name).unwrap(),
                    info,
                })
            } else {
                None
            }
        })
    }

    /// Returns information about the entry point with the provided name and execution model.
    /// Returns `None` if no entry and execution model exists in the shader module.
    #[inline]
    pub fn entry_point_with_execution<'a>(
        &'a self,
        name: &str,
        execution: ExecutionModel,
    ) -> Option<EntryPoint<'a>> {
        self.entry_points.get(name).and_then(|infos| {
            infos.get(&execution).map(|info| EntryPoint {
                module: self,
                name: CString::new(name).unwrap(),
                info,
            })
        })
    }
}

impl Drop for ShaderModule {
    #[inline]
    fn drop(&mut self) {
        unsafe {
            let fns = self.device.fns();
            (fns.v1_0.destroy_shader_module)(self.device.handle(), self.handle, ptr::null());
        }
    }
}

unsafe impl VulkanObject for ShaderModule {
    type Handle = ash::vk::ShaderModule;

    #[inline]
    fn handle(&self) -> Self::Handle {
        self.handle
    }
}

unsafe impl DeviceOwned for ShaderModule {
    #[inline]
    fn device(&self) -> &Arc<Device> {
        &self.device
    }
}

crate::impl_id_counter!(ShaderModule);

/// Error that can happen when creating a new shader module.
#[derive(Clone, Debug)]
pub enum ShaderCreationError {
    OomError(OomError),
    SpirvCapabilityNotSupported {
        capability: Capability,
        reason: ShaderSupportError,
    },
    SpirvError(SpirvError),
    SpirvExtensionNotSupported {
        extension: String,
        reason: ShaderSupportError,
    },
    SpirvVersionNotSupported {
        version: Version,
        reason: ShaderSupportError,
    },
}

impl Error for ShaderCreationError {
    fn source(&self) -> Option<&(dyn Error + 'static)> {
        match self {
            Self::OomError(err) => Some(err),
            Self::SpirvCapabilityNotSupported { reason, .. } => Some(reason),
            Self::SpirvError(err) => Some(err),
            Self::SpirvExtensionNotSupported { reason, .. } => Some(reason),
            Self::SpirvVersionNotSupported { reason, .. } => Some(reason),
        }
    }
}

impl Display for ShaderCreationError {
    fn fmt(&self, f: &mut Formatter<'_>) -> Result<(), FmtError> {
        match self {
            Self::OomError(_) => write!(f, "not enough memory available"),
            Self::SpirvCapabilityNotSupported { capability, .. } => write!(
                f,
                "the SPIR-V capability {:?} enabled by the shader is not supported by the device",
                capability,
            ),
            Self::SpirvError(_) => write!(f, "the SPIR-V module could not be read"),
            Self::SpirvExtensionNotSupported { extension, .. } => write!(
                f,
                "the SPIR-V extension {} enabled by the shader is not supported by the device",
                extension,
            ),
            Self::SpirvVersionNotSupported { version, .. } => write!(
                f,
                "the shader uses SPIR-V version {}.{}, which is not supported by the device",
                version.major, version.minor,
            ),
        }
    }
}

impl From<VulkanError> for ShaderCreationError {
    fn from(err: VulkanError) -> Self {
        Self::OomError(err.into())
    }
}

impl From<SpirvError> for ShaderCreationError {
    fn from(err: SpirvError) -> Self {
        Self::SpirvError(err)
    }
}

/// Error that can happen when checking whether a shader is supported by a device.
#[derive(Clone, Copy, Debug)]
pub enum ShaderSupportError {
    NotSupportedByVulkan,
    RequirementsNotMet(&'static [&'static str]),
}

impl Error for ShaderSupportError {}

impl Display for ShaderSupportError {
    fn fmt(&self, f: &mut Formatter<'_>) -> Result<(), FmtError> {
        match self {
            Self::NotSupportedByVulkan => write!(f, "not supported by Vulkan"),
            Self::RequirementsNotMet(requirements) => write!(
                f,
                "at least one of the following must be available/enabled on the device: {}",
                requirements.join(", "),
            ),
        }
    }
}

/// The information associated with a single entry point in a shader.
#[derive(Clone, Debug)]
pub struct EntryPointInfo {
    pub execution: ShaderExecution,
    pub descriptor_requirements: HashMap<(u32, u32), DescriptorRequirements>,
    pub push_constant_requirements: Option<PushConstantRange>,
    pub specialization_constant_requirements: HashMap<u32, SpecializationConstantRequirements>,
    pub input_interface: ShaderInterface,
    pub output_interface: ShaderInterface,
}

/// Represents a shader entry point in a shader module.
///
/// Can be obtained by calling [`entry_point`](ShaderModule::entry_point) on the shader module.
#[derive(Clone, Debug)]
pub struct EntryPoint<'a> {
    module: &'a ShaderModule,
    name: CString,
    info: &'a EntryPointInfo,
}

impl<'a> EntryPoint<'a> {
    /// Returns the module this entry point comes from.
    #[inline]
    pub fn module(&self) -> &'a ShaderModule {
        self.module
    }

    /// Returns the name of the entry point.
    #[inline]
    pub fn name(&self) -> &CStr {
        &self.name
    }

    /// Returns the execution model of the shader.
    #[inline]
    pub fn execution(&self) -> &ShaderExecution {
        &self.info.execution
    }

    /// Returns the descriptor requirements.
    #[inline]
    pub fn descriptor_requirements(
        &self,
    ) -> impl ExactSizeIterator<Item = ((u32, u32), &DescriptorRequirements)> {
        self.info
            .descriptor_requirements
            .iter()
            .map(|(k, v)| (*k, v))
    }

    /// Returns the push constant requirements.
    #[inline]
    pub fn push_constant_requirements(&self) -> Option<&PushConstantRange> {
        self.info.push_constant_requirements.as_ref()
    }

    /// Returns the specialization constant requirements.
    #[inline]
    pub fn specialization_constant_requirements(
        &self,
    ) -> impl ExactSizeIterator<Item = (u32, &SpecializationConstantRequirements)> {
        self.info
            .specialization_constant_requirements
            .iter()
            .map(|(k, v)| (*k, v))
    }

    /// Returns the input attributes used by the shader stage.
    #[inline]
    pub fn input_interface(&self) -> &ShaderInterface {
        &self.info.input_interface
    }

    /// Returns the output attributes used by the shader stage.
    #[inline]
    pub fn output_interface(&self) -> &ShaderInterface {
        &self.info.output_interface
    }
}

/// The mode in which a shader executes. This includes both information about the shader type/stage,
/// and additional data relevant to particular shader types.
#[derive(Clone, Copy, Debug, Eq, PartialEq)]
pub enum ShaderExecution {
    Vertex,
    TessellationControl,
    TessellationEvaluation,
    Geometry(GeometryShaderExecution),
    Fragment,
    Compute,
    RayGeneration,
    AnyHit,
    ClosestHit,
    Miss,
    Intersection,
    Callable,
}

/*#[derive(Clone, Copy, Debug)]
pub struct TessellationShaderExecution {
    pub num_output_vertices: u32,
    pub point_mode: bool,
    pub subdivision: TessellationShaderSubdivision,
}

#[derive(Clone, Copy, Debug, PartialEq, Eq, Hash)]
pub enum TessellationShaderSubdivision {
    Triangles,
    Quads,
    Isolines,
}*/

/// The mode in which a geometry shader executes.
#[derive(Clone, Copy, Debug, Eq, PartialEq)]
pub struct GeometryShaderExecution {
    pub input: GeometryShaderInput,
    /*pub max_output_vertices: u32,
    pub num_invocations: u32,
    pub output: GeometryShaderOutput,*/
}

/// The input primitive type that is expected by a geometry shader.
#[derive(Clone, Copy, Debug, PartialEq, Eq, Hash)]
pub enum GeometryShaderInput {
    Points,
    Lines,
    LinesWithAdjacency,
    Triangles,
    TrianglesWithAdjacency,
}

impl GeometryShaderInput {
    /// Returns true if the given primitive topology can be used as input for this geometry shader.
    #[inline]
    pub fn is_compatible_with(&self, topology: PrimitiveTopology) -> bool {
        match self {
            Self::Points => matches!(topology, PrimitiveTopology::PointList),
            Self::Lines => matches!(
                topology,
                PrimitiveTopology::LineList | PrimitiveTopology::LineStrip
            ),
            Self::LinesWithAdjacency => matches!(
                topology,
                PrimitiveTopology::LineListWithAdjacency
                    | PrimitiveTopology::LineStripWithAdjacency
            ),
            Self::Triangles => matches!(
                topology,
                PrimitiveTopology::TriangleList
                    | PrimitiveTopology::TriangleStrip
                    | PrimitiveTopology::TriangleFan,
            ),
            Self::TrianglesWithAdjacency => matches!(
                topology,
                PrimitiveTopology::TriangleListWithAdjacency
                    | PrimitiveTopology::TriangleStripWithAdjacency,
            ),
        }
    }
}

/*#[derive(Clone, Copy, Debug, PartialEq, Eq, Hash)]
pub enum GeometryShaderOutput {
    Points,
    LineStrip,
    TriangleStrip,
}*/

/// The requirements imposed by a shader on a descriptor within a descriptor set layout, and on any
/// resource that is bound to that descriptor.
#[derive(Clone, Debug, Default)]
pub struct DescriptorRequirements {
    /// The descriptor types that are allowed.
    pub descriptor_types: Vec<DescriptorType>,

    /// The number of descriptors (array elements) that the shader requires. The descriptor set
    /// layout can declare more than this, but never less.
    ///
    /// `None` means that the shader declares this as a runtime-sized array, and could potentially
    /// access every array element provided in the descriptor set.
    pub descriptor_count: Option<u32>,

    /// The image format that is required for image views bound to this descriptor. If this is
    /// `None`, then any image format is allowed.
    pub image_format: Option<Format>,

    /// Whether image views bound to this descriptor must have multisampling enabled or disabled.
    pub image_multisampled: bool,

    /// The base scalar type required for the format of image views bound to this descriptor.
    /// This is `None` for non-image descriptors.
    pub image_scalar_type: Option<ShaderScalarType>,

    /// The view type that is required for image views bound to this descriptor.
    /// This is `None` for non-image descriptors.
    pub image_view_type: Option<ImageViewType>,

    /// For sampler bindings, the descriptor indices that require a depth comparison sampler.
    pub sampler_compare: HashSet<u32>,

    /// For sampler bindings, the descriptor indices that perform sampling operations that are not
    /// permitted with unnormalized coordinates. This includes sampling with `ImplicitLod`,
    /// `Dref` or `Proj` SPIR-V instructions or with an LOD bias or offset.
    pub sampler_no_unnormalized_coordinates: HashSet<u32>,

    /// For sampler bindings, the descriptor indices that perform sampling operations that are not
    /// permitted with a sampler YCbCr conversion. This includes sampling with `Gather` SPIR-V
    /// instructions or with an offset.
    pub sampler_no_ycbcr_conversion: HashSet<u32>,

    /// For sampler bindings, the sampled image descriptors that are used in combination with each
    /// sampler descriptor index.
    pub sampler_with_images: HashMap<u32, HashSet<DescriptorIdentifier>>,

    /// The shader stages that the descriptor must be declared for.
    pub stages: ShaderStages,

    /// For storage image bindings, the descriptor indices that atomic operations are used with.
    pub storage_image_atomic: HashSet<u32>,

    /// For storage images and storage texel buffers, the descriptor indices that perform read
    /// operations on the bound resource.
    pub storage_read: HashSet<u32>,

    /// For storage buffers, storage images and storage texel buffers, the descriptor indices that
    /// perform write operations on the bound resource.
    pub storage_write: HashSet<u32>,
}

#[derive(Clone, Copy, Debug, PartialEq, Eq, Hash)]
pub struct DescriptorIdentifier {
    pub set: u32,
    pub binding: u32,
    pub index: u32,
}

impl DescriptorRequirements {
    /// Produces the intersection of two descriptor requirements, so that the requirements of both
    /// are satisfied. An error is returned if the requirements conflict.
    #[inline]
    pub fn intersection(&self, other: &Self) -> Result<Self, DescriptorRequirementsIncompatible> {
        let descriptor_types: Vec<_> = self
            .descriptor_types
            .iter()
            .copied()
            .filter(|ty| other.descriptor_types.contains(ty))
            .collect();

        if descriptor_types.is_empty() {
            return Err(DescriptorRequirementsIncompatible::DescriptorType);
        }

        if let (Some(first), Some(second)) = (self.image_format, other.image_format) {
            if first != second {
                return Err(DescriptorRequirementsIncompatible::ImageFormat);
            }
        }

        if let (Some(first), Some(second)) = (self.image_scalar_type, other.image_scalar_type) {
            if first != second {
                return Err(DescriptorRequirementsIncompatible::ImageScalarType);
            }
        }

        if let (Some(first), Some(second)) = (self.image_view_type, other.image_view_type) {
            if first != second {
                return Err(DescriptorRequirementsIncompatible::ImageViewType);
            }
        }

        if self.image_multisampled != other.image_multisampled {
            return Err(DescriptorRequirementsIncompatible::ImageMultisampled);
        }

        let sampler_with_images = {
            let mut result = self.sampler_with_images.clone();

            for (&index, other_identifiers) in &other.sampler_with_images {
                result.entry(index).or_default().extend(other_identifiers);
            }

            result
        };

        Ok(Self {
            descriptor_types,
            descriptor_count: self.descriptor_count.max(other.descriptor_count),
            image_format: self.image_format.or(other.image_format),
            image_multisampled: self.image_multisampled,
            image_scalar_type: self.image_scalar_type.or(other.image_scalar_type),
            image_view_type: self.image_view_type.or(other.image_view_type),
            sampler_compare: &self.sampler_compare | &other.sampler_compare,
            sampler_no_unnormalized_coordinates: &self.sampler_no_unnormalized_coordinates
                | &other.sampler_no_unnormalized_coordinates,
            sampler_no_ycbcr_conversion: &self.sampler_no_ycbcr_conversion
                | &other.sampler_no_ycbcr_conversion,
            sampler_with_images,
            stages: self.stages | other.stages,
            storage_image_atomic: &self.storage_image_atomic | &other.storage_image_atomic,
            storage_read: &self.storage_read | &other.storage_read,
            storage_write: &self.storage_write | &other.storage_write,
        })
    }
}

/// An error that can be returned when trying to create the intersection of two
/// `DescriptorRequirements` values.
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
pub enum DescriptorRequirementsIncompatible {
    /// The allowed descriptor types of the descriptors do not overlap.
    DescriptorType,
    /// The descriptors require different formats.
    ImageFormat,
    /// The descriptors require different scalar types.
    ImageScalarType,
    /// The multisampling requirements of the descriptors differ.
    ImageMultisampled,
    /// The descriptors require different image view types.
    ImageViewType,
}

impl Error for DescriptorRequirementsIncompatible {}

impl Display for DescriptorRequirementsIncompatible {
    fn fmt(&self, f: &mut Formatter<'_>) -> Result<(), FmtError> {
        match self {
            DescriptorRequirementsIncompatible::DescriptorType => write!(
                f,
                "the allowed descriptor types of the two descriptors do not overlap",
            ),
            DescriptorRequirementsIncompatible::ImageFormat => {
                write!(f, "the descriptors require different formats",)
            }
            DescriptorRequirementsIncompatible::ImageMultisampled => write!(
                f,
                "the multisampling requirements of the descriptors differ",
            ),
            DescriptorRequirementsIncompatible::ImageScalarType => {
                write!(f, "the descriptors require different scalar types",)
            }
            DescriptorRequirementsIncompatible::ImageViewType => {
                write!(f, "the descriptors require different image view types",)
            }
        }
    }
}

/// The requirements imposed by a shader on a specialization constant.
#[derive(Clone, Copy, Debug)]
pub struct SpecializationConstantRequirements {
    pub size: DeviceSize,
}

/// Trait for types that contain specialization data for shaders.
///
/// Shader modules can contain what is called *specialization constants*. They are the same as
/// constants except that their values can be defined when you create a compute pipeline or a
/// graphics pipeline. Doing so is done by passing a type that implements the
/// `SpecializationConstants` trait and that stores the values in question. The `descriptors()`
/// method of this trait indicates how to grab them.
///
/// Boolean specialization constants must be stored as 32bits integers, where `0` means `false` and
/// any non-zero value means `true`. Integer and floating-point specialization constants are
/// stored as their Rust equivalent.
///
/// This trait is implemented on `()` for shaders that don't have any specialization constant.
///
/// # Examples
///
/// ```rust
/// use vulkano::shader::SpecializationConstants;
/// use vulkano::shader::SpecializationMapEntry;
///
/// #[repr(C)]      // `#[repr(C)]` guarantees that the struct has a specific layout
/// struct MySpecConstants {
///     my_integer_constant: i32,
///     a_boolean: u32,
///     floating_point: f32,
/// }
///
/// unsafe impl SpecializationConstants for MySpecConstants {
///     fn descriptors() -> &'static [SpecializationMapEntry] {
///         static DESCRIPTORS: [SpecializationMapEntry; 3] = [
///             SpecializationMapEntry {
///                 constant_id: 0,
///                 offset: 0,
///                 size: 4,
///             },
///             SpecializationMapEntry {
///                 constant_id: 1,
///                 offset: 4,
///                 size: 4,
///             },
///             SpecializationMapEntry {
///                 constant_id: 2,
///                 offset: 8,
///                 size: 4,
///             },
///         ];
///
///         &DESCRIPTORS
///     }
/// }
/// ```
///
/// # Safety
///
/// - The `SpecializationMapEntry` returned must contain valid offsets and sizes.
/// - The size of each `SpecializationMapEntry` must match the size of the corresponding constant
///   (`4` for booleans).
pub unsafe trait SpecializationConstants {
    /// Returns descriptors of the struct's layout.
    fn descriptors() -> &'static [SpecializationMapEntry];
}

unsafe impl SpecializationConstants for () {
    #[inline]
    fn descriptors() -> &'static [SpecializationMapEntry] {
        &[]
    }
}

/// Describes an individual constant to set in the shader. Also a field in the struct.
// Implementation note: has the same memory representation as a `VkSpecializationMapEntry`.
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
#[repr(C)]
pub struct SpecializationMapEntry {
    /// Identifier of the constant in the shader that corresponds to this field.
    ///
    /// For SPIR-V, this must be the value of the `SpecId` decoration applied to the specialization
    /// constant.
    /// For GLSL, this must be the value of `N` in the `layout(constant_id = N)` attribute applied
    /// to a constant.
    pub constant_id: u32,

    /// Offset within the struct where the data can be found.
    pub offset: u32,

    /// Size of the data in bytes. Must match the size of the constant (`4` for booleans).
    pub size: usize,
}

impl From<SpecializationMapEntry> for ash::vk::SpecializationMapEntry {
    #[inline]
    fn from(val: SpecializationMapEntry) -> Self {
        Self {
            constant_id: val.constant_id,
            offset: val.offset,
            size: val.size,
        }
    }
}

/// Type that contains the definition of an interface between two shader stages, or between
/// the outside and a shader stage.
#[derive(Clone, Debug)]
pub struct ShaderInterface {
    elements: Vec<ShaderInterfaceEntry>,
}

impl ShaderInterface {
    /// Constructs a new `ShaderInterface`.
    ///
    /// # Safety
    ///
    /// - Must only provide one entry per location.
    /// - The format of each element must not be larger than 128 bits.
    // TODO: 4x64 bit formats are possible, but they require special handling.
    // TODO: could this be made safe?
    #[inline]
    pub unsafe fn new_unchecked(elements: Vec<ShaderInterfaceEntry>) -> ShaderInterface {
        ShaderInterface { elements }
    }

    /// Creates a description of an empty shader interface.
    #[inline]
    pub const fn empty() -> ShaderInterface {
        ShaderInterface {
            elements: Vec::new(),
        }
    }

    /// Returns a slice containing the elements of the interface.
    #[inline]
    pub fn elements(&self) -> &[ShaderInterfaceEntry] {
        self.elements.as_ref()
    }

    /// Checks whether the interface is potentially compatible with another one.
    ///
    /// Returns `Ok` if the two interfaces are compatible.
    #[inline]
    pub fn matches(&self, other: &ShaderInterface) -> Result<(), ShaderInterfaceMismatchError> {
        if self.elements().len() != other.elements().len() {
            return Err(ShaderInterfaceMismatchError::ElementsCountMismatch {
                self_elements: self.elements().len() as u32,
                other_elements: other.elements().len() as u32,
            });
        }

        for a in self.elements() {
            let location_range = a.location..a.location + a.ty.num_locations();
            for loc in location_range {
                let b = match other
                    .elements()
                    .iter()
                    .find(|e| loc >= e.location && loc < e.location + e.ty.num_locations())
                {
                    None => {
                        return Err(ShaderInterfaceMismatchError::MissingElement { location: loc })
                    }
                    Some(b) => b,
                };

                if a.ty != b.ty {
                    return Err(ShaderInterfaceMismatchError::TypeMismatch {
                        location: loc,
                        self_ty: a.ty,
                        other_ty: b.ty,
                    });
                }

                // TODO: enforce this?
                /*match (a.name, b.name) {
                    (Some(ref an), Some(ref bn)) => if an != bn { return false },
                    _ => ()
                };*/
            }
        }

        // Note: since we check that the number of elements is the same, we don't need to iterate
        // over b's elements.

        Ok(())
    }
}

/// Entry of a shader interface definition.
#[derive(Debug, Clone)]
pub struct ShaderInterfaceEntry {
    /// The location slot that the variable starts at.
    pub location: u32,

    /// The component slot that the variable starts at. Must be in the range 0..=3.
    pub component: u32,

    /// Name of the element, or `None` if the name is unknown.
    pub name: Option<Cow<'static, str>>,

    /// The type of the variable.
    pub ty: ShaderInterfaceEntryType,
}

/// The type of a variable in a shader interface.
#[derive(Clone, Copy, Debug, PartialEq, Eq, Hash)]
pub struct ShaderInterfaceEntryType {
    /// The base numeric type.
    pub base_type: ShaderScalarType,

    /// The number of vector components. Must be in the range 1..=4.
    pub num_components: u32,

    /// The number of array elements or matrix columns.
    pub num_elements: u32,

    /// Whether the base type is 64 bits wide. If true, each item of the base type takes up two
    /// component slots instead of one.
    pub is_64bit: bool,
}

impl ShaderInterfaceEntryType {
    pub(crate) fn to_format(&self) -> Format {
        assert!(!self.is_64bit); // TODO: implement
        match self.base_type {
            ShaderScalarType::Float => match self.num_components {
                1 => Format::R32_SFLOAT,
                2 => Format::R32G32_SFLOAT,
                3 => Format::R32G32B32_SFLOAT,
                4 => Format::R32G32B32A32_SFLOAT,
                _ => unreachable!(),
            },
            ShaderScalarType::Sint => match self.num_components {
                1 => Format::R32_SINT,
                2 => Format::R32G32_SINT,
                3 => Format::R32G32B32_SINT,
                4 => Format::R32G32B32A32_SINT,
                _ => unreachable!(),
            },
            ShaderScalarType::Uint => match self.num_components {
                1 => Format::R32_UINT,
                2 => Format::R32G32_UINT,
                3 => Format::R32G32B32_UINT,
                4 => Format::R32G32B32A32_UINT,
                _ => unreachable!(),
            },
        }
    }

    pub(crate) fn num_locations(&self) -> u32 {
        assert!(!self.is_64bit); // TODO: implement
        self.num_elements
    }
}

/// The numeric base type of a shader variable.
#[derive(Clone, Copy, Debug, PartialEq, Eq, Hash)]
pub enum ShaderScalarType {
    Float,
    Sint,
    Uint,
}

// https://registry.khronos.org/vulkan/specs/1.3-extensions/html/chap43.html#formats-numericformat
impl From<NumericType> for ShaderScalarType {
    #[inline]
    fn from(val: NumericType) -> Self {
        match val {
            NumericType::SFLOAT => Self::Float,
            NumericType::UFLOAT => Self::Float,
            NumericType::SINT => Self::Sint,
            NumericType::UINT => Self::Uint,
            NumericType::SNORM => Self::Float,
            NumericType::UNORM => Self::Float,
            NumericType::SSCALED => Self::Float,
            NumericType::USCALED => Self::Float,
            NumericType::SRGB => Self::Float,
        }
    }
}

/// Error that can happen when the interface mismatches between two shader stages.
#[derive(Clone, Debug, PartialEq, Eq)]
pub enum ShaderInterfaceMismatchError {
    /// The number of elements is not the same between the two shader interfaces.
    ElementsCountMismatch {
        /// Number of elements in the first interface.
        self_elements: u32,
        /// Number of elements in the second interface.
        other_elements: u32,
    },

    /// An element is missing from one of the interfaces.
    MissingElement {
        /// Location of the missing element.
        location: u32,
    },

    /// The type of an element does not match.
    TypeMismatch {
        /// Location of the element that mismatches.
        location: u32,
        /// Type in the first interface.
        self_ty: ShaderInterfaceEntryType,
        /// Type in the second interface.
        other_ty: ShaderInterfaceEntryType,
    },
}

impl Error for ShaderInterfaceMismatchError {}

impl Display for ShaderInterfaceMismatchError {
    fn fmt(&self, f: &mut Formatter<'_>) -> Result<(), FmtError> {
        write!(
            f,
            "{}",
            match self {
                ShaderInterfaceMismatchError::ElementsCountMismatch { .. } => {
                    "the number of elements mismatches"
                }
                ShaderInterfaceMismatchError::MissingElement { .. } => "an element is missing",
                ShaderInterfaceMismatchError::TypeMismatch { .. } => {
                    "the type of an element does not match"
                }
            }
        )
    }
}

vulkan_enum! {
    /// A single shader stage.
    #[non_exhaustive]
    ShaderStage = ShaderStageFlags(u32);

    // TODO: document
    Vertex = VERTEX,

    // TODO: document
    TessellationControl = TESSELLATION_CONTROL,

    // TODO: document
    TessellationEvaluation = TESSELLATION_EVALUATION,

    // TODO: document
    Geometry = GEOMETRY,

    // TODO: document
    Fragment = FRAGMENT,

    // TODO: document
    Compute = COMPUTE,

    // TODO: document
    Raygen = RAYGEN_KHR {
        device_extensions: [khr_ray_tracing_pipeline, nv_ray_tracing],
    },

    // TODO: document
    AnyHit = ANY_HIT_KHR {
        device_extensions: [khr_ray_tracing_pipeline, nv_ray_tracing],
    },

    // TODO: document
    ClosestHit = CLOSEST_HIT_KHR {
        device_extensions: [khr_ray_tracing_pipeline, nv_ray_tracing],
    },

    // TODO: document
    Miss = MISS_KHR {
        device_extensions: [khr_ray_tracing_pipeline, nv_ray_tracing],
    },

    // TODO: document
    Intersection = INTERSECTION_KHR {
        device_extensions: [khr_ray_tracing_pipeline, nv_ray_tracing],
    },

    // TODO: document
    Callable = CALLABLE_KHR {
        device_extensions: [khr_ray_tracing_pipeline, nv_ray_tracing],
    },

    /*
    // TODO: document
    Task = TASK_NV {
        device_extensions: [nv_mesh_shader],
    },

    // TODO: document
    Mesh = MESH_NV {
        device_extensions: [nv_mesh_shader],
    },

    // TODO: document
    SubpassShading = SUBPASS_SHADING_HUAWEI {
        device_extensions: [huawei_subpass_shading],
    },
     */
}

impl From<ShaderExecution> for ShaderStage {
    #[inline]
    fn from(val: ShaderExecution) -> Self {
        match val {
            ShaderExecution::Vertex => Self::Vertex,
            ShaderExecution::TessellationControl => Self::TessellationControl,
            ShaderExecution::TessellationEvaluation => Self::TessellationEvaluation,
            ShaderExecution::Geometry(_) => Self::Geometry,
            ShaderExecution::Fragment => Self::Fragment,
            ShaderExecution::Compute => Self::Compute,
            ShaderExecution::RayGeneration => Self::Raygen,
            ShaderExecution::AnyHit => Self::AnyHit,
            ShaderExecution::ClosestHit => Self::ClosestHit,
            ShaderExecution::Miss => Self::Miss,
            ShaderExecution::Intersection => Self::Intersection,
            ShaderExecution::Callable => Self::Callable,
        }
    }
}

impl From<ShaderStage> for ShaderStages {
    #[inline]
    fn from(val: ShaderStage) -> Self {
        match val {
            ShaderStage::Vertex => Self {
                vertex: true,
                ..Self::empty()
            },
            ShaderStage::TessellationControl => Self {
                tessellation_control: true,
                ..Self::empty()
            },
            ShaderStage::TessellationEvaluation => Self {
                tessellation_evaluation: true,
                ..Self::empty()
            },
            ShaderStage::Geometry => Self {
                geometry: true,
                ..Self::empty()
            },
            ShaderStage::Fragment => Self {
                fragment: true,
                ..Self::empty()
            },
            ShaderStage::Compute => Self {
                compute: true,
                ..Self::empty()
            },
            ShaderStage::Raygen => Self {
                raygen: true,
                ..Self::empty()
            },
            ShaderStage::AnyHit => Self {
                any_hit: true,
                ..Self::empty()
            },
            ShaderStage::ClosestHit => Self {
                closest_hit: true,
                ..Self::empty()
            },
            ShaderStage::Miss => Self {
                miss: true,
                ..Self::empty()
            },
            ShaderStage::Intersection => Self {
                intersection: true,
                ..Self::empty()
            },
            ShaderStage::Callable => Self {
                callable: true,
                ..Self::empty()
            },
        }
    }
}

vulkan_bitflags! {
    /// A set of shader stages.
    #[non_exhaustive]
    ShaderStages = ShaderStageFlags(u32);

    // TODO: document
    vertex = VERTEX,

    // TODO: document
    tessellation_control = TESSELLATION_CONTROL,

    // TODO: document
    tessellation_evaluation = TESSELLATION_EVALUATION,

    // TODO: document
    geometry = GEOMETRY,

    // TODO: document
    fragment = FRAGMENT,

    // TODO: document
    compute = COMPUTE,

    // TODO: document
    raygen = RAYGEN_KHR {
        device_extensions: [khr_ray_tracing_pipeline, nv_ray_tracing],
    },

    // TODO: document
    any_hit = ANY_HIT_KHR {
        device_extensions: [khr_ray_tracing_pipeline, nv_ray_tracing],
    },

    // TODO: document
    closest_hit = CLOSEST_HIT_KHR {
        device_extensions: [khr_ray_tracing_pipeline, nv_ray_tracing],
    },

    // TODO: document
    miss = MISS_KHR {
        device_extensions: [khr_ray_tracing_pipeline, nv_ray_tracing],
    },

    // TODO: document
    intersection = INTERSECTION_KHR {
        device_extensions: [khr_ray_tracing_pipeline, nv_ray_tracing],
    },

    // TODO: document
    callable = CALLABLE_KHR {
        device_extensions: [khr_ray_tracing_pipeline, nv_ray_tracing],
    },

    /*
    // TODO: document
    task = TASK_NV {
        device_extensions: [nv_mesh_shader],
    },

    // TODO: document
    mesh = MESH_NV {
        device_extensions: [nv_mesh_shader],
    },

    // TODO: document
    subpass_shading = SUBPASS_SHADING_HUAWEI {
        device_extensions: [huawei_subpass_shading],
    },
     */
}

impl ShaderStages {
    /// Creates a `ShaderStages` struct with all graphics stages set to `true`.
    #[inline]
    pub const fn all_graphics() -> ShaderStages {
        ShaderStages {
            vertex: true,
            tessellation_control: true,
            tessellation_evaluation: true,
            geometry: true,
            fragment: true,
            ..ShaderStages::empty()
        }
    }

    /// Creates a `ShaderStages` struct with the compute stage set to `true`.
    #[inline]
    pub const fn compute() -> ShaderStages {
        ShaderStages {
            compute: true,
            ..ShaderStages::empty()
        }
    }
}

impl From<ShaderStages> for PipelineStages {
    #[inline]
    fn from(stages: ShaderStages) -> PipelineStages {
        let ShaderStages {
            vertex,
            tessellation_control,
            tessellation_evaluation,
            geometry,
            fragment,
            compute,
            raygen,
            any_hit,
            closest_hit,
            miss,
            intersection,
            callable,
            _ne: _,
        } = stages;

        PipelineStages {
            vertex_shader: vertex,
            tessellation_control_shader: tessellation_control,
            tessellation_evaluation_shader: tessellation_evaluation,
            geometry_shader: geometry,
            fragment_shader: fragment,
            compute_shader: compute,
            ray_tracing_shader: raygen | any_hit | closest_hit | miss | intersection | callable,
            ..PipelineStages::empty()
        }
    }
}

fn check_spirv_version(device: &Device, mut version: Version) -> Result<(), ShaderSupportError> {
    version.patch = 0; // Ignore the patch version

    match version {
        Version::V1_0 => {}
        Version::V1_1 | Version::V1_2 | Version::V1_3 => {
            if !(device.api_version() >= Version::V1_1) {
                return Err(ShaderSupportError::RequirementsNotMet(&[
                    "Vulkan API version 1.1",
                ]));
            }
        }
        Version::V1_4 => {
            if !(device.api_version() >= Version::V1_2 || device.enabled_extensions().khr_spirv_1_4)
            {
                return Err(ShaderSupportError::RequirementsNotMet(&[
                    "Vulkan API version 1.2",
                    "extension `khr_spirv_1_4`",
                ]));
            }
        }
        Version::V1_5 => {
            if !(device.api_version() >= Version::V1_2) {
                return Err(ShaderSupportError::RequirementsNotMet(&[
                    "Vulkan API version 1.2",
                ]));
            }
        }
        Version::V1_6 => {
            if !(device.api_version() >= Version::V1_3) {
                return Err(ShaderSupportError::RequirementsNotMet(&[
                    "Vulkan API version 1.3",
                ]));
            }
        }
        _ => return Err(ShaderSupportError::NotSupportedByVulkan),
    }
    Ok(())
}