/// An error which can be emitted when a write fails.
#[derive(Debug, PartialEq)]
pub enum Error {
    /// Reached the end of the buffer.
    UnexpectedEof,
}

impl core::fmt::Display for Error {
    fn fmt(&self, f: &mut core::fmt::Formatter) -> core::fmt::Result {
        match *self {
            Error::UnexpectedEof => write!(f, "unexpected end of the buffer"),
        }
    }
}

impl std::error::Error for Error {}

/// A type alias for a `Result` that returns an `Error`.
pub type Result<T> = core::result::Result<T, Error>;

/// Types which can get written to a `Writer`.
pub trait Pack {
    /// Pack this type into the `Writer`.
    fn pack<T: Target>(&self, writer: &mut Writer<T>) -> Result<()>;

    /// Calculate the size of the data.
    fn size(&self) -> usize {
        let mut writer = Writer::new(());
        self.pack(&mut writer).unwrap(); // in normal operation, this would not happen
        writer.finish().unwrap()
    }
}

/// Pack the type to a buffer directly.
///
/// This is a short-hand for creating a `Writer` then calling `write` on it.
///
/// # Position
///
/// This advances the writer only if all the underlying writes succeed.
///
/// On failure, there may be incomplete data written to the buffer, the writer's position
/// will not include that data.
///
/// # Returns
///
/// On success, this returns an `Ok` with the amount of bytes written.
/// If there is not enough space to write all the data, this returns `Err(Error::UnexpectedEof)`.
///
/// # Examples
///
/// ```rust
/// let mut buf = [0; 6];
///
/// let written = pigeon::pack_buffer(&12345678u32, &mut buf[..]).unwrap();
///
/// assert_eq!(written, 4);
/// assert_eq!(&buf, &[0, 188, 97, 78, 0, 0]);
/// assert_eq!(&buf[0..written], &[0, 188, 97, 78]);
/// ```
pub fn pack_buffer(data: impl Pack, buffer: &mut [u8]) -> Result<usize> {
    let mut writer = Writer::new(buffer);
    writer.write(data)?;
    writer.finish()
}

/// A type which can be written to by a `Writer`.
pub trait Target {
    /// Write these bytes at a specific position.
    fn write_bytes_at(&mut self, pos: usize, bytes: &[u8]);

    /// Get the size of the buffer.
    fn size(&self) -> usize;
}

impl Target for () {
    #[inline(always)]
    fn write_bytes_at(&mut self, _pos: usize, _bytes: &[u8]) {}

    fn size(&self) -> usize {
        usize::MAX
    }
}

impl<'a> Target for &'a mut [u8] {
    fn write_bytes_at(&mut self, pos: usize, bytes: &[u8]) {
        let len = bytes.len();
        self[pos..pos + len].copy_from_slice(bytes);
    }

    fn size(&self) -> usize {
        self.len()
    }
}

/// A type which can be used for writing to a type that implements `Target`.
#[derive(Debug)]
pub struct Writer<T: Target> {
    /// The index the writer is at now
    idx: usize,
    /// The current bit alignment
    bit_align: u8,
    /// Current unaligned bits
    bit_buffer: u16,
    /// The underlying buffer
    buf: T,
}

impl<T: Target> Writer<T> {
    /// Create a new `Writer` from a type implementing `Target`.
    ///
    /// # Examples
    ///
    /// ```rust
    /// # use {
    /// #     pigeon::Writer,
    /// #     std::convert::TryInto,
    /// # };
    /// let mut buf = [0; 8];
    /// let mut writer = Writer::new(&mut buf[..]);
    ///
    /// writer.write(12u8).unwrap();
    /// writer.write(24u8).unwrap();
    /// writer.write(554u16).unwrap();
    /// writer.write(12345678u32).unwrap();
    ///
    /// writer.finish();
    ///
    /// assert_eq!(buf[0], 12);
    /// assert_eq!(buf[1], 24);
    /// assert_eq!(u16::from_be_bytes(buf[2..4].try_into().unwrap()), 554);
    /// assert_eq!(u32::from_be_bytes(buf[4..8].try_into().unwrap()), 12345678);
    /// ```
    pub fn new(buf: T) -> Writer<T> {
        Writer {
            idx: 0,
            bit_align: 0,
            bit_buffer: 0,
            buf,
        }
    }

