#![allow(clippy::missing_const_for_fn)]

use crate::types::Bitboard;

use std::arch::x86_64::{
    __m256i, _mm256_add_epi64, _mm256_and_si256, _mm256_andnot_si256, _mm256_broadcastsi128_si256,
    _mm256_castsi128_si256, _mm256_castsi256_si128, _mm256_cmpeq_epi64, _mm256_extracti128_si256,
    _mm256_inserti128_si256, _mm256_or_si256, _mm256_set1_epi64x, _mm256_setr_epi8,
    _mm256_setzero_si256, _mm256_shuffle_epi8, _mm256_sll_epi64, _mm256_srl_epi64,
    _mm256_sub_epi64, _mm256_testz_si256, _mm256_unpackhi_epi64, _mm256_unpacklo_epi64,
    _mm256_xor_si256, _mm_cvtsi64_si128, _mm_or_si128,
};

#[derive(Copy, Clone)]
#[repr(C, align(32))]
union U64x4Repr {
    p: [u64; 4],
    m: __m256i,
}

#[derive(Copy, Clone)]
#[repr(align(32))]
pub struct U64x4 {
    repr: U64x4Repr,
}

impl PartialEq for U64x4 {
    fn eq(&self, other: &Self) -> bool {
        let neq = unsafe { _mm256_xor_si256(self.as_m256(), other.as_m256()) };
        unsafe { _mm256_testz_si256(neq, neq) != 0 }
    }
}

impl Eq for U64x4 {}

impl std::fmt::Debug for U64x4 {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        f.debug_tuple("U64x4").field(&self.parts()).finish()
    }
}

impl U64x4 {
    #[must_use]
    pub const fn new(p0: u64, p1: u64, p2: u64, p3: u64) -> Self {
        Self { repr: U64x4Repr { p: [p0, p1, p2, p3] } }
    }

    #[must_use]
    pub const fn parts(&self) -> [u64; 4] {
        unsafe { self.repr.p }
    }

    #[inline]
    pub(crate) fn as_m256(self) -> __m256i {
        unsafe { self.repr.m }
    }

    #[inline]
    pub(crate) fn from_m256(value: __m256i) -> Self {
        Self { repr: U64x4Repr { m: value } }
    }

    #[must_use]
    pub fn from_bitboards(b1: Bitboard, b2: Bitboard) -> Self {
        let mut value = unsafe { _mm256_castsi128_si256(b1.as_m128()) };
        value = unsafe { _mm256_inserti128_si256(value, b2.as_m128(), 1) };
        Self::from_m256(value)
    }

    #[must_use]
    pub fn splat_bitboard(bb: Bitboard) -> Self {
        let value = unsafe { _mm256_broadcastsi128_si256(bb.as_m128()) };
        Self::from_m256(value)
    }

    #[must_use]
    pub fn and(self, other: Self) -> Self {
        // SAFETY: AVX2対応ビルドでのみ使用される。
        unsafe { and_simd(self, other) }
    }

    #[must_use]
    pub fn wrapping_add(self, other: Self) -> Self {
        // SAFETY: AVX2対応ビルドでのみ使用される。
        unsafe { add_simd(self, other) }
    }

    #[must_use]
    pub fn wrapping_sub(self, other: Self) -> Self {
        // SAFETY: AVX2対応ビルドでのみ使用される。
        unsafe { sub_simd(self, other) }
    }

    #[must_use]
    pub fn shift_left(self, shift: u32) -> Self {
        // SAFETY: AVX2対応ビルドでのみ使用される。
        unsafe { shl_simd(self, shift) }
    }

    #[must_use]
    pub fn shift_right(self, shift: u32) -> Self {
        // SAFETY: AVX2対応ビルドでのみ使用される。
        unsafe { shr_simd(self, shift) }
    }

    #[must_use]
    pub fn or(self, other: Self) -> Self {
        // SAFETY: AVX2対応ビルドでのみ使用される。
        unsafe { or_simd(self, other) }
    }

    #[must_use]
    pub fn xor(self, other: Self) -> Self {
        // SAFETY: AVX2対応ビルドでのみ使用される。
        unsafe { xor_simd(self, other) }
    }

    #[must_use]
    pub fn and_not(self, other: Self) -> Self {
        // SAFETY: AVX2対応ビルドでのみ使用される。
        unsafe { andnot_simd(self, other) }
    }

    #[must_use]
    pub fn byte_reverse(self) -> Self {
        // SAFETY: AVX2対応ビルドでのみ使用される。
        unsafe { byte_reverse_simd(self) }
    }

    #[must_use]
    pub fn unpack(hi_in: Self, lo_in: Self) -> (Self, Self) {
        // SAFETY: AVX2対応ビルドでのみ使用される。
        unsafe { unpack_simd(hi_in, lo_in) }
    }

