use core::{
    fmt, hint,
    marker::PhantomData,
    ops::{Deref, DerefMut},
    ptr, slice,
};

use bytes::{buf::UninitSlice, BufMut};

/// An optimized encoder for pre-allocated bytes.
pub struct EncoderSlice<'a> {
    start: ptr::NonNull<u8>,
    ptr: ptr::NonNull<u8>,
    end: ptr::NonNull<u8>,
    _marker: PhantomData<&'a mut [u8]>,
}

unsafe impl<'a> Send for EncoderSlice<'a> {}
unsafe impl<'a> Sync for EncoderSlice<'a> {}

impl<'a> fmt::Debug for EncoderSlice<'a> {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        fmt::Debug::fmt(&**self, f)
    }
}

impl<'a> PartialEq for EncoderSlice<'a> {
    #[inline(always)]
    fn eq(&self, other: &Self) -> bool {
        **self == **other
    }
}

impl<'a> Eq for EncoderSlice<'a> {}

impl<'a> EncoderSlice<'a> {
    /// Creates a new instance of the encoder with the provided buffer.
    #[inline(always)]
    pub const fn new(buffer: &'a mut [u8]) -> Self {
        let len = buffer.len();
        // Safety: buffer.as_mut_ptr() is guaranteed non-null by &mut [u8]
        let start = unsafe { ptr::NonNull::new_unchecked(buffer.as_mut_ptr()) };
        let end = unsafe { ptr::NonNull::new_unchecked(start.as_ptr().add(len)) };

        Self {
            start,
            ptr: start,
            end,
            _marker: PhantomData,
        }
    }

    #[inline(always)]
    fn written_len(&self) -> usize {
        self.ptr.as_ptr() as usize - self.start.as_ptr() as usize
    }

    /// Returns a shared slice of the data written so far.
    #[inline(always)]
    pub fn written(&self) -> &[u8] {
        unsafe { slice::from_raw_parts(self.start.as_ptr(), self.written_len()) }
    }

    /// Returns a mutable slice of the data written so far.
    #[inline(always)]
    pub fn written_mut(&mut self) -> &mut [u8] {
        unsafe { slice::from_raw_parts_mut(self.start.as_ptr(), self.written_len()) }
    }

    /// Consumes the encoder and returns the written part of the buffer
    /// with the original lifetime `'a`.
    #[inline(always)]
    pub fn into_written(self) -> &'a mut [u8] {
        unsafe { slice::from_raw_parts_mut(self.start.as_ptr(), self.written_len()) }
    }

    /// Resets the encoder, allowing the underlying buffer to be reused from the start.
    /// This does not zero the memory.
    #[inline(always)]
    pub const fn reset(&mut self) {
        self.ptr = self.start;
    }

    #[inline(always)]
    fn remaining(&self) -> usize {
        self.end.as_ptr() as usize - self.ptr.as_ptr() as usize
    }

    /// Writes a fixed-size array by value.
    ///
    /// Prefer this over `put_slice` when the size is known at compile time: the
    /// const `N` lets LLVM emit scalar `MOV` instructions instead of a memcpy
    /// call, and avoids the fat-pointer (ptr + len) that `put_slice` receives.
    #[inline(always)]
    pub fn put_arr<const N: usize>(&mut self, src: [u8; N]) {
        debug_assert!(self.remaining() >= N);
        #[cfg(feature = "assertions")]
        assert!(self.remaining() >= N);

        unsafe {
            hint::assert_unchecked(self.remaining() >= N);
            ptr::copy_nonoverlapping(src.as_ptr(), self.ptr.as_ptr(), N);
            self.ptr = self.ptr.add(N);
        }
    }

    /// Writes a format tag byte followed by a single value byte.
    ///
    /// Two direct `write` calls to the destination — no intermediate stack buffer.
    #[inline(always)]
    pub fn put_tag_u8(&mut self, tag: u8, val: u8) {
        debug_assert!(self.remaining() >= 2);
        #[cfg(feature = "assertions")]
        assert!(self.remaining() >= 2);

        unsafe {
            hint::assert_unchecked(self.remaining() >= 2);
            self.ptr.as_ptr().write(tag);
            self.ptr.as_ptr().add(1).write(val);
            self.ptr = self.ptr.add(2);
        }
    }

    /// Writes a format tag byte followed by a big-endian `u16` (3 bytes total).
    ///
    /// Uses `write_unaligned` for the 16-bit field: on x86-64 this compiles to
    /// `MOV byte, BSWAP + MOV word` — two instructions with no stack array.
    #[inline(always)]
    pub fn put_tag_u16_be(&mut self, tag: u8, val: u16) {
        debug_assert!(self.remaining() >= 3);
        #[cfg(feature = "assertions")]
        assert!(self.remaining() >= 3);

        unsafe {
            hint::assert_unchecked(self.remaining() >= 3);
            self.ptr.as_ptr().write(tag);
            (self.ptr.as_ptr().add(1) as *mut u16).write_unaligned(val.to_be());
            self.ptr = self.ptr.add(3);
        }
    }

    /// Writes a format tag byte followed by a big-endian `u32` (5 bytes total).
    #[inline(always)]
    pub fn put_tag_u32_be(&mut self, tag: u8, val: u32) {
        debug_assert!(self.remaining() >= 5);
        #[cfg(feature = "assertions")]
        assert!(self.remaining() >= 5);

        unsafe {
            hint::assert_unchecked(self.remaining() >= 5);
            self.ptr.as_ptr().write(tag);
            (self.ptr.as_ptr().add(1) as *mut u32).write_unaligned(val.to_be());
            self.ptr = self.ptr.add(5);
        }
    }

    /// Writes a format tag byte followed by a big-endian `u64` (9 bytes total).
    #[inline(always)]
    pub fn put_tag_u64_be(&mut self, tag: u8, val: u64) {
        debug_assert!(self.remaining() >= 9);
        #[cfg(feature = "assertions")]
        assert!(self.remaining() >= 9);

        unsafe {
            hint::assert_unchecked(self.remaining() >= 9);
            self.ptr.as_ptr().write(tag);
            (self.ptr.as_ptr().add(1) as *mut u64).write_unaligned(val.to_be());
            self.ptr = self.ptr.add(9);
        }
    }
}

