use crate::{BitBuf, Error, Result, StorageMut};

impl<S: StorageMut> BitBuf<S> {
  /// Write a [`u8`] in BE-bit-order at `byte_offset` without performing bound checks
  ///
  /// # Safety
  ///
  /// * This is UB if [`byte_offset >= self.bytes().len()`][Self::bytes]
  ///
  /// # Panics
  ///
  /// * Panics in debug mode if [`byte_offset >= self.bytes().len()`][Self::bytes]
  #[inline(always)]
  pub unsafe fn write_u8_be_aligned_full_at_unchecked(
    &mut self,
    byte_offset: usize,
    v: u8,
  ) -> &mut Self {
    let bytes = self.bytes_mut();

    debug_assert!(
      byte_offset < bytes.len(),
      "BitBuf::write_u8_be_aligned_full_at_unchecked: index out of bounds! len is {}, offset is {}",
      bytes.len(),
      byte_offset,
    );

    unsafe { *bytes.get_unchecked_mut(byte_offset) = v };
    self
  }

  /// Write a [`u8`] in BE-bit-order without performing bound checks, advancing the internal cursor
  ///
  /// # Safety
  ///
  /// * The internal cursor must be byte-aligned ([`self.is_aligned()`][Self::is_aligned])
  /// * This is UB if the internal cursor points past the end of storage (<code>[self.byte_pos()][Self::byte_pos] >= [self.bytes().len()][Self::bytes]</code>)
  ///
  /// # Panics
  ///
  /// * Panics in debug mode if the internal cursor is not byte-aligned ([`!self.is_aligned()`][Self::is_aligned])
  /// * Panics in debug mode if the internal cursor points past the end of storage (<code>[self.byte_pos()][Self::byte_pos] >= [self.bytes().len()][Self::bytes]</code>)
  #[inline(always)]
  pub unsafe fn write_u8_be_aligned_full_unchecked(&mut self, v: u8) -> &mut Self {
    debug_assert!(
      self.is_aligned(),
      "BitBuf::write_u8_be_aligned_full_unchecked called at unaligned bit position: {}",
      self.pos(),
    );

    unsafe { self.write_u8_be_aligned_full_at_unchecked(self.byte_pos(), v) };
    self.advance_bytes(1);
    self
  }

  /// Write a [`u8`] in BE-bit-order at `offset` without performing bound checks
  ///
  /// # Safety
  ///
  /// * This is UB if <code>(offset / 8) >= [`self.bytes().len()`][Self::bytes]</code>
  /// * This is UB if <code>(offset % 8 != 0) && (offset / 8) + 1 >= [`self.bytes().len()`][Self::bytes]</code>
  ///
  /// # Panics
  ///
  /// * Panics in debug mode if <code>(offset / 8) >= [`self.bytes().len()`][Self::bytes]</code>
  /// * Panics in debug mode if <code>(offset % 8 != 0) && (offset / 8) + 1 >= [`self.bytes().len()`][Self::bytes]</code>
  #[inline(always)]
  pub unsafe fn write_u8_be_full_at_unchecked(&mut self, offset: usize, v: u8) -> &mut Self {
    let byte_idx = offset / 8;
    let shift = offset % 8;

    let bytes = self.bytes_mut();

    debug_assert!(
      byte_idx < bytes.len(),
      "BitBuf::write_u8_be_full_at_unchecked: index out of bounds! len is {}, byte_idx is {}",
      bytes.len(),
      byte_idx,
    );

    if shift == 0 {
      unsafe { self.write_u8_be_aligned_full_at_unchecked(byte_idx, v) };
    } else {
      let next_idx = byte_idx + 1;

      debug_assert!(
        next_idx < bytes.len(),
        "BitBuf::write_u8_be_full_at_unchecked: lookahead index out of bounds! len is {}, lookahead byte_idx is {}",
        bytes.len(),
        next_idx,
      );

      // first byte, higher part of value
      unsafe {
        // mask to keep high bits of the current byte
        let mask = !0 << (8 - shift);
        let b = bytes.get_unchecked(byte_idx) & mask;
        // high bits of the value, shifted down
        let new = v >> shift;
        *bytes.get_unchecked_mut(byte_idx) = b | new;
      }

      // second byte, lower part of value
      unsafe {
        // mask to keep low bits of the current byte
        let mask = !0 >> shift;
        let b = bytes.get_unchecked(next_idx) & mask;
        // low bits of the value, shifted up
        let new = v << (8 - shift);
        *bytes.get_unchecked_mut(next_idx) = b | new;
      }
    }

    self
  }

  /// Write a BE-bit-order [`u8`] at `offset` while performing bound checks
  ///
  /// # Errors
  ///
  /// * Returns [`Error::OutOfBounds`]
  ///   * if <code>(offset / 8) >= [`self.bytes().len()`][Self::bytes]</code>
  ///   * if <code>(offset % 8 != 0) && (offset / 8) + 1 >= [`self.bytes().len()`][Self::bytes]</code>
  #[inline(always)]
  pub fn try_write_u8_be_full_at(&mut self, offset: usize, v: u8) -> Result<&mut Self> {
    let byte_idx = offset / 8;
    let shift = offset % 8;

    let len = self.bytes().len();

    if byte_idx >= len {
      return Err(Error::OutOfBounds);
    }

    if shift != 0 && byte_idx + 1 >= len {
      return Err(Error::OutOfBounds);
    }

    Ok(unsafe { self.write_u8_be_full_at_unchecked(offset, v) })
  }

  /// Write a BE-bit-order [`u8`] at `offset`, panicking on out of bounds
  ///
  /// # Panics
  ///
  /// * Panics if <code>(offset / 8) >= [`self.bytes().len()`][Self::bytes]</code>
  /// * Panics if <code>(offset % 8 != 0) && (offset / 8) + 1 >= [`self.bytes().len()`][Self::bytes]</code>
  #[inline(always)]
  pub fn write_u8_be_full_at(&mut self, offset: usize, v: u8) -> &mut Self {
    self
      .try_write_u8_be_full_at(offset, v)
      .expect("BitBuf::write_u8_be_full_at out of bounds")
  }
}
