//! Optimized implementations for `wasm32`.

use core::arch::wasm32::*;
use core::mem::MaybeUninit;

use crate::backend::generic::encode_generic_unchecked;
use crate::util::lut16;

#[target_feature(enable = "simd128")]
/// ## Safety
///
/// We assume that:
///
/// 1. `simd128` instructions are supported.
/// 2. `src.len() <= dst.len()`.
pub(crate) unsafe fn encode_simd128_unchecked<const UPPER: bool>(
    mut src: &[u8],
    mut dst: &mut [[MaybeUninit<u8>; 2]],
) {
    /// The batch size in bytes.
    const BATCH: usize = size_of::<v128>();

    if src.len() >= BATCH {
        let lut = v128_load(lut16::<UPPER>().as_ptr().cast());
        let and = u8x16_splat(0b1111);

        while src.len() >= BATCH {
            let invec = v128_load(src.as_ptr().cast());

            // Split each byte into nibbles.
            let mut lo = v128_and(invec, and);
            let mut hi = u8x16_shr(invec, 4);

            // Lookup the ASCII values for each nibble.
            lo = u8x16_swizzle(lut, lo);
            hi = u8x16_swizzle(lut, hi);

            // Interleave the nibbles ([hi[0], lo[0], hi[1], lo[1], ...]).
            let hex_lo = u8x16_shuffle::<
                0x00,
                0x10,
                0x01,
                0x11,
                0x02,
                0x12,
                0x03,
                0x13,
                0x04,
                0x14,
                0x05,
                0x15,
                0x06,
                0x16,
                0x07,
                0x17,
            >(hi, lo);
            let hex_hi = u8x16_shuffle::<
                0x08,
                0x18,
                0x09,
                0x19,
                0x0A,
                0x1A,
                0x0B,
                0x1B,
                0x0C,
                0x1C,
                0x0D,
                0x1D,
                0x0E,
                0x1E,
                0x0F,
                0x1F,
            >(hi, lo);

            // Finally, store the result.
            {
                let dst = dst.as_mut_ptr().cast::<v128>();
                v128_store(dst, hex_lo);
                v128_store(dst.add(1), hex_hi);
            }

            src = &src[BATCH..];
            dst = dst.get_unchecked_mut(BATCH..);
        }
    }

    encode_generic_unchecked::<UPPER>(src, dst);
}

#[cfg(test)]
mod smoking {
    use alloc::string::String;
    use alloc::vec;
    use core::mem::MaybeUninit;
    use core::{slice, str};

    use super::*;
    use crate::util::{HEX_CHARS_LOWER, HEX_CHARS_UPPER};

    macro_rules! test {
        (
            Encode = $encode_f:ident;
            Validate = $($validate_f:ident),*;
            Decode = $($decode_f:ident),*;
            Case = $i:expr
        ) => {{
            let input = $i;

            let expected_lower = input
                .iter()
                .flat_map(|b| [
                    HEX_CHARS_LOWER[(*b >> 4) as usize] as char,
                    HEX_CHARS_LOWER[(*b & 0b1111) as usize] as char,
                ])
                .collect::<String>();
            let expected_upper = input
                .iter()
                .flat_map(|b| [
                    HEX_CHARS_UPPER[(*b >> 4) as usize] as char,
                    HEX_CHARS_UPPER[(*b & 0b1111) as usize] as char,
                ])
                .collect::<String>();

            let mut output_lower = vec![[MaybeUninit::<u8>::uninit(); 2]; input.len()];
            let mut output_upper = vec![[MaybeUninit::<u8>::uninit(); 2]; input.len()];

            unsafe {
                $encode_f::<false>(input, &mut output_lower);
                $encode_f::<true>(input, &mut output_upper);
            }

            let output_lower = unsafe {
                slice::from_raw_parts(
                    output_lower.as_ptr().cast::<[u8; 2]>(),
                    output_lower.len(),
                )
            };
            let output_upper = unsafe {
                slice::from_raw_parts(
                    output_upper.as_ptr().cast::<[u8; 2]>(),
                    output_upper.len(),
                )
            };

            assert_eq!(
                output_lower.as_flattened(),
                expected_lower.as_bytes(),
                "Encode error, expect \"{expected_lower}\", got \"{}\" ({:?})",
                str::from_utf8(output_lower.as_flattened()).unwrap_or("<invalid utf-8>"),
                output_lower.as_flattened()
            );
            assert_eq!(
                output_upper.as_flattened(),
                expected_upper.as_bytes(),
                "Encode error, expect \"{expected_upper}\", got \"{}\" ({:?})",
                str::from_utf8(output_upper.as_flattened()).unwrap_or("<invalid utf-8>"),
                output_upper.as_flattened()
            );

            $(
                unsafe {
                    $validate_f(output_lower)
                        .unwrap_or_else(|_| panic!("validation failed for {}", stringify!($validate_f)));
                    $validate_f(output_upper)
                        .unwrap_or_else(|_| panic!("validation failed for {}", stringify!($validate_f)));
                }
            )*

            $({
                let mut decoded_lower = vec![MaybeUninit::<u8>::uninit(); input.len()];
                let mut decoded_upper = vec![MaybeUninit::<u8>::uninit(); input.len()];

                unsafe {
                    $decode_f(output_lower, &mut decoded_lower);
                    $decode_f(output_upper, &mut decoded_upper);

                    assert_eq!(
                        decoded_lower.assume_init_ref(),
                        input,
                        "Decode error for {}, expect {:?}, got {:?}",
                        stringify!($decode_f),
                        input,
                        decoded_lower.assume_init_ref()
                    );
                    assert_eq!(
                        decoded_upper.assume_init_ref(),
                        input,
                        "Decode error for {}, expect {:?}, got {:?}",
                        stringify!($decode_f),
                        input,
                        decoded_upper.assume_init_ref()
                    );
                }
            })*
        }};
    }

    #[test]
    #[cfg_attr(miri, ignore)]
    fn test_simd128() {
        const CASE: &[u8; 17] = &[
            0x12, 0x77, 0x4C, 0x16, 0x16, 0x2B, 0x99, 0x97, 0x37, 0x62, 0x24, 0x24, 0x36, 0x83,
            0xA4, 0xF1, 0xDD,
        ];

        test! {
            Encode = encode_simd128_unchecked;
            Validate = ;
            Decode = ;
            Case = &[]
        }

        test! {
            Encode = encode_simd128_unchecked;
            Validate = ;
            Decode = ;
            Case = &CASE[..15]
        }

        test! {
            Encode = encode_simd128_unchecked;
            Validate = ;
            Decode = ;
            Case = &CASE[..16]
        }

        test! {
            Encode = encode_simd128_unchecked;
            Validate = ;
            Decode = ;
            Case = &CASE[..17]
        };
    }
}
