use bitflags::bitflags;

pub struct Function<'a>
{
    pub(super) inner: &'a shader_slang::reflection::Function,
}

impl Function<'_>
{
    #[must_use]
    pub fn name(&self) -> Option<&str>
    {
        self.inner.name()
    }
}

pub struct EntryPointReflection<'a>
{
    inner: &'a shader_slang::reflection::EntryPoint,
}

impl<'a> EntryPointReflection<'a>
{
    #[must_use]
    pub fn name(&self) -> Option<&str>
    {
        self.inner.name()
    }

    #[must_use]
    pub fn name_override(&self) -> Option<&str>
    {
        self.inner.name_override()
    }

    #[must_use]
    pub fn stage(&self) -> Stage
    {
        Stage::from_slang_stage(self.inner.stage())
    }

    #[must_use]
    pub fn parameters(&self) -> impl ExactSizeIterator<Item = VariableLayout<'a>>
    {
        self.inner
            .parameters()
            .map(|param| VariableLayout { inner: param })
    }

    #[must_use]
    pub fn var_layout(&self) -> Option<VariableLayout<'a>>
    {
        Some(VariableLayout { inner: self.inner.var_layout()? })
    }
}

pub struct Program<'a>
{
    pub(super) inner: &'a shader_slang::reflection::Shader,
}

impl<'a> Program<'a>
{
    #[must_use]
    pub fn get_entry_point_by_index(&self, index: u32)
        -> Option<EntryPointReflection<'a>>
    {
        Some(EntryPointReflection {
            inner: self.inner.entry_point_by_index(index)?,
        })
    }

    #[must_use]
    pub fn get_entry_point_by_name(&self, name: &str)
        -> Option<EntryPointReflection<'a>>
    {
        Some(EntryPointReflection {
            inner: self.inner.find_entry_point_by_name(name)?,
        })
    }

    #[must_use]
    pub fn entry_points(
        &self,
    ) -> impl ExactSizeIterator<Item = EntryPointReflection<'a>> + use<'a>
    {
        self.inner
            .entry_points()
            .map(|entry_point| EntryPointReflection { inner: entry_point })
    }

    #[must_use]
    pub fn global_params_type_layout(&self) -> Option<TypeLayout<'a>>
    {
        Some(TypeLayout {
            inner: self.inner.global_params_type_layout()?,
        })
    }

    #[must_use]
    pub fn global_params_var_layout(&self) -> Option<VariableLayout<'a>>
    {
        Some(VariableLayout {
            inner: self.inner.global_params_var_layout()?,
        })
    }

    #[must_use]
    pub fn get_type(&self, name: &str) -> Option<TypeReflection<'a>>
    {
        Some(TypeReflection {
            inner: self.inner.find_type_by_name(name)?,
        })
    }

    #[must_use]
    pub fn get_type_layout(&self, ty: &TypeReflection<'a>) -> Option<TypeLayout<'a>>
    {
        Some(TypeLayout {
            inner: self
                .inner
                .type_layout(ty.inner, shader_slang::LayoutRules::Default)?,
        })
    }
}

#[derive(Clone, Copy)]
pub struct VariableLayout<'a>
{
    inner: &'a shader_slang::reflection::VariableLayout,
}

impl<'a> VariableLayout<'a>
{
    #[must_use]
    pub fn name(&self) -> Option<&'a str>
    {
        self.inner.name()
    }

    #[must_use]
    pub fn semantic_name(&self) -> Option<&str>
    {
        self.inner.semantic_name()
    }

    #[must_use]
    pub fn binding_index(&self) -> u32
    {
        self.inner
            .offset(shader_slang::ParameterCategory::DescriptorTableSlot) as u32

        // self.inner.binding_index()
    }

    #[must_use]
    pub fn varying_input_offset(&self) -> Option<usize>
    {
        if !self
            .inner
            .categories()
            .any(|category| category == shader_slang::ParameterCategory::VaryingInput)
        {
            return None;
        }

        Some(
            self.inner
                .offset(shader_slang::ParameterCategory::VaryingInput),
        )
    }

    #[must_use]
    pub fn binding_space(&self) -> u32
    {
        self.inner.binding_space()
    }

    #[must_use]
    pub fn semantic_index(&self) -> usize
    {
        self.inner.semantic_index()
    }

    #[must_use]
    pub fn offset(&self) -> usize
    {
        self.inner.offset(shader_slang::ParameterCategory::Uniform)
    }