    /// Get the bit offset that the writer is at.
    ///
    /// This is 0 when the writer is aligned.
    ///
    /// # Examples
    ///
    /// ```rust
    /// # use {
    /// #     pigeon::Writer,
    /// # };
    /// let mut buf = [0; 2];
    /// let mut writer = Writer::new(&mut buf[..]);
    ///
    /// assert_eq!(writer.bit_offset(), 0);
    /// writer.write_u8_bits(1, 0xFF).unwrap();
    /// assert_eq!(writer.bit_offset(), 1);
    /// writer.write_u8_bits(4, 0x00).unwrap();
    /// assert_eq!(writer.bit_offset(), 5);
    /// writer.write_u8_bits(5, 0xFF).unwrap();
    /// assert_eq!(writer.bit_offset(), 2);
    /// writer.write_u8_bits(8, 0xAA).unwrap();
    /// assert_eq!(writer.bit_offset(), 2);
    /// ```
    pub fn bit_offset(&self) -> u8 {
        self.bit_align
    }

    /// Get the position of the `Writer` in its buffer.
    ///
    /// This rounds upwards, so, for example, if you're at position 3 bytes and 2 bits, this will
    /// return 4.
    ///
    /// # Examples
    ///
    /// ```rust
    /// # use {
    /// #     pigeon::Writer,
    /// # };
    /// let mut buf = [0; 10];
    /// let mut writer = Writer::new(&mut buf[..]);
    ///
    /// assert_eq!(writer.position(), 0);
    ///
    /// writer.write(0u8).unwrap();
    ///
    /// assert_eq!(writer.position(), 1);
    ///
    /// writer.write(0u16).unwrap();
    ///
    /// assert_eq!(writer.position(), 3);
    ///
    /// writer.write_u8_bits(3, 0x05).unwrap();
    ///
    /// assert_eq!(writer.position(), 4);
    ///
    /// writer.write_bytes(b"coucou").unwrap();
    ///
    /// assert_eq!(writer.position(), 10);
    /// ```
    pub fn position(&self) -> usize {
        let offset = if self.bit_align == 0 { 0 } else { 1 };
        self.idx + offset
    }

    /// Write a byte.
    ///
    /// # Position
    ///
    /// This advances the writer only if the write succeeds.
    ///
    /// No data will be written if the write does not succeed.
    ///
    /// # Returns
    ///
    /// On success, this returns `Ok(())`.
    ///
    /// If there is not enough space to write the whole buffer, this returns
    /// `Err(Error::UnexpectedEof)`.
    ///
    /// # Examples
    /// ```rust
    /// # use {
    /// #     pigeon::Writer,
    /// # };
    /// let mut buf = [0; 4];
    /// let mut writer = Writer::new(&mut buf[..]);
    ///
    /// writer.write_u8(0b1011_0101).unwrap();
    /// writer.write_u8_bits(1, 0b0000_0001).unwrap();
    /// writer.write_u8(0b1100_1010).unwrap();
    /// writer.write_u8_bits(5, 0b0001_0111).unwrap();
    /// writer.write_u8(0b0110_1001).unwrap();
    ///
    /// assert!(writer.write_u8(0x00).is_err());
    ///
    /// writer.finish();
    ///
    /// assert_eq!(&buf[..], &[0b1011_0101, 0b1110_0101, 0b0101_1101, 0b1010_0100]);
    /// ```
    pub fn write_u8(&mut self, byte: u8) -> Result<()> {
        let bytes_count = if self.bit_align == 0 { 1 } else { 2 };
        if !self.fits_bytes_count(bytes_count) {
            return Err(Error::UnexpectedEof);
        }
        if self.bit_align == 0 {
            self.buf.write_bytes_at(self.idx, &[byte]);
            self.idx += 1;
            Ok(())
        } else {
            let bi = self.bit_align as u32;
            let byte_fst = byte >> bi;
            let byte_snd = byte << (8 - bi);
            let bit_buffer_pop = (self.bit_buffer >> 8) as u8;
            let output_byte = bit_buffer_pop | byte_fst;
            self.buf.write_bytes_at(self.idx, &[output_byte]);
            self.idx += 1;
            self.bit_buffer = (byte_snd as u16) << 8;
            Ok(())
        }
    }

