use crate::{BufferAddress, Extent3d, TexelCopyBufferLayout, TextureAspect, TextureFormat};

impl TexelCopyBufferLayout {
    /// Extract a variety of information about the given copy operation.
    ///
    /// Returns an error if the size of the copy overflows a `u64`, or if the arguments are
    /// not valid in conjunction with the `bytes_per_row` or `rows_per_image` parameters in
    /// `self`.
    ///
    /// This is public for use by `wgpu-core` and `wgpu-hal`, it is not a stable API.
    ///
    /// Although WebGPU requires that `bytes_per_row` and `rows_per_image` be specified in
    /// cases where they apply, we are more lenient here (although it's not clear if that is
    /// necessary). Our caller, `validate_linear_texture_data`, enforces this and other
    /// WebGPU requirements on the copy parameters that we do not check here.
    #[doc(hidden)]
    #[inline(always)]
    pub fn get_buffer_texture_copy_info(
        &self,
        format: TextureFormat,
        aspect: TextureAspect,
        copy_size: &Extent3d,
    ) -> Result<BufferTextureCopyInfo, Error> {
        let copy_width = BufferAddress::from(copy_size.width);
        let copy_height = BufferAddress::from(copy_size.height);
        let depth_or_array_layers = BufferAddress::from(copy_size.depth_or_array_layers);

        let block_size_bytes = BufferAddress::from(format.block_copy_size(Some(aspect)).unwrap());
        let (block_width, block_height) = format.block_dimensions();
        let block_width_texels = BufferAddress::from(block_width);
        let block_height_texels = BufferAddress::from(block_height);

        let width_blocks = copy_width.div_ceil(block_width_texels);
        let height_blocks = copy_height.div_ceil(block_height_texels);

        // The spec calls this bytesInLastRow.
        let row_bytes_dense = width_blocks * block_size_bytes;
        let row_stride_bytes = match self.bytes_per_row.map(BufferAddress::from) {
            Some(bytes_per_row) if bytes_per_row >= row_bytes_dense => bytes_per_row,
            Some(_) => return Err(Error::InvalidBytesPerRow),
            None => row_bytes_dense,
        };

        let image_rows_dense = height_blocks;
        let image_stride_rows = match self.rows_per_image.map(BufferAddress::from) {
            Some(rows_per_image) if rows_per_image >= image_rows_dense => rows_per_image,
            Some(_) => return Err(Error::InvalidRowsPerImage),
            None => image_rows_dense,
        };

        let image_bytes_dense = match image_rows_dense.checked_sub(1) {
            Some(rows_minus_one) => rows_minus_one
                .checked_mul(row_stride_bytes)
                .ok_or(Error::ImageBytesOverflow(false))?
                .checked_add(row_bytes_dense)
                .ok_or(Error::ImageBytesOverflow(true))?,
            None => 0,
        };

        // It is possible that `image_stride_bytes` overflows, but the actual
        // copy size does not, when the copy only has a single layer and
        // `image_size_bytes` is not used. We don't worry about handling this
        // gracefully because WebGPU texture size limits should keep things out
        // of this realm entirely.
        let image_stride_bytes = row_stride_bytes
            .checked_mul(image_stride_rows)
            .ok_or(Error::ImageStrideOverflow)?;

        let bytes_in_copy = if depth_or_array_layers <= 1 {
            depth_or_array_layers * image_bytes_dense
        } else {
            (depth_or_array_layers - 1)
                .checked_mul(image_stride_bytes)
                .ok_or(Error::ArraySizeOverflow(false))?
                .checked_add(image_bytes_dense)
                .ok_or(Error::ArraySizeOverflow(true))?
        };

        Ok(BufferTextureCopyInfo {
            copy_width,
            copy_height,
            depth_or_array_layers,

            offset: self.offset,

            block_size_bytes,
            block_width_texels,
            block_height_texels,

            width_blocks,
            height_blocks,

            row_bytes_dense,
            row_stride_bytes,

            image_stride_rows,
            image_stride_bytes,

            image_rows_dense,
            image_bytes_dense,

            bytes_in_copy,
        })
    }
}

