use crate::Channel;
use commonware_codec::{Buf, EncodeSize, Error, RangeCfg, Read, ReadExt as _, Write, varint::UInt};
use commonware_runtime::{BufMut, IoBuf, IoBufMut, IoBufs};
use commonware_stream::Upgrader;

/// Data is an arbitrary message sent between peers.
#[derive(Clone, Debug, PartialEq)]
pub struct Data {
    /// A unique identifier for the channel the message is sent on.
    ///
    /// This is used to route the message to the correct handler.
    pub channel: u64,

    /// The payload of the message.
    pub message: IoBuf,
}

impl EncodeSize for Data {
    fn encode_size(&self) -> usize {
        UInt(self.channel).encode_size() + self.message.encode_size()
    }
}

impl Write for Data {
    fn write(&self, buf: &mut impl BufMut) {
        UInt(self.channel).write(buf);
        self.message.write(buf);
    }
}

impl Read for Data {
    type Cfg = RangeCfg<usize>;

    fn read_cfg(buf: &mut impl Buf, range: &Self::Cfg) -> Result<Self, Error> {
        let channel = UInt::read(buf)?.into();
        let message = IoBuf::read_cfg(buf, range)?;
        Ok(Self { channel, message })
    }
}

/// Prefix byte identifying a data frame on the wire.
pub(crate) const DATA_PREFIX: u8 = 0;

/// The maximum overhead (in bytes) when encoding a `message` into a data frame.
///
/// The byte overhead is calculated as the sum of the following:
/// - 1: Frame discriminant
/// - 10: Channel varint
/// - 5: Message length varint (lengths longer than 32 bits are forbidden by the codec)
pub const MAX_PAYLOAD_OVERHEAD: u32 = 1 + 10 + 5;

/// Maximum application payload size supported by the upgrader's streams.
///
/// # Panics
///
/// Panics if the stream's limit cannot accommodate p2p framing overhead.
pub const fn max_size<U: Upgrader>() -> u32 {
    U::MAX_SIZE
        .checked_sub(MAX_PAYLOAD_OVERHEAD)
        .expect("stream message limit too small for p2p framing")
}

/// Pre-encoded data ready for transmission.
///
/// Contains the channel ID (for metrics) and the pre-encoded payload bytes.
/// The `payload` field contains the fully encoded data frame bytes,
/// stored as one or more buffers ready to be sent directly to the stream layer.
#[derive(Clone, Debug)]
pub struct EncodedData {
    /// The channel this data belongs to (used for metrics/logging).
    pub channel: Channel,

    /// Pre-encoded data frame bytes ready for transmission.
    pub payload: IoBufs,
}

impl EncodedData {
    /// Encode data frame bytes in-place as:
    /// `DATA_PREFIX || channel || message_len || message`.
    pub fn new(channel: Channel, mut message: IoBufs) -> Self {
        let payload_len = message.len();
        let header_len =
            DATA_PREFIX.encode_size() + UInt(channel).encode_size() + payload_len.encode_size();
        let mut header = IoBufMut::with_capacity(header_len);
        DATA_PREFIX.write(&mut header);
        UInt(channel).write(&mut header);
        payload_len.write(&mut header);
        assert_eq!(header.len(), header_len, "data header size mismatch");
        message.prepend(header.freeze());

        Self {
            channel,
            payload: message,
        }
    }
}

#[cfg(feature = "arbitrary")]
impl arbitrary::Arbitrary<'_> for Data {
    fn arbitrary(u: &mut arbitrary::Unstructured<'_>) -> arbitrary::Result<Self> {
        let channel = u.arbitrary::<u64>()?;
        let message = {
            let size = u.int_in_range(0..=1024)?;
            let bytes = u.bytes(size)?;
            IoBuf::copy_from_slice(bytes)
        };
        Ok(Self { channel, message })
    }
}

/// Counts global allocations on demand, so tests can show that a decode path allocates nothing.
#[cfg(test)]
pub(crate) mod allocations {
    use commonware_runtime::{BufMut as _, IoBufMut, IoBufs};
    use std::{
        alloc::{GlobalAlloc, Layout, System},
        cell::Cell,
    };

    /// Counts global allocations while a test asks it to, delegating every call to the
    /// system allocator.
    struct TrackingAllocator;

    std::thread_local! {
        static ALLOCATION_CALLS: Cell<Option<usize>> = const { Cell::new(None) };
    }

    #[global_allocator]
    static ALLOCATOR: TrackingAllocator = TrackingAllocator;

    /// Counts an allocation when tracking is enabled on this thread.
    fn record_allocation() {
        ALLOCATION_CALLS.with(|calls| {
            if let Some(count) = calls.get() {
                calls.set(Some(count + 1));
            }
        });
    }

    // SAFETY: Every operation delegates to the system allocator with unchanged
    // arguments. Recording an allocation only updates a thread-local cell.
    unsafe impl GlobalAlloc for TrackingAllocator {
        unsafe fn alloc(&self, layout: Layout) -> *mut u8 {
            record_allocation();
            // SAFETY: The caller supplies the layout required by GlobalAlloc.
            unsafe { System.alloc(layout) }
        }