    /// Write the `bits_count` least significant bits of a byte.
    ///
    /// # Position
    ///
    /// This advances the writer only if the write succeeds.
    ///
    /// No data will be written if the write does not succeed.
    ///
    /// # Returns
    ///
    /// On success, this returns `Ok(())`.
    ///
    /// If there is not enough space to write the whole buffer, this returns
    /// `Err(Error::UnexpectedEof)`.
    ///
    /// # Panics
    ///
    /// This panics if `bits_count` is less than 1 or greater than 8.
    ///
    /// # Examples
    ///
    /// ```rust
    /// # use {
    /// #     pigeon::Writer,
    /// # };
    /// let mut buf = [0; 3];
    /// let mut writer = Writer::new(&mut buf[..]);
    ///
    /// writer.write_u8_bits(4, 0b0000_1010).unwrap();
    /// writer.write_u8_bits(1, 0b0000_0001).unwrap();
    /// writer.write_u8_bits(7, 0b0011_1100).unwrap();
    /// writer.write_u8_bits(4, 0b0000_1100).unwrap();
    /// writer.write_u8_bits(3, 0b0000_0101).unwrap();
    /// writer.write_u8_bits(5, 0b0001_1011).unwrap();
    ///
    /// assert!(writer.write_u8_bits(1, 0x00).is_err());
    ///
    /// writer.finish();
    ///
    /// assert_eq!(&buf[..], &[0b1010_1011, 0b1100_1100, 0b1011_1011]);
    /// ```
    pub fn write_u8_bits(&mut self, bits_count: u8, byte: u8) -> Result<()> {
        assert!(bits_count <= 8);
        assert!(bits_count > 0);
        if !self.fits_bytes_count(1) {
            return Err(Error::UnexpectedEof);
        }
        let mask = 0xFF >> (8 - bits_count);
        self.bit_buffer = self.bit_buffer
            | u16::wrapping_shl(
                (byte & mask) as u16,
                (16 - bits_count - self.bit_align) as u32,
            );
        self.bit_align += bits_count;
        if self.bit_align >= 8 {
            let pop_byte = (self.bit_buffer >> 8) as u8;
            self.bit_buffer = self.bit_buffer << 8;
            self.bit_align -= 8;
            self.buf.write_bytes_at(self.idx, &[pop_byte]);
            self.idx += 1;
        }
        Ok(())
    }

    /// Pad until aligned.
    ///
    /// # Position
    ///
    /// This advances the position by either 0 or 1, depending on whether the
    /// writer was aligned.
    ///
    /// # Examples
    ///
    /// ```rust
    /// # use {
    /// #     pigeon::Writer,
    /// # };
    /// let mut buf = [0; 3];
    /// let mut writer = Writer::new(&mut buf[..]);
    ///
    /// writer.pad_align();
    /// writer.write_u8_bits(2, 0xFF).unwrap();
    /// writer.pad_align();
    /// writer.write_u8_bits(4, 0xFF).unwrap();
    /// writer.pad_align();
    /// writer.write_u8_bits(8, 0xFF).unwrap();
    /// writer.pad_align();
    ///
    /// writer.finish();
    ///
    /// assert_eq!(&buf[..], &[0xC0, 0xF0, 0xFF]);
    /// ```
    pub fn pad_align(&mut self) {
        if self.bit_align != 0 {
            let byte = (self.bit_buffer >> 8) as u8;
            self.buf.write_bytes_at(self.idx, &[byte]);
            self.idx += 1;
            self.bit_align = 0;
            self.bit_buffer = 0;
        }
    }

    /// Write a byte slice into the `Writer` at the current offset.
    ///
    /// # Position
    ///
    /// This advances the writer only if the write succeeds.
    ///
    /// No data will be written if the write does not succeed.
    ///
    /// # Returns
    ///
    /// On success, this returns `Ok(())`.
    /// If there is not enough space to write the whole buffer, this returns
    /// `Err(Error::UnexpectedEof)`.
    ///
    /// # Examples
    ///
    /// ```rust
    /// # use {
    /// #     pigeon::Writer,
    /// # };
    /// let mut buf = [0; 13];
    /// let mut writer = Writer::new(&mut buf[..]);
    ///
    /// writer.write_bytes(b"Hello, world!").unwrap();
    ///
    /// assert!(writer.write_bytes(b"Hewwo").is_err());
    ///
    /// writer.finish();
    ///
    /// assert_eq!(&buf, b"Hello, world!");
    /// ```
    pub fn write_bytes(&mut self, bytes: &[u8]) -> Result<()> {
        if self.bit_align == 0 {
            if self.fits_bytes(bytes) {
                self.buf.write_bytes_at(self.idx, bytes);
                self.idx += bytes.len();
                Ok(())
            } else {
                Err(Error::UnexpectedEof)
            }
        } else {
            if self.fits_bytes_count(bytes.len() + 1) {
                for &byte in bytes {
                    self.write_u8(byte).unwrap(); // already checked
                }
                Ok(())
            } else {
                Err(Error::UnexpectedEof)
            }
        }
    }