    #[must_use]
    pub fn decrement_pair(hi_in: Self, lo_in: Self) -> (Self, Self) {
        // SAFETY: AVX2対応ビルドでのみ使用される。
        unsafe { decrement_pair_simd(hi_in, lo_in) }
    }

    #[must_use]
    pub fn merge(self) -> [u64; 2] {
        // SAFETY: AVX2対応ビルドでのみ使用される。
        unsafe { merge_simd(self) }
    }
}

#[inline]
unsafe fn and_simd(lhs: U64x4, rhs: U64x4) -> U64x4 {
    let l = lhs.as_m256();
    let r = rhs.as_m256();
    U64x4::from_m256(_mm256_and_si256(l, r))
}

#[inline]
unsafe fn add_simd(lhs: U64x4, rhs: U64x4) -> U64x4 {
    let l = lhs.as_m256();
    let r = rhs.as_m256();
    U64x4::from_m256(_mm256_add_epi64(l, r))
}

#[inline]
unsafe fn sub_simd(lhs: U64x4, rhs: U64x4) -> U64x4 {
    let l = lhs.as_m256();
    let r = rhs.as_m256();
    U64x4::from_m256(_mm256_sub_epi64(l, r))
}

#[inline]
unsafe fn shl_simd(value: U64x4, shift: u32) -> U64x4 {
    let v = value.as_m256();
    let count = _mm_cvtsi64_si128(i64::from(shift));
    U64x4::from_m256(_mm256_sll_epi64(v, count))
}

#[inline]
unsafe fn shr_simd(value: U64x4, shift: u32) -> U64x4 {
    let v = value.as_m256();
    let count = _mm_cvtsi64_si128(i64::from(shift));
    U64x4::from_m256(_mm256_srl_epi64(v, count))
}

#[inline]
unsafe fn or_simd(lhs: U64x4, rhs: U64x4) -> U64x4 {
    let l = lhs.as_m256();
    let r = rhs.as_m256();
    U64x4::from_m256(_mm256_or_si256(l, r))
}

#[inline]
unsafe fn xor_simd(lhs: U64x4, rhs: U64x4) -> U64x4 {
    let l = lhs.as_m256();
    let r = rhs.as_m256();
    U64x4::from_m256(_mm256_xor_si256(l, r))
}

#[inline]
unsafe fn andnot_simd(lhs: U64x4, rhs: U64x4) -> U64x4 {
    let l = lhs.as_m256();
    let r = rhs.as_m256();
    U64x4::from_m256(_mm256_andnot_si256(r, l))
}

#[inline]
unsafe fn byte_reverse_simd(value: U64x4) -> U64x4 {
    let v = value.as_m256();
    let shuffle = _mm256_setr_epi8(
        15, 14, 13, 12, 11, 10, 9, 8, 7, 6, 5, 4, 3, 2, 1, 0, 15, 14, 13, 12, 11, 10, 9, 8, 7, 6,
        5, 4, 3, 2, 1, 0,
    );
    U64x4::from_m256(_mm256_shuffle_epi8(v, shuffle))
}

#[inline]
unsafe fn unpack_simd(hi_in: U64x4, lo_in: U64x4) -> (U64x4, U64x4) {
    let hi = hi_in.as_m256();
    let lo = lo_in.as_m256();
    let hi_out = _mm256_unpackhi_epi64(lo, hi);
    let lo_out = _mm256_unpacklo_epi64(lo, hi);
    (U64x4::from_m256(hi_out), U64x4::from_m256(lo_out))
}

#[inline]
unsafe fn decrement_pair_simd(hi_in: U64x4, lo_in: U64x4) -> (U64x4, U64x4) {
    let hi = hi_in.as_m256();
    let lo = lo_in.as_m256();
    let borrow = _mm256_cmpeq_epi64(lo, _mm256_setzero_si256());
    let hi_out = _mm256_add_epi64(hi, borrow);
    let lo_out = _mm256_add_epi64(lo, _mm256_set1_epi64x(-1));
    (U64x4::from_m256(hi_out), U64x4::from_m256(lo_out))
}

#[inline]
unsafe fn merge_simd(value: U64x4) -> [u64; 2] {
    let v = value.as_m256();
    let lo = _mm256_castsi256_si128(v);
    let hi = _mm256_extracti128_si256(v, 1);
    let merged = _mm_or_si128(lo, hi);
    let mut out = [0u64; 2];
    std::arch::x86_64::_mm_storeu_si128(
        out.as_mut_ptr().cast::<std::arch::x86_64::__m128i>(),
        merged,
    );
    out
}
