//! Shared aliasable memory allocation helpers for buffers and images.

use std::sync::Arc;

use ash::vk;

use crate::{
    Device,
    buffers::Buffer,
    device::{DeviceObject, GpuAllocation, MemoryBinding},
    images::Image,
};

/// Vulkan memory requirements returned by the device layer.
pub use crate::device::MemoryRequirements;

/// Errors from resource requirement queries, validation, and shared binding.
#[derive(Debug)]
pub enum Error {
    RequirementQuery(crate::Error),
    Bind(crate::Error),
    Allocation(crate::Error),
    MisalignedOffset {
        offset: u64,
        alignment: u64,
    },
    IncompatibleMemoryType {
        memory_type_index: u32,
        memory_type_bits: u32,
    },
    OutOfBounds {
        offset: u64,
        end: u64,
        allocation_size: u64,
    },
    AlreadyBound,
    AlreadyAllocated,
    RequiresDedicatedAllocation,
    GroupMemoryTypeMismatch {
        memory_type_index: u32,
    },
    GroupCapacityExceeded {
        requested_end: u64,
        capacity: u64,
    },
    EmptyGroup,
    NotAllocated,
    InvalidGroupRequirements,
    OffsetOverflow,
}

impl core::fmt::Display for Error {
    fn fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result {
        match self {
            Self::RequirementQuery(error) => {
                write!(f, "failed to query resource memory requirements: {error}")
            }
            Self::Bind(error) => write!(f, "failed to bind resource memory: {error}"),
            Self::Allocation(error) => write!(f, "failed to allocate group memory: {error}"),
            Self::MisalignedOffset { offset, alignment } => write!(
                f,
                "allocation offset {offset} is not aligned to {alignment}"
            ),
            Self::IncompatibleMemoryType {
                memory_type_index,
                memory_type_bits,
            } => write!(
                f,
                "allocation memory type {memory_type_index} is not allowed by mask {memory_type_bits:#010x}"
            ),
            Self::OutOfBounds {
                offset,
                end,
                allocation_size,
            } => write!(
                f,
                "resource range [{offset}, {end}) does not fit allocation size {allocation_size}"
            ),
            Self::AlreadyBound => f.write_str("resource is already bound to an allocation"),
            Self::AlreadyAllocated => f.write_str("memory group has already been allocated"),
            Self::RequiresDedicatedAllocation => f.write_str(
                "resources requiring dedicated allocations cannot be added to a memory group",
            ),
            Self::GroupMemoryTypeMismatch { memory_type_index } => write!(
                f,
                "resource memory requirements do not include memory type {memory_type_index}"
            ),
            Self::GroupCapacityExceeded {
                requested_end,
                capacity,
            } => write!(
                f,
                "memory group capacity exceeded: requested end {requested_end}, capacity {capacity}"
            ),
            Self::EmptyGroup => f.write_str("memory group has no members"),
            Self::NotAllocated => f.write_str("memory group has not been allocated"),
            Self::InvalidGroupRequirements => {
                f.write_str("memory group size or alignment is invalid")
            }
            Self::OffsetOverflow => f.write_str("memory group offset arithmetic overflowed"),
        }
    }
}

impl std::error::Error for Error {
    fn source(&self) -> Option<&(dyn std::error::Error + 'static)> {
        match self {
            Self::RequirementQuery(error) | Self::Bind(error) | Self::Allocation(error) => {
                Some(error)
            }
            _ => None,
        }
    }
}

pub type Result<T> = core::result::Result<T, Error>;

pub(crate) fn validate_binding(
    requirements: &MemoryRequirements,
    allocation: &DeviceObject<GpuAllocation>,
    offset: u64,
) -> Result<()> {
    let alignment = requirements.alignment;
    if alignment == 0 || !offset.is_multiple_of(alignment) {
        return Err(Error::MisalignedOffset { offset, alignment });
    }

    let memory_type_index = allocation.memory_type_index();
    if requirements.memory_type_bits & (1 << memory_type_index) == 0 {
        return Err(Error::IncompatibleMemoryType {
            memory_type_index,
            memory_type_bits: requirements.memory_type_bits,
        });
    }

    let end = offset
        .checked_add(requirements.size)
        .ok_or(Error::OutOfBounds {
            offset,
            end: u64::MAX,
            allocation_size: allocation.size(),
        })?;

    if end > allocation.size() {
        return Err(Error::OutOfBounds {
            offset,
            end,
            allocation_size: allocation.size(),
        });
    }

    Ok(())
}