/// Information about a copy between a buffer and a texture.
///
/// Mostly used for internal calculations, but useful nonetheless.
/// Generated by [`TexelCopyBufferLayout::get_buffer_texture_copy_info`].
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub struct BufferTextureCopyInfo {
    /// The width of the copy region in pixels.
    pub copy_width: u64,
    /// The height of the copy region in pixels.
    pub copy_height: u64,
    /// The depth of the copy region in pixels.
    pub depth_or_array_layers: u64,

    /// The offset in the buffer where the copy starts.
    pub offset: u64,

    /// The size of a single texture texel block in bytes.
    pub block_size_bytes: u64,
    /// The number of texel in a texel block in the x direction.
    pub block_width_texels: u64,
    /// The number of texel in a texel block in the y direction.
    pub block_height_texels: u64,

    /// The width of the copy region in blocks.
    pub width_blocks: u64,
    /// The height of the copy region in blocks.
    pub height_blocks: u64,

    /// The number of bytes in the last row of the copy region.
    pub row_bytes_dense: u64,
    /// The stride in bytes between the start of one row in an image and the next row in the same image.
    ///
    /// This includes any padding between one row and the next row.
    pub row_stride_bytes: u64,

    /// The stride in rows between the start of one image and the next image.
    pub image_stride_rows: u64,
    /// The stride in bytes between the start of one image and the next image.
    pub image_stride_bytes: u64,

    /// The number of rows in a densely packed list of images.
    ///
    /// This is the number of rows in the image that are actually used for texel data,
    /// and does not include any padding rows, unlike `image_stride_rows`.
    pub image_rows_dense: u64,
    /// The number of bytes in a densely packed list of images.
    ///
    /// This is the number of bytes in the image that are actually used for texel data,
    /// or are used for padding between _rows_. Padding at the end of the last row and
    /// between _images_ is not included.
    pub image_bytes_dense: u64,

    /// The total number of bytes in the copy region.
    ///
    /// This includes all padding except the padding after the last row in the copy.
    pub bytes_in_copy: u64,
}

