use super::{TypeLike, shaped_type_like::ShapedTypeLike};
use crate::{
    Error,
    ir::{Attribute, Location, Type, affine_map::AffineMap, attribute::AttributeLike},
};
use mlir_sys::{
    MlirType, mlirMemRefTypeContiguousGet, mlirMemRefTypeContiguousGetChecked, mlirMemRefTypeGet,
    mlirMemRefTypeGetAffineMap, mlirMemRefTypeGetChecked, mlirMemRefTypeGetLayout,
    mlirMemRefTypeGetMemorySpace, mlirMemRefTypeGetStridesAndOffset,
};

/// A mem-ref type.
#[derive(Clone, Copy, Debug, Hash)]
pub struct MemRefType<'c> {
    r#type: Type<'c>,
}

impl<'c> MemRefType<'c> {
    /// Creates a mem-ref type.
    pub fn new(
        r#type: Type<'c>,
        dimensions: &[i64],
        layout: Option<Attribute<'c>>,
        memory_space: Option<Attribute<'c>>,
    ) -> Self {
        unsafe {
            Self::from_raw(mlirMemRefTypeGet(
                r#type.to_raw(),
                dimensions.len() as _,
                dimensions.as_ptr() as *const _,
                layout.unwrap_or_else(|| Attribute::null()).to_raw(),
                memory_space.unwrap_or_else(|| Attribute::null()).to_raw(),
            ))
        }
    }

    /// Creates a mem-ref type with diagnostics.
    pub fn checked(
        location: Location<'c>,
        r#type: Type<'c>,
        dimensions: &[u64],
        layout: Attribute<'c>,
        memory_space: Attribute<'c>,
    ) -> Option<Self> {
        unsafe {
            Self::from_option_raw(mlirMemRefTypeGetChecked(
                location.to_raw(),
                r#type.to_raw(),
                dimensions.len() as isize,
                dimensions.as_ptr() as *const i64,
                layout.to_raw(),
                memory_space.to_raw(),
            ))
        }
    }

    /// Creates a contiguous mem-ref type (identity layout).
    pub fn contiguous(
        element_type: Type<'c>,
        dimensions: &[i64],
        memory_space: Option<Attribute<'c>>,
    ) -> Self {
        unsafe {
            Self::from_raw(mlirMemRefTypeContiguousGet(
                element_type.to_raw(),
                dimensions.len() as isize,
                dimensions.as_ptr(),
                memory_space.unwrap_or_else(|| Attribute::null()).to_raw(),
            ))
        }
    }

    /// Creates a contiguous mem-ref type with diagnostics.
    pub fn contiguous_checked(
        location: Location<'c>,
        element_type: Type<'c>,
        dimensions: &[i64],
        memory_space: Option<Attribute<'c>>,
    ) -> Option<Self> {
        unsafe {
            Self::from_option_raw(mlirMemRefTypeContiguousGetChecked(
                location.to_raw(),
                element_type.to_raw(),
                dimensions.len() as isize,
                dimensions.as_ptr(),
                memory_space.unwrap_or_else(|| Attribute::null()).to_raw(),
            ))
        }
    }

    /// Returns the strides and offset of the mem-ref type.
    ///
    /// Returns a tuple of `(strides, offset)` where `strides` has length equal
    /// to the rank.
    pub fn strides_and_offset(&self) -> (Vec<i64>, i64) {
        let rank = self.rank();
        let mut strides = vec![0i64; rank];
        let mut offset = 0i64;
        unsafe {
            mlirMemRefTypeGetStridesAndOffset(
                self.r#type.to_raw(),
                strides.as_mut_ptr(),
                &mut offset,
            );
        }
        (strides, offset)
    }

    /// Returns a layout.
    pub fn layout(&self) -> Attribute<'c> {
        unsafe { Attribute::from_raw(mlirMemRefTypeGetLayout(self.r#type.to_raw())) }
    }

    /// Returns an affine map.
    pub fn affine_map(&self) -> AffineMap<'c> {
        unsafe { AffineMap::from_raw(mlirMemRefTypeGetAffineMap(self.r#type.to_raw())) }
    }

    /// Returns a memory space.
    pub fn memory_space(&self) -> Option<Attribute<'c>> {
        unsafe { Attribute::from_option_raw(mlirMemRefTypeGetMemorySpace(self.r#type.to_raw())) }
    }

    unsafe fn from_option_raw(raw: MlirType) -> Option<Self> {
        if raw.ptr.is_null() {
            None
        } else {
            Some(unsafe { Self::from_raw(raw) })
        }
    }
}

impl<'c> ShapedTypeLike<'c> for MemRefType<'c> {}

type_traits!(MemRefType, is_mem_ref, "mem ref");

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

    #[test]
    fn new() {
        let context = Context::new();

        assert_eq!(
            Type::from(MemRefType::new(Type::float64(&context), &[42], None, None,)),
            Type::parse(&context, "memref<42xf64>").unwrap()
        );
    }

    #[test]
    fn dynamic_dimension() {
        let context = Context::new();

        assert_eq!(
            Type::from(MemRefType::new(
                Type::float64(&context),
                &[i64::MIN],
                None,
                None,
            )),
            Type::parse(&context, "memref<?xf64>").unwrap()
        );
    }

    #[test]
    fn layout() {
        let context = Context::new();

        assert_eq!(
            MemRefType::new(Type::index(&context), &[42, 42], None, None,).layout(),
            Attribute::parse(&context, "affine_map<(d0, d1) -> (d0, d1)>").unwrap(),
        );
    }

    #[test]
    fn affine_map() {
        let context = Context::new();

        assert_eq!(
            MemRefType::new(Type::index(&context), &[42, 42], None, None,)
                .affine_map()
                .to_string(),
            "(d0, d1) -> (d0, d1)"
        );
    }

    #[test]
    fn memory_space() {
        let context = Context::new();

        assert_eq!(
            MemRefType::new(Type::index(&context), &[42, 42], None, None).memory_space(),
            None,
        );
    }

    #[test]
    fn contiguous() {
        let context = Context::new();

        assert_eq!(
            Type::from(MemRefType::contiguous(Type::float64(&context), &[42], None)),
            Type::parse(&context, "memref<42xf64>").unwrap()
        );
    }

    #[test]
    fn contiguous_checked() {
        let context = Context::new();

        assert_eq!(
            MemRefType::contiguous_checked(
                Location::unknown(&context),
                Type::float64(&context),
                &[42],
                None,
            )
            .map(Type::from),
            Type::parse(&context, "memref<42xf64>")
        );
    }

    #[test]
    fn contiguous_checked_fail() {
        let context = Context::new();

        // index type is not a valid element type for memref in the checked variant
        assert_eq!(
            MemRefType::contiguous_checked(
                Location::unknown(&context),
                Type::index(&context),
                &[42],
                None,
            )
            .map(Type::from),
            Type::parse(&context, "memref<42xindex>")
        );
    }

    #[test]
    fn strides_and_offset() {
        let context = Context::new();

        let memref = MemRefType::new(Type::float64(&context), &[4, 8], None, None);
        let (strides, offset) = memref.strides_and_offset();
        assert_eq!(strides.len(), 2);
        assert_eq!(offset, 0);
        // Row-major strides for 4x8: strides[0]=8, strides[1]=1
        assert_eq!(strides[0], 8);
        assert_eq!(strides[1], 1);
    }
}