    #[must_use]
    pub fn ty(&self) -> Option<TypeReflection<'a>>
    {
        self.inner.ty().map(|ty| TypeReflection { inner: ty })
    }

    #[must_use]
    pub fn type_layout(&self) -> Option<TypeLayout<'a>>
    {
        Some(TypeLayout { inner: self.inner.type_layout()? })
    }
}

#[derive(Clone, Copy)]
pub struct TypeLayout<'a>
{
    inner: &'a shader_slang::reflection::TypeLayout,
}

impl<'a> TypeLayout<'a>
{
    #[must_use]
    pub fn kind(&self) -> TypeKind
    {
        TypeKind::from_slang_type_kind(self.inner.kind())
    }

    #[must_use]
    pub fn scalar_type(&self) -> Option<ScalarType>
    {
        Some(ScalarType::from_slang_scalar_type(
            self.inner.scalar_type()?,
        ))
    }

    #[must_use]
    pub fn resource_shape(&self) -> Option<ResourceShape>
    {
        Some(ResourceShape::from_bits_retain(
            self.inner.resource_shape()? as u32,
        ))
    }

    #[must_use]
    pub fn get_field_by_name(&self, name: &str) -> Option<VariableLayout<'a>>
    {
        let index = self.inner.find_field_index_by_name(name);

        if index < 0 {
            return None;
        }

        let index = u32::try_from(index.cast_unsigned()).expect("Should not happend");

        let field = self.inner.field_by_index(index)?;

        Some(VariableLayout { inner: field })
    }

    #[must_use]
    pub fn parameter_category(&self) -> ParameterCategory
    {
        ParameterCategory::from_slang_parameter_category(self.inner.parameter_category())
    }

    #[must_use]
    pub fn binding_range_descriptor_set_index(&self, index: i64) -> i64
    {
        self.inner.binding_range_descriptor_set_index(index)
    }

    #[must_use]
    pub fn get_field_binding_range_offset_by_name(&self, name: &str) -> Option<u64>
    {
        let field_index = self.inner.find_field_index_by_name(name);

        if field_index < 0 {
            return None;
        }

        let field_binding_range_offset =
            self.inner.field_binding_range_offset(field_index);

        if field_binding_range_offset < 0 {
            return None;
        }

        Some(field_binding_range_offset.cast_unsigned())
    }

    #[must_use]
    pub fn ty(&self) -> Option<TypeReflection<'a>>
    {
        self.inner.ty().map(|ty| TypeReflection { inner: ty })
    }

    #[must_use]
    pub fn fields(&self) -> FieldIter<'a>
    {
        FieldIter {
            type_layout: *self,
            cnt: self.field_cnt(),
            index: 0,
        }
    }

    #[must_use]
    pub fn field_cnt(&self) -> u32
    {
        self.inner.field_count()
    }

    #[must_use]
    pub fn element_cnt(&self) -> Option<usize>
    {
        self.inner.element_count()
    }

    #[must_use]
    pub fn row_cnt(&self) -> Option<u32>
    {
        self.inner.row_count()
    }

    #[must_use]
    pub fn column_cnt(&self) -> Option<u32>
    {
        self.inner.column_count()
    }

    #[must_use]
    pub fn element_type_layout(&self) -> Option<TypeLayout<'a>>
    {
        self.inner
            .element_type_layout()
            .map(|type_layout| TypeLayout { inner: type_layout })
    }

    #[must_use]
    pub fn element_var_layout(&self) -> Option<VariableLayout<'a>>
    {
        self.inner
            .element_var_layout()
            .map(|var_layout| VariableLayout { inner: var_layout })
    }

    #[must_use]
    pub fn container_var_layout(&self) -> Option<VariableLayout<'a>>
    {
        self.inner
            .container_var_layout()
            .map(|var_layout| VariableLayout { inner: var_layout })
    }

    #[must_use]
    pub fn uniform_size(&self) -> Option<usize>
    {
        if !self
            .inner
            .categories()
            .any(|category| category == shader_slang::ParameterCategory::Uniform)
        {
            return None;
        }

        Some(self.inner.size(shader_slang::ParameterCategory::Uniform))
    }

    #[must_use]
    pub fn stride(&self) -> usize
    {
        self.inner.stride(self.inner.categories().next().unwrap())
    }
}

pub struct FieldIter<'a>
{
    type_layout: TypeLayout<'a>,
    cnt: u32,
    index: u32,
}

impl<'a> Iterator for FieldIter<'a>
{
    type Item = VariableLayout<'a>;

