dambi 0.1.2

use std::ops::{Bound, Deref, DerefMut, RangeBounds};

use super::byte_buffer::ByteBuffer;

#[derive(Clone)]
pub struct Bytes {
    repr: BytesRepr,
}

#[derive(Clone)]
enum BytesRepr {
    Static(&'static [u8]),
    Owned {
        buf: ByteBuffer,
        start: u32,
        len: u32,
    },
}

impl Bytes {
    #[must_use]
    pub const fn new() -> Self {
        Self {
            repr: BytesRepr::Static(&[]),
        }
    }

    #[must_use]
    pub const fn from_static(s: &'static [u8]) -> Self {
        Self {
            repr: BytesRepr::Static(s),
        }
    }

    #[must_use]
    pub fn from_vec(v: Vec<u8>) -> Self {
        Self::copy_from_slice(&v)
    }

    pub fn from_byte_buffer_range(
        buf: super::byte_buffer::ByteBuffer,
        start: u32,
        len: u32,
    ) -> Self {
        debug_assert!(
            start
                .checked_add(len)
                .is_some_and(|end| end <= buf.capacity())
        );
        Self {
            repr: BytesRepr::Owned { buf, start, len },
        }
    }

    #[must_use]
    pub fn copy_from_slice(s: &[u8]) -> Self {
        if s.is_empty() {
            return Self::new();
        }
        let len = s.len();
        assert!(
            len <= u32::MAX as usize,
            "Bytes: payload too large ({len}, max {})",
            u32::MAX
        );
        Self {
            repr: BytesRepr::Owned {
                buf: ByteBuffer::from_slice(s),
                start: 0,
                len: len as u32,
            },
        }
    }

    pub fn len(&self) -> usize {
        match &self.repr {
            BytesRepr::Static(s) => s.len(),
            BytesRepr::Owned { len, .. } => *len as usize,
        }
    }

    pub fn is_empty(&self) -> bool {
        self.len() == 0
    }

    pub fn as_slice(&self) -> &[u8] {
        match &self.repr {
            BytesRepr::Static(s) => s,
            // SAFETY: every constructor preserves `start + len <= buf.capacity()`,
            // so the resulting slice stays inside the payload region.
            BytesRepr::Owned { buf, start, len } => unsafe {
                std::slice::from_raw_parts(buf.data_ptr().add(*start as usize), *len as usize)
            },
        }
    }

    #[must_use]
    pub fn slice(&self, range: impl RangeBounds<usize>) -> Self {
        let len = self.len();
        let start = match range.start_bound() {
            Bound::Included(&n) => n,
            Bound::Excluded(&n) => n + 1,
            Bound::Unbounded => 0,
        };
        let end = match range.end_bound() {
            Bound::Included(&n) => n + 1,
            Bound::Excluded(&n) => n,
            Bound::Unbounded => len,
        };
        assert!(
            start <= end && end <= len,
            "Bytes::slice: range out of bounds"
        );
        if start == end {
            return Self::new();
        }
        match &self.repr {
            BytesRepr::Static(s) => Self::from_static(&s[start..end]),
            BytesRepr::Owned {
                buf, start: cur, ..
            } => Self {
                repr: BytesRepr::Owned {
                    buf: buf.clone(),
                    start: *cur + start as u32,
                    len: (end - start) as u32,
                },
            },
        }
    }

    pub fn advance(&mut self, n: usize) {
        let len = self.len();
        assert!(n <= len, "Bytes::advance: out of bounds");
        match &mut self.repr {
            BytesRepr::Static(s) => *s = &s[n..],
            BytesRepr::Owned { start, len, .. } => {
                *start += n as u32;
                *len -= n as u32;
            }
        }
    }

    #[must_use]
    pub fn split_to(&mut self, at: usize) -> Self {
        let head = self.slice(..at);
        self.advance(at);
        head
    }

    pub fn clear(&mut self) {
        self.repr = BytesRepr::Static(&[]);
    }

    #[must_use]
    pub fn freeze(self) -> Self {
        self
    }

    pub fn truncate(&mut self, n: usize) {
        let len = self.len();
        if n >= len {
            return;
        }
        match &mut self.repr {
            BytesRepr::Static(s) => *s = &s[..n],
            BytesRepr::Owned { len, .. } => *len = n as u32,
        }
    }
}

impl Default for Bytes {
    fn default() -> Self {
        Self::new()
    }
}

impl AsRef<[u8]> for Bytes {
    fn as_ref(&self) -> &[u8] {
        self.as_slice()
    }
}

impl Deref for Bytes {
    type Target = [u8];
    fn deref(&self) -> &[u8] {
        self.as_slice()
    }
}

impl From<&'static [u8]> for Bytes {
    fn from(value: &'static [u8]) -> Self {
        Self::from_static(value)
    }
}