    /// Write a type implementing `Pack` into the `Writer` at the current offset.
    ///
    /// # Position
    ///
    /// This advances the writer only if all the underlying writes succeed.
    ///
    /// On failure, there may be incomplete data written to the buffer, the writer's position
    /// will not include that data.
    ///
    /// # Returns
    ///
    /// On success, this returns an `Ok(())`.
    /// If there is not enough space to write all the data, this returns `Err(Error::UnexpectedEof)`.
    ///
    /// # Examples
    ///
    /// ```rust
    /// # use {
    /// #   pigeon::Writer,
    /// # };
    /// let mut buf = [0; 6];
    /// let mut writer = Writer::new(&mut buf[..]);
    ///
    /// writer.write(4242u16).unwrap();
    /// writer.write(56789u32).unwrap();
    ///
    /// writer.finish();
    ///
    /// assert_eq!(&buf[..], &[16, 146, 0, 0, 221, 213]);
    /// ```
    pub fn write<D: Pack>(&mut self, data: D) -> Result<()> {
        let last_idx = self.idx;
        let last_bit_align = self.bit_align;
        let last_bit_buffer = self.bit_buffer;
        match data.pack(self) {
            Ok(()) => Ok(()),
            Err(err) => {
                self.idx = last_idx;
                self.bit_align = last_bit_align;
                self.bit_buffer = last_bit_buffer;
                Err(err)
            }
        }
    }

    /// Check whether a type implementing `Pack` fits into the rest of the buffer.
    ///
    /// # Examples
    ///
    /// ```rust
    /// # use {
    /// #     pigeon::Writer,
    /// # };
    /// let mut buf = [0; 4];
    /// let mut writer = Writer::new(&mut buf[..]);
    ///
    /// assert!(writer.fits(0u32));
    /// assert!(!writer.fits(0u64));
    ///
    /// writer.write(0u8).unwrap();
    ///
    /// assert!(writer.fits(0u16));
    /// assert!(!writer.fits(0u32));
    ///
    /// writer.write(0u16).unwrap();
    ///
    /// assert!(writer.fits(0u8));
    /// assert!(!writer.fits(0u16));
    ///
    /// writer.write(0u8).unwrap();
    ///
    /// assert!(!writer.fits(0u8));
    /// ```
    pub fn fits<D: Pack>(&mut self, data: D) -> bool {
        let size = data.size();
        self.idx + size <= self.buf.size()
    }

    /// Check whether a byteslice can fit into the buffer.
    ///
    /// # Examples
    ///
    /// ```rust
    /// # use {
    /// #     pigeon::Writer,
    /// # };
    /// let mut buf = [0; 4];
    /// let mut writer = Writer::new(&mut buf[..]);
    ///
    /// assert!(writer.fits_bytes(b"hewo"));
    /// assert!(!writer.fits_bytes(b"hallo, iedereen!"));
    ///
    /// writer.write_bytes(b"mew").unwrap();
    ///
    /// assert!(writer.fits_bytes(&[12]));
    /// assert!(!writer.fits_bytes(b"hewo"));
    /// ```
    pub fn fits_bytes(&mut self, bytes: &[u8]) -> bool {
        self.idx + bytes.len() <= self.buf.size()
    }

    /// Check whether an amount of bytes can fit into the buffer.
    ///
    /// # Examples
    ///
    /// ```rust
    /// # use {
    /// #     pigeon::Writer,
    /// # };
    /// let mut buf = [0; 4];
    /// let mut writer = Writer::new(&mut buf[..]);
    ///
    /// assert!(writer.fits_bytes_count(4));
    /// assert!(!writer.fits_bytes_count(5));
    ///
    /// writer.write_bytes(b"mew").unwrap();
    ///
    /// assert!(writer.fits_bytes_count(1));
    /// assert!(!writer.fits_bytes_count(4));
    /// ```
    pub fn fits_bytes_count(&mut self, count: usize) -> bool {
        self.idx + count <= self.buf.size()
    }

    /// Create a `Writer` on a buffer and pass it into a closure.
    ///
    /// This will implicitly finalize the inner writer.
    ///
    /// # Position
    ///
    /// This advances the writer only if all the underlying writes succeed.
    ///
    /// On failure, there may be incomplete data written to the buffer, the writer's position
    /// will not include that data.
    ///
    /// # Returns
    ///
    /// On success, this returns an `Ok` with the amount of bytes written.
    /// If there is not enough space to write all the data, this returns `Err(Error::UnexpectedEof)`.
    ///
    /// # Examples
    ///
    /// ```rust
    /// # use {
    /// #     pigeon::{
    /// #         Writer,
    /// #         U4,
    /// #     },
    /// # };
    /// let mut buf = [0; 32];
    ///
    /// let size = Writer::with(&mut buf[..], |writer| {
    ///     writer.write(U4(0x1))?;
    ///     writer.write(0x23u8)?;
    ///     writer.write(U4(0x4))?;
    ///     Ok(())
    /// }).unwrap();
    ///
    /// assert_eq!(&buf[0..size], &[0x12, 0x34]);
    /// ```
    pub fn with(buf: T, cb: impl FnOnce(&mut Writer<T>) -> Result<()>) -> Result<usize> {
        let mut writer = Writer::new(buf);
        cb(&mut writer)?;
        writer.finish()
    }

