use super::block::Block;

pub struct BlockAccessor<'a> {
    data: &'a mut [u8],
    m_uv_len: usize,
    m_u_len: usize,
    m_v_len: usize,
    num_block_pairs: usize,
}

const BIG_CHUNK_SIZE: usize = 8 * 32;

impl<'a> BlockAccessor<'a> {
    pub fn new(message: &'a mut [u8]) -> Self {
        let num_block_pairs = (message.len() - 16 - 16) / 32;
        let m_uv_len = (message.len() % 32) * 8;
        Self {
            data: message,
            m_uv_len,
            m_u_len: 128.min(m_uv_len),
            m_v_len: m_uv_len.saturating_sub(128),
            num_block_pairs,
        }
    }

    pub fn m_uv_len(&self) -> usize {
        self.m_uv_len
    }

    fn suffix_start(&self) -> usize {
        self.num_block_pairs * 32
    }

    pub fn m_u(&self) -> Block {
        let start = self.suffix_start();
        Block::from_slice(&self.data[start..start + self.m_u_len / 8])
    }

    pub fn set_m_u(&mut self, m_u: Block) {
        let start = self.suffix_start();
        self.data[start..start + self.m_u_len / 8]
            .copy_from_slice(&m_u.bytes()[..self.m_u_len / 8]);
    }

    pub fn m_v(&self) -> Block {
        let start = self.suffix_start();
        Block::from_slice(&self.data[start + self.m_u_len / 8..start + self.m_uv_len / 8])
    }

    pub fn set_m_v(&mut self, m_v: Block) {
        let start = self.suffix_start();
        self.data[start + self.m_u_len / 8..start + self.m_uv_len / 8]
            .copy_from_slice(&m_v.bytes()[..self.m_v_len / 8]);
    }

    pub fn m_x(&self) -> Block {
        let start = self.suffix_start() + self.m_uv_len / 8;
        Block::from_slice(&self.data[start..start + 16])
    }

    pub fn set_m_x(&mut self, m_x: Block) {
        let start = self.suffix_start() + self.m_uv_len / 8;
        self.data[start..start + 16].copy_from_slice(&m_x.bytes());
    }

    pub fn m_y(&self) -> Block {
        let start = self.suffix_start() + self.m_uv_len / 8;
        Block::from_slice(&self.data[start + 16..start + 32])
    }

    pub fn set_m_y(&mut self, m_y: Block) {
        let start = self.suffix_start() + self.m_uv_len / 8;
        self.data[start + 16..start + 32].copy_from_slice(&m_y.bytes());
    }

    pub fn pairs_mut(&mut self) -> impl Iterator<Item = (&mut [u8; 16], &mut [u8; 16])> {
        let stop = self.suffix_start();
        self.data[..stop]
            .chunks_exact_mut(32)
            .map(move |x| x.split_at_mut(16))
            .map(move |(x, y)| (x.try_into().unwrap(), y.try_into().unwrap()))
    }

    pub fn suffix_8_mut(&mut self) -> impl Iterator<Item = (&mut [u8; 16], &mut [u8; 16])> {
        let start = self.suffix_start() / BIG_CHUNK_SIZE * BIG_CHUNK_SIZE;
        let stop = self.suffix_start();
        self.data[start..stop]
            .chunks_exact_mut(32)
            .map(move |x| x.split_at_mut(16))
            .map(move |(x, y)| (x.try_into().unwrap(), y.try_into().unwrap()))
    }

    pub fn pairs_8_mut(
        &mut self,
    ) -> impl Iterator<Item = ([&mut [u8; 16]; 8], [&mut [u8; 16]; 8])> {
        let stop = self.suffix_start() / BIG_CHUNK_SIZE * BIG_CHUNK_SIZE;
        self.data[..stop]
            .chunks_exact_mut(BIG_CHUNK_SIZE)
            .map(move |x| {
                let (b0, b1) = x.split_at_mut(BIG_CHUNK_SIZE / 2);
                let (b00, b01) = b0.split_at_mut(BIG_CHUNK_SIZE / 4);
                let (b10, b11) = b1.split_at_mut(BIG_CHUNK_SIZE / 4);
                let (b000, b001) = b00.split_at_mut(BIG_CHUNK_SIZE / 8);
                let (b010, b011) = b01.split_at_mut(BIG_CHUNK_SIZE / 8);
                let (b100, b101) = b10.split_at_mut(BIG_CHUNK_SIZE / 8);
                let (b110, b111) = b11.split_at_mut(BIG_CHUNK_SIZE / 8);
                let (b0000, b0001) = b000.split_at_mut(16);
                let (b0010, b0011) = b001.split_at_mut(16);
                let (b0100, b0101) = b010.split_at_mut(16);
                let (b0110, b0111) = b011.split_at_mut(16);
                let (b1000, b1001) = b100.split_at_mut(16);
                let (b1010, b1011) = b101.split_at_mut(16);
                let (b1100, b1101) = b110.split_at_mut(16);
                let (b1110, b1111) = b111.split_at_mut(16);
                (
                    [
                        b0000.try_into().unwrap(),
                        b0010.try_into().unwrap(),
                        b0100.try_into().unwrap(),
                        b0110.try_into().unwrap(),
                        b1000.try_into().unwrap(),
                        b1010.try_into().unwrap(),
                        b1100.try_into().unwrap(),
                        b1110.try_into().unwrap(),
                    ],
                    [
                        b0001.try_into().unwrap(),
                        b0011.try_into().unwrap(),
                        b0101.try_into().unwrap(),
                        b0111.try_into().unwrap(),
                        b1001.try_into().unwrap(),
                        b1011.try_into().unwrap(),
                        b1101.try_into().unwrap(),
                        b1111.try_into().unwrap(),
                    ],
                )
            })
    }
}