impl<'a> Deref for EncoderSlice<'a> {
    type Target = [u8];
    #[inline(always)]
    fn deref(&self) -> &Self::Target {
        unsafe { slice::from_raw_parts(self.start.as_ptr(), self.written_len()) }
    }
}

impl<'a> DerefMut for EncoderSlice<'a> {
    #[inline(always)]
    fn deref_mut(&mut self) -> &mut Self::Target {
        unsafe { slice::from_raw_parts_mut(self.start.as_ptr(), self.written_len()) }
    }
}

unsafe impl<'a> BufMut for EncoderSlice<'a> {
    #[inline(always)]
    fn remaining_mut(&self) -> usize {
        self.remaining()
    }

    #[inline(always)]
    fn chunk_mut(&mut self) -> &mut UninitSlice {
        unsafe { UninitSlice::from_raw_parts_mut(self.ptr.as_ptr(), self.remaining()) }
    }

    #[inline(always)]
    unsafe fn advance_mut(&mut self, cnt: usize) {
        debug_assert!(self.remaining() >= cnt);
        #[cfg(feature = "assertions")]
        assert!(self.remaining() >= cnt);

        unsafe {
            hint::assert_unchecked(self.remaining() >= cnt);
            self.ptr = self.ptr.add(cnt);
        }
    }

    #[inline(always)]
    fn put_u8(&mut self, val: u8) {
        debug_assert!(self.ptr.as_ptr() < self.end.as_ptr());
        #[cfg(feature = "assertions")]
        assert!(self.ptr.as_ptr() < self.end.as_ptr());

        unsafe {
            hint::assert_unchecked(self.ptr.as_ptr() < self.end.as_ptr());
            self.ptr.as_ptr().write(val);
            self.ptr = self.ptr.add(1);
        }
    }

    #[inline(always)]
    fn put_slice(&mut self, src: &[u8]) {
        let n = src.len();
        debug_assert!(self.remaining() >= n);
        #[cfg(feature = "assertions")]
        assert!(self.remaining() >= n);

        unsafe {
            hint::assert_unchecked(self.remaining() >= n);
            ptr::copy_nonoverlapping(src.as_ptr(), self.ptr.as_ptr(), n);
            self.ptr = self.ptr.add(n);
        }
    }

    #[inline(always)]
    fn put_bytes(&mut self, val: u8, cnt: usize) {
        debug_assert!(self.remaining() >= cnt);
        #[cfg(feature = "assertions")]
        assert!(self.remaining() >= cnt);

        unsafe {
            hint::assert_unchecked(self.remaining() >= cnt);
            ptr::write_bytes(self.ptr.as_ptr(), val, cnt);
            self.ptr = self.ptr.add(cnt);
        }
    }

    #[inline(always)]
    fn put_u16(&mut self, n: u16) {
        debug_assert!(self.remaining() >= 2);
        #[cfg(feature = "assertions")]
        assert!(self.remaining() >= 2);

        unsafe {
            hint::assert_unchecked(self.remaining() >= 2);
            (self.ptr.as_ptr() as *mut u16).write_unaligned(n.to_be());
            self.ptr = self.ptr.add(2);
        }
    }

    #[inline(always)]
    fn put_u32(&mut self, n: u32) {
        debug_assert!(self.remaining() >= 4);
        #[cfg(feature = "assertions")]
        assert!(self.remaining() >= 4);

        unsafe {
            hint::assert_unchecked(self.remaining() >= 4);
            (self.ptr.as_ptr() as *mut u32).write_unaligned(n.to_be());
            self.ptr = self.ptr.add(4);
        }
    }

    #[inline(always)]
    fn put_u64(&mut self, n: u64) {
        debug_assert!(self.remaining() >= 8);
        #[cfg(feature = "assertions")]
        assert!(self.remaining() >= 8);

        unsafe {
            hint::assert_unchecked(self.remaining() >= 8);
            (self.ptr.as_ptr() as *mut u64).write_unaligned(n.to_be());
            self.ptr = self.ptr.add(8);
        }
    }

    #[inline(always)]
    fn put_i8(&mut self, n: i8) {
        self.put_u8(n as u8);
    }

    #[inline(always)]
    fn put_i16(&mut self, n: i16) {
        self.put_u16(n as u16);
    }

    #[inline(always)]
    fn put_i32(&mut self, n: i32) {
        self.put_u32(n as u32);
    }

    #[inline(always)]
    fn put_i64(&mut self, n: i64) {
        self.put_u64(n as u64);
    }
}