    fn next(&mut self) -> Option<Self::Item>
    {
        if self.index == self.cnt {
            return None;
        }

        let Some(field) = self.type_layout.inner.field_by_index(self.index) else {
            unreachable!();
        };

        self.index += 1;

        Some(VariableLayout { inner: field })
    }

    fn size_hint(&self) -> (usize, Option<usize>)
    {
        let len = (self.cnt - self.index) as usize;

        (len, Some(len))
    }
}

impl ExactSizeIterator for FieldIter<'_>
{
    fn len(&self) -> usize
    {
        (self.cnt - self.index) as usize
    }
}

impl DoubleEndedIterator for FieldIter<'_>
{
    fn next_back(&mut self) -> Option<Self::Item>
    {
        if self.cnt == 0 || self.index == self.cnt - 1 {
            return None;
        }

        let Some(field) = self.type_layout.inner.field_by_index(self.cnt - 1) else {
            unreachable!();
        };

        self.cnt -= 1;

        Some(VariableLayout { inner: field })
    }
}

pub struct TypeReflection<'a>
{
    inner: &'a shader_slang::reflection::Type,
}

impl TypeReflection<'_>
{
    #[must_use]
    pub fn kind(&self) -> TypeKind
    {
        TypeKind::from_slang_type_kind(self.inner.kind())
    }
}

#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Hash)]
#[non_exhaustive]
pub enum TypeKind
{
    None,
    Struct,
    Enum,
    Array,
    Matrix,
    Vector,
    Scalar,
    ConstantBuffer,
    Resource,
    SamplerState,
    TextureBuffer,
    ShaderStorageBuffer,
    ParameterBlock,
    GenericTypeParameter,
    Interface,
    OutputStream,
    MeshOutput,
    Specialized,
    Feedback,
    Pointer,
    DynamicResource,
    Count,
}

impl TypeKind
{
    fn from_slang_type_kind(type_kind: shader_slang::TypeKind) -> Self
    {
        match type_kind {
            shader_slang::TypeKind::None => Self::None,
            shader_slang::TypeKind::Struct => Self::Struct,
            shader_slang::TypeKind::Enum => Self::Enum,
            shader_slang::TypeKind::Array => Self::Array,
            shader_slang::TypeKind::Matrix => Self::Matrix,
            shader_slang::TypeKind::Vector => Self::Vector,
            shader_slang::TypeKind::Scalar => Self::Scalar,
            shader_slang::TypeKind::ConstantBuffer => Self::ConstantBuffer,
            shader_slang::TypeKind::Resource => Self::Resource,
            shader_slang::TypeKind::SamplerState => Self::SamplerState,
            shader_slang::TypeKind::TextureBuffer => Self::TextureBuffer,
            shader_slang::TypeKind::ShaderStorageBuffer => Self::ShaderStorageBuffer,
            shader_slang::TypeKind::ParameterBlock => Self::ParameterBlock,
            shader_slang::TypeKind::GenericTypeParameter => Self::GenericTypeParameter,
            shader_slang::TypeKind::Interface => Self::Interface,
            shader_slang::TypeKind::OutputStream => Self::OutputStream,
            shader_slang::TypeKind::MeshOutput => Self::MeshOutput,
            shader_slang::TypeKind::Specialized => Self::Specialized,
            shader_slang::TypeKind::Feedback => Self::Feedback,
            shader_slang::TypeKind::Pointer => Self::Pointer,
            shader_slang::TypeKind::DynamicResource => Self::DynamicResource,
            shader_slang::TypeKind::Count => Self::Count,
        }
    }
}

#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord)]
#[non_exhaustive]
pub enum ScalarType
{
    None,
    Void,
    Bool,
    Int32,
    Uint32,
    Int64,
    Uint64,
    Float16,
    Float32,
    Float64,
    Int8,
    Uint8,
    Int16,
    Uint16,
    Intptr,
    Uintptr,
    Bfloat16,
    FloatE4m3,
    FloatE5m2,
}

impl ScalarType
{
    fn from_slang_scalar_type(scalar_type: shader_slang::ScalarType) -> Self
    {
        match scalar_type {
            shader_slang::ScalarType::None => Self::None,
            shader_slang::ScalarType::Void => Self::Void,
            shader_slang::ScalarType::Bool => Self::Bool,
            shader_slang::ScalarType::Int32 => Self::Int32,
            shader_slang::ScalarType::Uint32 => Self::Uint32,
            shader_slang::ScalarType::Int64 => Self::Int64,
            shader_slang::ScalarType::Uint64 => Self::Uint64,
            shader_slang::ScalarType::Float16 => Self::Float16,
            shader_slang::ScalarType::Float32 => Self::Float32,
            shader_slang::ScalarType::Float64 => Self::Float64,
            shader_slang::ScalarType::Int8 => Self::Int8,
            shader_slang::ScalarType::Uint8 => Self::Uint8,
            shader_slang::ScalarType::Int16 => Self::Int16,
            shader_slang::ScalarType::Uint16 => Self::Uint16,
            shader_slang::ScalarType::Intptr => Self::Intptr,
            shader_slang::ScalarType::Uintptr => Self::Uintptr,
            shader_slang::ScalarType::Bfloat16 => Self::Bfloat16,
            shader_slang::ScalarType::FloatE4m3 => Self::FloatE4m3,
            shader_slang::ScalarType::FloatE5m2 => Self::FloatE5m2,
            #[allow(unreachable_patterns)]
            _ => unimplemented!("conversion from slang scalar type"),
        }
    }
}

bitflags! {
#[derive(Debug, Copy, Clone)]
pub struct ResourceShape: u32 {
    const BASE = shader_slang::ResourceShape::SlangResourceBaseShapeMask as u32;
    const NONE = shader_slang::ResourceShape::SlangResourceNone as u32;
    const TEXTURE_1D = shader_slang::ResourceShape::SlangTexture1d as u32;
    const TEXTURE_2D = shader_slang::ResourceShape::SlangTexture2d as u32;
    const TEXTURE_3D = shader_slang::ResourceShape::SlangTexture3d as u32;
    const TEXTURE_CUBE = shader_slang::ResourceShape::SlangTextureCube as u32;
    const TEXTURE_BUFFER = shader_slang::ResourceShape::SlangTextureBuffer as u32;
    const STRUCTURED_BUFFER = shader_slang::ResourceShape::SlangStructuredBuffer as u32;
    const BYTE_ADDRESS_BUFFER = shader_slang::ResourceShape::SlangByteAddressBuffer as u32;
    const UNKNOWN = shader_slang::ResourceShape::SlangResourceUnknown as u32;
    const ACCELERATION_STRUCTURE = shader_slang::ResourceShape::SlangAccelerationStructure as u32;
    const TEXTURE_SUBPASS = shader_slang::ResourceShape::SlangTextureSubpass as u32;
    const EXT_SHAPE = shader_slang::ResourceShape::SlangResourceExtShapeMask as u32;
    const TEXTURE_FEEDBACK_FLAG = shader_slang::ResourceShape::SlangTextureFeedbackFlag as u32;
    const TEXTURE_SHADOW_FLAG = shader_slang::ResourceShape::SlangTextureShadowFlag as u32;
    const TEXTURE_ARRAY_FLAG = shader_slang::ResourceShape::SlangTextureArrayFlag as u32;
    const TEXTURE_MULTISAMPLE_FLAG = shader_slang::ResourceShape::SlangTextureMultisampleFlag as u32;
    const TEXTURE_COMBINED_FLAG = shader_slang::ResourceShape::SlangTextureCombinedFlag as u32;
    const TEXTURE_1D_ARRAY = shader_slang::ResourceShape::SlangTexture1dArray as u32;
    const TEXTURE_2D_ARRAY = shader_slang::ResourceShape::SlangTexture2dArray as u32;
    const TEXTURE_CUBE_ARRAY = shader_slang::ResourceShape::SlangTextureCubeArray as u32;
    const TEXTURE_2D_MULTISAMPLE = shader_slang::ResourceShape::SlangTexture2dMultisample as u32;
    const TEXTURE_2D_MULTISAMPLE_ARRAY = shader_slang::ResourceShape::SlangTexture2dMultisampleArray as u32;
    const TEXTURE_SUBPASS_MULTISAMPLE = shader_slang::ResourceShape::SlangTextureSubpassMultisample as u32;
}
}

#[derive(Debug, Copy, Clone, Hash, PartialEq, Eq)]
pub enum ParameterCategory
{
    None,
    Mixed,
    ConstantBuffer,
    ShaderResource,
    UnorderedAccess,
    VaryingInput,
    VaryingOutput,
    SamplerState,
    Uniform,
    DescriptorTableSlot,
    SpecializationConstant,
    PushConstantBuffer,
    RegisterSpace,
    Generic,
    RayPayload,
    HitAttributes,
    CallablePayload,
    ShaderRecord,
    ExistentialTypeParam,
    ExistentialObjectParam,
    SubElementRegisterSpace,
    Subpass,
    MetalArgumentBufferElement,
    MetalAttribute,
    MetalPayload,
    Count,
}

impl ParameterCategory
{
    fn from_slang_parameter_category(
        parameter_category: shader_slang::ParameterCategory,
    ) -> Self
    {
        match parameter_category {
            shader_slang::ParameterCategory::None => Self::None,
            shader_slang::ParameterCategory::Mixed => Self::Mixed,
            shader_slang::ParameterCategory::ConstantBuffer => Self::ConstantBuffer,
            shader_slang::ParameterCategory::ShaderResource => Self::ShaderResource,
            shader_slang::ParameterCategory::UnorderedAccess => Self::UnorderedAccess,
            shader_slang::ParameterCategory::VaryingInput => Self::VaryingInput,
            shader_slang::ParameterCategory::VaryingOutput => Self::VaryingOutput,
            shader_slang::ParameterCategory::SamplerState => Self::SamplerState,
            shader_slang::ParameterCategory::Uniform => Self::Uniform,
            shader_slang::ParameterCategory::DescriptorTableSlot => {
                Self::DescriptorTableSlot
            }
            shader_slang::ParameterCategory::SpecializationConstant => {
                Self::SpecializationConstant
            }
            shader_slang::ParameterCategory::PushConstantBuffer => {
                Self::PushConstantBuffer
            }
            shader_slang::ParameterCategory::RegisterSpace => Self::RegisterSpace,
            shader_slang::ParameterCategory::Generic => Self::Generic,
            shader_slang::ParameterCategory::RayPayload => Self::RayPayload,
            shader_slang::ParameterCategory::HitAttributes => Self::HitAttributes,
            shader_slang::ParameterCategory::CallablePayload => Self::CallablePayload,
            shader_slang::ParameterCategory::ShaderRecord => Self::ShaderRecord,
            shader_slang::ParameterCategory::ExistentialTypeParam => {
                Self::ExistentialTypeParam
            }
            shader_slang::ParameterCategory::ExistentialObjectParam => {
                Self::ExistentialObjectParam
            }
            shader_slang::ParameterCategory::SubElementRegisterSpace => {
                Self::SubElementRegisterSpace
            }
            shader_slang::ParameterCategory::Subpass => Self::Subpass,
            shader_slang::ParameterCategory::MetalArgumentBufferElement => {
                Self::MetalArgumentBufferElement
            }
            shader_slang::ParameterCategory::MetalAttribute => Self::MetalAttribute,
            shader_slang::ParameterCategory::MetalPayload => Self::MetalPayload,
            shader_slang::ParameterCategory::Count => Self::Count,
        }
    }
}

#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Hash)]
#[non_exhaustive]
pub enum Stage
{
    None,
    Vertex,
    Hull,
    Domain,
    Geometry,
    Fragment,
    Compute,
    RayGeneration,
    Intersection,
    AnyHit,
    ClosestHit,
    Miss,
    Callable,
    Mesh,
    Amplification,
    Dispatch,
    Count,
}

impl Stage
{
    fn from_slang_stage(stage: shader_slang::Stage) -> Self
    {
        match stage {
            shader_slang::Stage::None => Self::None,
            shader_slang::Stage::Vertex => Self::Vertex,
            shader_slang::Stage::Hull => Self::Hull,
            shader_slang::Stage::Domain => Self::Domain,
            shader_slang::Stage::Geometry => Self::Geometry,
            shader_slang::Stage::Fragment => Self::Fragment,
            shader_slang::Stage::Compute => Self::Compute,
            shader_slang::Stage::RayGeneration => Self::RayGeneration,
            shader_slang::Stage::Intersection => Self::Intersection,
            shader_slang::Stage::AnyHit => Self::AnyHit,
            shader_slang::Stage::ClosestHit => Self::ClosestHit,
            shader_slang::Stage::Miss => Self::Miss,
            shader_slang::Stage::Callable => Self::Callable,
            shader_slang::Stage::Mesh => Self::Mesh,
            shader_slang::Stage::Amplification => Self::Amplification,
            shader_slang::Stage::Dispatch => Self::Dispatch,
            shader_slang::Stage::Count => Self::Count,
        }
    }
}