impl<const N: usize> From<&'static [u8; N]> for Bytes {
    fn from(value: &'static [u8; N]) -> Self {
        Self::from_static(value)
    }
}

impl From<Vec<u8>> for Bytes {
    fn from(value: Vec<u8>) -> Self {
        Self::from_vec(value)
    }
}

impl From<BytesMut> for Bytes {
    fn from(value: BytesMut) -> Self {
        Self::from_vec(value.into_vec())
    }
}

impl From<String> for Bytes {
    fn from(value: String) -> Self {
        Self::from_vec(value.into_bytes())
    }
}

impl From<&str> for Bytes {
    fn from(value: &str) -> Self {
        Self::from_vec(value.as_bytes().to_vec())
    }
}

impl PartialEq for Bytes {
    fn eq(&self, other: &Self) -> bool {
        self.as_slice() == other.as_slice()
    }
}

impl PartialEq<[u8]> for Bytes {
    fn eq(&self, other: &[u8]) -> bool {
        self.as_slice() == other
    }
}

impl PartialEq<&[u8]> for Bytes {
    fn eq(&self, other: &&[u8]) -> bool {
        self.as_slice() == *other
    }
}

impl Eq for Bytes {}

impl std::hash::Hash for Bytes {
    fn hash<H: std::hash::Hasher>(&self, state: &mut H) {
        self.as_slice().hash(state);
    }
}

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

#[derive(Clone, Default, PartialEq, Eq, Hash)]
pub struct BytesMut {
    inner: Vec<u8>,
}

impl BytesMut {
    #[must_use]
    pub const fn new() -> Self {
        Self { inner: Vec::new() }
    }

    #[must_use]
    pub fn with_capacity(cap: usize) -> Self {
        Self {
            inner: Vec::with_capacity(cap),
        }
    }

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

    pub fn is_empty(&self) -> bool {
        self.inner.is_empty()
    }

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

    pub fn as_slice(&self) -> &[u8] {
        &self.inner
    }

    pub fn as_mut_slice(&mut self) -> &mut [u8] {
        &mut self.inner
    }

    pub fn extend_from_slice(&mut self, src: &[u8]) {
        self.inner.extend_from_slice(src);
    }

    pub fn reserve(&mut self, additional: usize) {
        self.inner.reserve(additional);
    }

    pub fn clear(&mut self) {
        self.inner.clear();
    }

    pub fn truncate(&mut self, len: usize) {
        self.inner.truncate(len);
    }

    pub fn push(&mut self, byte: u8) {
        self.inner.push(byte);
    }

    /// # Safety
    ///
    /// Caller asserts `len` bytes within `capacity()` are initialized.
    pub unsafe fn set_len(&mut self, len: usize) {
        unsafe { self.inner.set_len(len) };
    }

    pub fn spare_capacity_mut(&mut self) -> &mut [std::mem::MaybeUninit<u8>] {
        self.inner.spare_capacity_mut()
    }

    #[must_use]
    pub fn split(&mut self) -> Bytes {
        Bytes::from_vec(std::mem::take(&mut self.inner))
    }

    #[must_use]
    pub fn split_to(&mut self, at: usize) -> BytesMut {
        let mut full = std::mem::take(&mut self.inner);
        let tail = full.split_off(at);
        self.inner = tail;
        BytesMut { inner: full }
    }

    #[must_use]
    pub fn split_off(&mut self, at: usize) -> BytesMut {
        BytesMut {
            inner: self.inner.split_off(at),
        }
    }

    #[must_use]
    pub fn freeze(self) -> Bytes {
        Bytes::from_vec(self.inner)
    }

    pub fn into_vec(self) -> Vec<u8> {
        self.inner
    }
}

impl AsRef<[u8]> for BytesMut {
    fn as_ref(&self) -> &[u8] {
        &self.inner
    }
}

impl AsMut<[u8]> for BytesMut {
    fn as_mut(&mut self) -> &mut [u8] {
        &mut self.inner
    }
}

impl Deref for BytesMut {
    type Target = [u8];
    fn deref(&self) -> &[u8] {
        &self.inner
    }
}

impl DerefMut for BytesMut {
    fn deref_mut(&mut self) -> &mut [u8] {
        &mut self.inner
    }
}

impl From<Vec<u8>> for BytesMut {
    fn from(value: Vec<u8>) -> Self {
        Self { inner: value }
    }
}

impl From<&[u8]> for BytesMut {
    fn from(value: &[u8]) -> Self {
        Self {
            inner: value.to_vec(),
        }
    }
}

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