enum Member<'a> {
    Buffer {
        resource: &'a mut Buffer,
        offset: u64,
    },
    Image {
        resource: &'a mut Image,
        offset: u64,
    },
}

/// Packs resources into a single aliasable allocation with a fixed memory type.
///
/// Members are exclusively borrowed until `bind_all` consumes the group. Their
/// allocation ownership is then transferred to each resource; the group retains
/// no allocation handle after binding.
pub struct MemoryGroup<'a> {
    memory_type_index: u32,
    property_flags: vk::MemoryPropertyFlags,
    size: u64,
    alignment: u64,
    next_offset: u64,
    members: Vec<Member<'a>>,
    allocation: Option<Arc<DeviceObject<GpuAllocation>>>,
}

impl<'a> MemoryGroup<'a> {
    /// Creates an empty group for a fixed memory type and allocation capacity.
    pub fn new(
        memory_type_index: u32,
        property_flags: vk::MemoryPropertyFlags,
        size: u64,
        alignment: u64,
    ) -> Result<Self> {
        if memory_type_index >= u32::BITS || size == 0 || alignment == 0 {
            return Err(Error::InvalidGroupRequirements);
        }
        Ok(Self {
            memory_type_index,
            property_flags,
            size,
            alignment,
            next_offset: 0,
            members: Vec::new(),
            allocation: None,
        })
    }

    /// Adds a buffer, placing it at the next suitably aligned offset.
    pub fn add_buffer(&mut self, resource: &'a mut Buffer) -> Result<u64> {
        if self.allocation.is_some() {
            return Err(Error::AlreadyAllocated);
        }

        let requirements = resource.memory_requirements()?;
        self.add_member_type_check(&requirements)?;
        let offset = self.next_member_offset(&requirements)?;

        let next_offset = offset
            .checked_add(requirements.size)
            .ok_or(Error::OffsetOverflow)?;

        self.members.push(Member::Buffer { resource, offset });

        self.next_offset = next_offset;
        Ok(offset)
    }

    /// Adds either resource kind through a uniform entry point.
    pub fn add(&mut self, member: MemoryMember<'a>) -> Result<u64> {
        match member {
            MemoryMember::Buffer(resource) => self.add_buffer(resource),
            MemoryMember::Image(resource) => self.add_image(resource),
        }
    }

    /// Adds an image, placing it at the next suitably aligned offset.
    pub fn add_image(&mut self, resource: &'a mut Image) -> Result<u64> {
        if self.allocation.is_some() {
            return Err(Error::AlreadyAllocated);
        }

        let requirements = resource.memory_requirements()?;
        self.add_member_type_check(&requirements)?;
        let offset = self.next_member_offset(&requirements)?;

        let next_offset = offset
            .checked_add(requirements.size)
            .ok_or(Error::OffsetOverflow)?;

        self.members.push(Member::Image { resource, offset });

        self.next_offset = next_offset;
        Ok(offset)
    }

    /// Allocates the group's shared backing memory.
    pub fn allocate(&mut self, device: &Device) -> Result<()> {
        if self.members.is_empty() {
            return Err(Error::EmptyGroup);
        }

        if self.allocation.is_some() {
            return Err(Error::AlreadyAllocated);
        }

        let requirements = MemoryRequirements {
            size: self.size,
            alignment: self.alignment,
            memory_type_bits: 1 << self.memory_type_index,
            requires_dedicated_allocation: false,
            prefers_dedicated_allocation: false,
        };

        let allocation = device
            .allocate_memory(&requirements, self.property_flags)
            .map_err(Error::Allocation)?;

        for member in &self.members {
            let (requirements, offset) = match member {
                Member::Buffer { resource, offset } => (resource.memory_requirements()?, *offset),
                Member::Image { resource, offset } => (resource.memory_requirements()?, *offset),
            };
            validate_binding(&requirements, allocation.as_ref(), offset)?;
        }

        self.allocation = Some(allocation);

        Ok(())
    }

