use core::{fmt, ops};

/// A bitmask, tracking which 4-byte slots have been written via `set_immediates`.
/// Bit N corresponds to bytes [N*4 .. N*4+4).
#[derive(Clone, Copy, Debug, Default, PartialEq, Eq)]
#[cfg_attr(feature = "serialize", derive(serde::Serialize))]
#[cfg_attr(feature = "deserialize", derive(serde::Deserialize))]
pub struct ImmediateSlots(u64);

impl ImmediateSlots {
    pub const fn from_raw(raw: u64) -> Self {
        Self(raw)
    }

    /// Compute the bitmask for a byte range [offset .. offset + size_bytes).
    pub const fn from_range(offset: u32, size_bytes: u32) -> Self {
        if size_bytes == 0 {
            return Self(0);
        }
        let lo = offset / 4;
        let hi = (offset + size_bytes).div_ceil(4);
        Self(u64::MAX << lo & u64::MAX >> (64 - hi))
    }

    /// Compute the slots occupied by a type at a given byte offset,
    /// excluding padding between struct members.
    pub fn from_type(
        ty: &crate::TypeInner,
        offset: u32,
        types: &crate::UniqueArena<crate::Type>,
        gctx: crate::proc::GlobalCtx,
    ) -> Self {
        match *ty {
            crate::TypeInner::Struct { ref members, .. } => {
                let mut slots = Self::default();
                for member in members {
                    let member_ty = &types[member.ty].inner;
                    slots |= Self::from_type(member_ty, offset + member.offset, types, gctx);
                }
                slots
            }
            _ => Self::from_range(offset, ty.size(gctx)),
        }
    }

    /// Returns true if `self` contains all bits in `other`.
    pub const fn contains(self, other: Self) -> bool {
        other.0 & !self.0 == 0
    }

    /// Returns the bits in `self` that are not set in `other`.
    pub const fn difference(self, other: Self) -> Self {
        Self(self.0 & !other.0)
    }

    /// Returns the byte size of the `var<immediate>` type in a module.
    /// Zero if the module has no `var<immediate>`.
    pub fn size_for_module(module: &crate::Module) -> u32 {
        module
            .global_variables
            .iter()
            .find(|&(_, var)| var.space == crate::AddressSpace::Immediate)
            .map(|(_, var)| module.types[var.ty].inner.size(module.to_ctx()))
            .unwrap_or(0)
    }

    /// Compute the immediate slot bitmask for a pointer expression that
    /// refers to (part of) an immediate global variable.
    ///
    /// `global` is the handle of the immediate global variable that this
    /// pointer derives from (obtained from `assignable_global`).
    pub(crate) fn for_pointer(
        pointer: crate::arena::Handle<crate::Expression>,
        global: crate::arena::Handle<crate::GlobalVariable>,
        expression_arena: &crate::Arena<crate::Expression>,
        global_vars: &crate::Arena<crate::GlobalVariable>,
        types: &crate::UniqueArena<crate::Type>,
    ) -> Self {
        use crate::Expression as E;
        use crate::TypeInner;

        let gctx = crate::proc::GlobalCtx {
            types,
            constants: &crate::Arena::new(),
            overrides: &crate::Arena::new(),
            global_expressions: &crate::Arena::new(),
        };

        let global_ty = &types[global_vars[global].ty].inner;

        match expression_arena[pointer] {
            E::GlobalVariable(_) => Self::from_type(global_ty, 0, types, gctx),
            E::AccessIndex { base, index } => {
                if let E::GlobalVariable(_) = expression_arena[base] {
                    if let TypeInner::Struct { ref members, .. } = *global_ty {
                        let member = &members[index as usize];
                        let member_ty = &types[member.ty].inner;
                        return Self::from_type(member_ty, member.offset, types, gctx);
                    }
                }
                Self::from_type(global_ty, 0, types, gctx)
            }
            _ => Self::from_type(global_ty, 0, types, gctx),
        }
    }
}

impl ops::BitOrAssign for ImmediateSlots {
    fn bitor_assign(&mut self, rhs: Self) {
        self.0 |= rhs.0;
    }
}

impl ops::BitOr for ImmediateSlots {
    type Output = Self;
    fn bitor(self, rhs: Self) -> Self {
        Self(self.0 | rhs.0)
    }
}

impl fmt::Display for ImmediateSlots {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        if self.0 == 0 {
            return write!(f, "(none)");
        }
        let mut first = true;
        let mut bit = 0u32;
        while bit < 64 {
            if self.0 & (1u64 << bit) != 0 {
                let start = bit * 4;
                while bit < 64 && self.0 & (1u64 << bit) != 0 {
                    bit += 1;
                }
                let end = bit * 4;
                if !first {
                    write!(f, ", ")?;
                }
                write!(f, "{start}..{end}")?;
                first = false;
            } else {
                bit += 1;
            }
        }
        Ok(())
    }
}

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

    #[test]
    fn range_single() {
        assert_eq!(
            ImmediateSlots::from_range(0, 4),
            ImmediateSlots::from_raw(0b1)
        );
        assert_eq!(
            ImmediateSlots::from_range(4, 4),
            ImmediateSlots::from_raw(0b10)
        );
        assert_eq!(
            ImmediateSlots::from_range(8, 4),
            ImmediateSlots::from_raw(0b100)
        );
    }

    #[test]
    fn range_vec4() {
        assert_eq!(
            ImmediateSlots::from_range(0, 16),
            ImmediateSlots::from_raw(0b1111)
        );
        assert_eq!(
            ImmediateSlots::from_range(16, 16),
            ImmediateSlots::from_raw(0b1111_0000)
        );
    }

    #[test]
    fn range_full_256() {
        assert_eq!(
            ImmediateSlots::from_range(0, 256),
            ImmediateSlots::from_raw(u64::MAX)
        );
    }

    #[test]
    fn from_type_excludes_struct_padding() {
        let module = crate::front::wgsl::parse_str("struct S { a: f32, b: vec4<f32> }").unwrap();
        let struct_ty = (module.types.iter().map(|ty| ty.1))
            .find(|ty| ty.name.as_deref() == Some("S"))
            .unwrap();
        let slots = ImmediateSlots::from_type(&struct_ty.inner, 0, &module.types, module.to_ctx());
        assert_eq!(slots, ImmediateSlots::from_raw(0b1111_0001));
    }

    #[test]
    fn range_unaligned() {
        assert_eq!(
            ImmediateSlots::from_range(0, 3),
            ImmediateSlots::from_raw(0b1)
        );
        assert_eq!(
            ImmediateSlots::from_range(0, 5),
            ImmediateSlots::from_raw(0b11)
        );
    }

    #[test]
    fn contains() {
        let required = ImmediateSlots::from_raw(0b1111_0001);
        let mut set = ImmediateSlots::default();
        assert!(!set.contains(required));
        set |= ImmediateSlots::from_range(0, 4);
        assert!(!set.contains(required));
        set |= ImmediateSlots::from_range(16, 16);
        assert!(set.contains(required));
    }

    #[test]
    fn difference() {
        let required = ImmediateSlots::from_raw(0b1111_0001);
        let set = ImmediateSlots::from_range(0, 4);
        assert_eq!(
            required.difference(set),
            ImmediateSlots::from_raw(0b1111_0000)
        );
    }
}