        unsafe fn alloc_zeroed(&self, layout: Layout) -> *mut u8 {
            record_allocation();
            // SAFETY: The caller supplies the layout required by GlobalAlloc.
            unsafe { System.alloc_zeroed(layout) }
        }

        unsafe fn dealloc(&self, ptr: *mut u8, layout: Layout) {
            // SAFETY: The pointer was allocated through this wrapper, which
            // delegates every allocation to System with the same layout.
            unsafe { System.dealloc(ptr, layout) }
        }

        unsafe fn realloc(&self, ptr: *mut u8, layout: Layout, new_size: usize) -> *mut u8 {
            record_allocation();
            // SAFETY: The pointer and layout came from System through this wrapper,
            // and the caller supplies the new size required by GlobalAlloc.
            unsafe { System.realloc(ptr, layout, new_size) }
        }
    }

    /// Runs `f` and returns its result with the number of allocations it made.
    pub(crate) fn allocation_calls_during<T>(f: impl FnOnce() -> T) -> (T, usize) {
        ALLOCATION_CALLS.with(|calls| assert!(calls.replace(Some(0)).is_none()));
        let result = f();
        let calls = ALLOCATION_CALLS.with(|calls| calls.replace(None).expect("tracking enabled"));
        (result, calls)
    }

    /// Returns `encoded` in one native buffer, as a received record arrives, with the address
    /// range of its bytes.
    pub(crate) fn frame(encoded: &[u8]) -> (IoBufs, std::ops::Range<*const u8>) {
        let mut buf = IoBufMut::with_capacity(encoded.len());
        buf.put_slice(encoded);
        let range = buf.as_ref().as_ptr_range();
        (IoBufs::from(buf.freeze()), range)
    }
}

#[cfg(test)]
mod tests {
    use super::{
        allocations::{allocation_calls_during, frame},
        *,
    };
    use commonware_codec::{Decode as _, Encode as _, Error};
    use commonware_cryptography::{ChaCha20Poly1305, ed25519};
    use commonware_stream::SakeCups;

    #[test]
    fn test_max_size_bounds() {
        assert_eq!(
            max_size::<SakeCups<ed25519::PrivateKey, ChaCha20Poly1305>>() + MAX_PAYLOAD_OVERHEAD,
            SakeCups::<ed25519::PrivateKey, ChaCha20Poly1305>::MAX_SIZE
        );
    }

    #[test]
    fn test_data_codec() {
        let original = Data {
            channel: 12345,
            message: IoBuf::from(b"Hello, world!"),
        };
        let encoded = original.encode();
        let decoded = Data::decode_cfg(encoded, &(13..=13).into()).unwrap();
        assert_eq!(original, decoded);

        let too_short = Data::decode_cfg(original.encode(), &(0..13).into());
        assert!(matches!(too_short, Err(Error::InvalidLength(13))));

        let too_long = Data::decode_cfg(original.encode(), &(14..).into());
        assert!(matches!(too_long, Err(Error::InvalidLength(13))));
    }

    #[test]
    fn test_decode_invalid() {
        let invalid_payload = [3, 0, 1, 2, 3, 4, 5, 6, 7, 8, 9];
        let result = Data::decode_cfg(commonware_codec::Copying(&invalid_payload), &(..).into());
        assert!(result.is_err());
    }

    /// Decoding a single-chunk frame neither copies the message nor allocates owner records
    /// for it: the reader hands the body out as a view of its own buffer.
    #[commonware_macros::test_group("miri")]
    #[test]
    fn test_decode_does_not_allocate() {
        let original = Data {
            channel: 12345,
            message: IoBuf::from(b"Hello, world!"),
        };
        let (frame, bytes) = frame(&original.encode());
        let range = (..).into();
        let (decoded, calls) = allocation_calls_during(|| Data::decode_cfg(frame, &range).unwrap());
        assert_eq!(calls, 0);
        assert_eq!(decoded, original);
        assert!(bytes.contains(&decoded.message.as_ptr()));
    }

    #[test]
    fn test_encoded_data_new_matches_data_encode() {
        let mut message = IoBufs::from(IoBuf::from(b"hello "));
        message.append(IoBuf::from(b"world"));
        message.append(IoBuf::from(b"!"));

        let data = Data {
            channel: 12345,
            message: message.clone().coalesce(),
        };

        let mut expected = IoBufs::from(data.encode());
        expected.prepend(IoBuf::from(vec![DATA_PREFIX]));

        let encoded = EncodedData::new(12345, message);
        assert_eq!(encoded.channel, 12345);
        assert_eq!(encoded.payload.coalesce(), expected.coalesce());
    }

    #[cfg(feature = "arbitrary")]
    mod conformance {
        use super::*;
        use commonware_codec::conformance::CodecConformance;

        commonware_conformance::conformance_tests! {
            CodecConformance<Data>,
        }
    }
}