    /// Binds every member and transfers the allocation reference into each resource.
    pub fn bind_all(mut self) -> Result<()> {
        let allocation = self.allocation.take().ok_or(Error::NotAllocated)?;

        // Validate resources first to preserve the same resource-level error mapping as
        // individual binding, then prepare the ordered device batch from cached requirements.
        let mut bindings = Vec::with_capacity(self.members.len());
        for member in &self.members {
            match member {
                Member::Buffer { resource, offset } => {
                    resource.validate_unallocated()?;
                    let requirements = resource.memory_requirements()?;
                    validate_binding(&requirements, allocation.as_ref(), *offset)?;
                    resource.validate_device(allocation.device())?;
                    bindings.push(MemoryBinding::Buffer {
                        buffer: resource.raw(),
                        offset: *offset,
                        requirements,
                    });
                }
                Member::Image { resource, offset } => {
                    resource.validate_unallocated()?;
                    let requirements = resource.memory_requirements()?;
                    validate_binding(&requirements, allocation.as_ref(), *offset)?;
                    resource.validate_device(allocation.device())?;
                    bindings.push(MemoryBinding::Image {
                        image: resource.raw(),
                        offset: *offset,
                        requirements,
                    });
                }
            }
        }

        if let Err(bind_error) = allocation
            .device()
            .bind_group_at_offset(&allocation, &bindings)
        {
            // A Vulkan bind failure leaves its earlier prefix bound; retain its allocation just
            // as repeated `bind_allocation` calls would have done.
            for (index, member) in self.members.into_iter().enumerate() {
                if index >= bind_error.bound_count {
                    break;
                }
                match member {
                    Member::Buffer { resource, .. } => {
                        resource.set_shared_allocation(allocation.clone());
                    }
                    Member::Image { resource, .. } => {
                        resource.set_shared_allocation(allocation.clone());
                    }
                }
            }
            return Err(Error::Bind(bind_error.error));
        }

        for member in self.members {
            match member {
                Member::Buffer { resource, .. } => {
                    resource.set_shared_allocation(allocation.clone());
                }
                Member::Image { resource, .. } => {
                    resource.set_shared_allocation(allocation.clone());
                }
            }
        }

        // `allocation` drops here. Only successfully bound resources retain it.
        Ok(())
    }

    fn add_member_type_check(&self, requirements: &MemoryRequirements) -> Result<()> {
        if requirements.requires_dedicated_allocation {
            return Err(Error::RequiresDedicatedAllocation);
        }
        if requirements.memory_type_bits & (1 << self.memory_type_index) == 0 {
            return Err(Error::GroupMemoryTypeMismatch {
                memory_type_index: self.memory_type_index,
            });
        }
        Ok(())
    }

    fn next_member_offset(&self, requirements: &MemoryRequirements) -> Result<u64> {
        if requirements.alignment == 0 {
            return Err(Error::InvalidGroupRequirements);
        }
        let alignment = checked_lcm(requirements.alignment, self.alignment)?;
        let offset = (self
            .next_offset
            .checked_add(alignment - 1)
            .ok_or(Error::OffsetOverflow)?
            / alignment)
            .checked_mul(alignment)
            .ok_or(Error::OffsetOverflow)?;
        let end = offset
            .checked_add(requirements.size)
            .ok_or(Error::OffsetOverflow)?;
        if end > self.size {
            return Err(Error::GroupCapacityExceeded {
                requested_end: end,
                capacity: self.size,
            });
        }
        Ok(offset)
    }
}

fn checked_lcm(lhs: u64, rhs: u64) -> Result<u64> {
    let gcd = {
        let (mut lhs, mut rhs) = (lhs, rhs);
        while rhs != 0 {
            (lhs, rhs) = (rhs, lhs % rhs);
        }
        lhs
    };
    lhs.checked_div(gcd)
        .and_then(|lhs| lhs.checked_mul(rhs))
        .ok_or(Error::OffsetOverflow)
}

/// A resource accepted by [`MemoryGroup::add`].
pub enum MemoryMember<'a> {
    Buffer(&'a mut Buffer),
    Image(&'a mut Image),
}