    /// Finalize writing. This will flush the bit buffer if the writer is not aligned.
    ///
    /// # Returns
    ///
    /// This returns the amount of bytes that were written to.
    ///
    /// Currently, this cannot fail, but it might have failure modes in the future.
    ///
    /// # Examples
    ///
    /// ```rust
    /// # use {
    /// #     pigeon::{
    /// #         Writer,
    /// #     },
    /// # };
    /// let mut buf = [0; 32];
    ///
    /// let mut writer = Writer::new(&mut buf[..]);
    ///
    /// writer.write_bytes(&[0xAB, 0xCD, 0xEF]).unwrap();
    /// writer.write_u8_bits(5, 0x0E).unwrap();
    ///
    /// let len = writer.finish().unwrap();
    ///
    /// assert_eq!(len, 4);
    /// ```
    pub fn finish(mut self) -> Result<usize> {
        self.pad_align();
        Ok(self.idx)
    }
}

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

    use super::*;

    proptest! {
        #[test]
        fn prop_write_u8_bits_and_write_u8_equivalent_for_8_bits(
            pre_byte in prop::num::u8::ANY, pre_bits_count in 0..8u8,
            byte in prop::num::u8::ANY,
            post_byte in prop::num::u8::ANY, post_bits_count in 0..8u8,
        ) {
            let mut buf_a = [0; 3];
            let mut buf_b = [0; 3];

            let len_a = {
                let mut writer_a = Writer::new(&mut buf_a[..]);
                if pre_bits_count > 0 {
                    writer_a.write_u8_bits(pre_bits_count, pre_byte).unwrap();
                }
                writer_a.write_u8_bits(8, byte).unwrap();
                if post_bits_count > 0 {
                    writer_a.write_u8_bits(post_bits_count, post_byte).unwrap();
                }
                writer_a.finish().unwrap()
            };

            let len_b = {
                let mut writer_b = Writer::new(&mut buf_b[..]);
                if pre_bits_count > 0 {
                    writer_b.write_u8_bits(pre_bits_count, pre_byte).unwrap();
                }
                writer_b.write_u8(byte).unwrap();
                if post_bits_count > 0 {
                    writer_b.write_u8_bits(post_bits_count, post_byte).unwrap();
                }
                writer_b.finish().unwrap()
            };

            assert_eq!(len_a, len_b);
            assert_eq!(&buf_a[..len_a], &buf_b[..len_b]);
        }
    }

    #[test]
    fn test_writing_1_a() {
        let mut buf = [0; 2];
        let mut writer = Writer::new(&mut buf[..]);
        writer.write(true).unwrap();
        assert_eq!(writer.bit_align, 1);
        assert_eq!(writer.bit_buffer, 0x8000);
        writer.write(0u8).unwrap();
        assert_eq!(writer.bit_align, 1);
        assert_eq!(writer.bit_buffer, 0x0000);
        writer.write(true).unwrap();
        assert_eq!(writer.bit_align, 2);
        assert_eq!(writer.bit_buffer, 0x4000);
        let len = writer.finish().unwrap();
        assert_eq!(len, 2);
        assert_eq!(&buf[..len], [0x80, 0x40]);
    }

    #[test]
    fn test_writing_1_b() {
        let mut buf = [0; 2];
        let mut writer = Writer::new(&mut buf[..]);
        writer.write(false).unwrap();
        assert_eq!(writer.bit_align, 1);
        assert_eq!(writer.bit_buffer, 0x0000);
        writer.write(0u8).unwrap();
        assert_eq!(writer.bit_align, 1);
        assert_eq!(writer.bit_buffer, 0x0000);
        writer.write(true).unwrap();
        assert_eq!(writer.bit_align, 2);
        assert_eq!(writer.bit_buffer, 0x4000);
        let len = writer.finish().unwrap();
        assert_eq!(len, 2);
        assert_eq!(&buf[..len], [0x00, 0x40]);
    }
}