/// Errors that can occur while populating `BufferTextureCopyInfo`.
//
// We use the additional detail provided by these errors (over wgpu-core's
// `TransferError`) to improve the reliability of the tests in this module. It
// doesn't seem worth plumbing them upwards, because at the API level it
// shouldn't be possible to exceed them without exceeding the WebGPU limits on
// texture dimension. But the WebGPU limits are not currently enforced, so we
// have to do something here to protect against overflows.
//
// Even when the WebGPU limits are enforced, it may still be useful to keep the
// checks here as a failsafe if the correctness of the primary limit enforcement
// is not immediately apparent.
#[derive(Clone, Copy, Debug, Eq, PartialEq)]
pub enum BufferTextureCopyInfoError {
    /// The `bytes_per_row` is too small for the texture width.
    InvalidBytesPerRow,
    /// The `rows_per_image` is too small for the texture height.
    InvalidRowsPerImage,
    /// The image stride overflows a `u64`.
    ImageStrideOverflow,
    /// The last-layer byte size overflows a `u64`.
    ///
    /// The bool value indicates whether the multiplication (false) or the
    /// addition (true) overflowed.
    ImageBytesOverflow(bool),
    /// The total size of the copy overflows a `u64`.
    ///
    /// The bool value indicates whether the multiplication (false) or the
    /// addition (true) overflowed.
    ArraySizeOverflow(bool),
}
type Error = BufferTextureCopyInfoError;

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

    #[derive(Clone)]
    struct LTDTest {
        layout: TexelCopyBufferLayout,
        format: TextureFormat,
        aspect: TextureAspect,
        copy_size: Extent3d,
        expected_result: BufferTextureCopyInfo,
        // Normally a Result<BufferTextureCopyInfo, Error> would be make sense,
        // but since the existing tests were written to mutate
        // `LTDTest.expected_result`, keeping this separate avoids a bunch of
        // `unwrap`s.
        expected_error: Option<Error>,
    }

    impl LTDTest {
        #[track_caller]
        fn run(&self) {
            let linear_texture_data =
                self.layout
                    .get_buffer_texture_copy_info(self.format, self.aspect, &self.copy_size);
            let expected = match self.expected_error {
                Some(err) => Err(err),
                None => Ok(self.expected_result),
            };
            assert_eq!(linear_texture_data, expected);
        }
    }

    #[test]
    fn linear_texture_data_1d_copy() {
        let mut test = LTDTest {
            layout: TexelCopyBufferLayout {
                offset: 0,
                bytes_per_row: None,
                rows_per_image: None,
            },
            format: TextureFormat::Rgba8Unorm,
            aspect: TextureAspect::All,
            copy_size: Extent3d {
                width: 4,
                height: 1,
                depth_or_array_layers: 1,
            },
            expected_result: BufferTextureCopyInfo {
                copy_width: 4,
                copy_height: 1,
                depth_or_array_layers: 1,
                offset: 0,
                block_size_bytes: 4,
                block_width_texels: 1,
                block_height_texels: 1,
                width_blocks: 4,
                height_blocks: 1,
                row_bytes_dense: 16,
                row_stride_bytes: 16,
                image_stride_rows: 1,
                image_stride_bytes: 16,
                image_rows_dense: 1,
                image_bytes_dense: 16,
                bytes_in_copy: 16,
            },
            expected_error: None,
        };

        test.run();

        // Changing bytes_per_row should only change the bytes_per_row, not the bytes_in_copy
        // as that is only affected by the last row size.
        test.layout.bytes_per_row = Some(32);
        test.expected_result.row_stride_bytes = 32;
        test.expected_result.image_stride_bytes = 32;

        test.run();

        // Changing rows_per_image should only change the rows_per_image and bytes_per_image, nothing else
        test.layout.rows_per_image = Some(4);
        test.expected_result.image_stride_bytes = 128; // 32 * 4
        test.expected_result.image_stride_rows = 4;

        test.run();

        // Changing the offset should change nothing.

        test.layout.offset = 4;
        test.expected_result.offset = 4;

        test.run();
    }

    #[test]
    fn linear_texture_data_2d_3d_copy() {
        let template = LTDTest {
            layout: TexelCopyBufferLayout {
                offset: 0,
                bytes_per_row: None,
                rows_per_image: None,
            },
            format: TextureFormat::Rgba8Unorm,
            aspect: TextureAspect::All,
            copy_size: Extent3d {
                width: 7,
                height: 12,
                depth_or_array_layers: 1,
            },
            expected_result: BufferTextureCopyInfo {
                copy_width: 7,
                copy_height: 12,
                depth_or_array_layers: 1,
                offset: 0,
                block_size_bytes: 4,
                block_width_texels: 1,
                block_height_texels: 1,
                width_blocks: 7,
                height_blocks: 12,
                row_bytes_dense: 4 * 7,
                row_stride_bytes: 4 * 7,
                image_stride_rows: 12,
                image_stride_bytes: 4 * 7 * 12,
                image_rows_dense: 12,
                image_bytes_dense: 4 * 7 * 12,
                bytes_in_copy: 4 * 7 * 12,
            },
            expected_error: None,
        };

        let mut test = template.clone();
        test.run();

        // Changing bytes_per_row changes a number of other properties.
        test.layout.bytes_per_row = Some(48);
        test.expected_result.row_stride_bytes = 48;
        test.expected_result.image_stride_bytes = 48 * 12;
        test.expected_result.image_bytes_dense = 48 * 11 + (4 * 7);
        test.expected_result.bytes_in_copy = 48 * 11 + (4 * 7);
        test.run();

        // Making this a 3D copy only changes the depth_or_array_layers and the bytes_in_copy.
        test.copy_size.depth_or_array_layers = 4;
        test.expected_result.depth_or_array_layers = 4;
        test.expected_result.bytes_in_copy = 48 * 12 * 3 + 48 * 11 + (4 * 7); // 4 layers
        test.run();

        // Changing rows_per_image
        test.layout.rows_per_image = Some(20);
        test.expected_result.image_stride_rows = 20;
        test.expected_result.image_stride_bytes = 20 * test.expected_result.row_stride_bytes;
        test.expected_result.bytes_in_copy = 48 * 20 * 3 + 48 * 11 + (4 * 7); // 4 layers
        test.run();

        // Invalid because the row stride is too small.
        let mut test = template.clone();
        test.layout.bytes_per_row = Some(20);
        test.expected_error = Some(Error::InvalidBytesPerRow);
        test.run();

        // Invalid because the image stride is too small.
        let mut test = template.clone();
        test.layout.rows_per_image = Some(8);
        test.expected_error = Some(Error::InvalidRowsPerImage);
        test.run();

        // Invalid because width * height * texel_size_bytes overflows.
        let mut test = template.clone();
        test.copy_size.width = u32::MAX;
        test.copy_size.height = u32::MAX;
        test.expected_error = Some(Error::ImageBytesOverflow(false));
        test.run();

        // Invalid because the addition of row_bytes_dense overflows.
        // (But the product rows_minus_one * row_stride_bytes does not overflow.)
        let mut test = template.clone();
        test.copy_size.width = 0x8000_0000;
        test.copy_size.height = 0x8000_0000;
        test.expected_error = Some(Error::ImageBytesOverflow(true));
        test.run();

        // Invalid because image_stride_bytes overflows.
        let mut test = template.clone();
        test.copy_size.width = 0x8000_0000;
        test.layout.rows_per_image = Some(0x8000_0000);
        test.expected_result.image_stride_rows = 0x8000_0000;
        test.expected_error = Some(Error::ImageStrideOverflow);
        test.run();

        // Invalid because (layers - 1) * image_stride_bytes overflows.
        let mut test = template.clone();
        test.copy_size.depth_or_array_layers = 0x8000_0000;
        test.copy_size.width = 0x1_0000;
        test.copy_size.height = 0x1_0000;
        test.expected_error = Some(Error::ArraySizeOverflow(false));
        test.run();

        // Invalid because the total size of the copy overflows (but the product
        // (layers - 1) * image_stride_bytes does not overflow).
        let mut test = template.clone();
        test.copy_size.depth_or_array_layers = 0x3fff_8001;
        test.copy_size.width = 0x1_0001;
        test.copy_size.height = 0x1_0001;
        test.expected_error = Some(Error::ArraySizeOverflow(true));
        test.run();
    }

    #[test]
    fn linear_texture_data_2d_3d_compressed_copy() {
        let mut test = LTDTest {
            layout: TexelCopyBufferLayout {
                offset: 0,
                bytes_per_row: None,
                rows_per_image: None,
            },
            format: TextureFormat::Bc1RgbaUnorm,
            aspect: TextureAspect::All,
            copy_size: Extent3d {
                width: 7,
                height: 13,
                depth_or_array_layers: 1,
            },
            expected_result: BufferTextureCopyInfo {
                copy_width: 7,
                copy_height: 13,
                depth_or_array_layers: 1,
                offset: 0,
                block_size_bytes: 8,
                block_width_texels: 4,
                block_height_texels: 4,
                width_blocks: 2,
                height_blocks: 4,
                row_bytes_dense: 8 * 2, // block size * width_blocks
                row_stride_bytes: 8 * 2,
                image_stride_rows: 4,
                image_stride_bytes: 8 * 2 * 4, // block size * width_blocks * height_blocks
                image_rows_dense: 4,
                image_bytes_dense: 8 * 2 * 4,
                bytes_in_copy: 8 * 2 * 4,
            },
            expected_error: None,
        };

        test.run();

        // Changing bytes_per_row.
        test.layout.bytes_per_row = Some(48);
        test.expected_result.row_stride_bytes = 48;
        test.expected_result.image_stride_bytes = 48 * 4;
        test.expected_result.image_bytes_dense = 48 * 3 + (8 * 2);
        test.expected_result.bytes_in_copy = 48 * 3 + (8 * 2);

        test.run();

        // Changing rows_per_image.
        test.layout.rows_per_image = Some(8);
        test.expected_result.image_stride_bytes = 48 * 8;
        test.expected_result.image_stride_rows = 8;

        test.run();

        // Making this a 3D copy only changes the depth_or_array_layers and the bytes_in_copy.
        test.copy_size.depth_or_array_layers = 4;
        test.expected_result.depth_or_array_layers = 4;
        test.expected_result.bytes_in_copy = 48 * 8 * 3 + 48 * 3 + (8 * 2); // 4 layers

        test.run();
    }
}
