use std::marker::PhantomData;

use bytes::{Buf, BytesMut};

use crate::{StreamBodyError, error::StreamBodyKind};

#[derive(Clone, Debug)]
pub(crate) struct ProtobufLenPrefixCodec<T> {
    max_length: usize,
    cursor: ProtobufCursor,
    _ph: PhantomData<T>,
}

#[derive(Clone, Debug)]
struct ProtobufCursor {
    current_obj_len: usize,
}

impl<T> ProtobufLenPrefixCodec<T> {
    pub(crate) const fn new_with_max_length(max_length: usize) -> Self {
        let initial_cursor = ProtobufCursor { current_obj_len: 0 };

        Self {
            max_length,
            cursor: initial_cursor,
            _ph: PhantomData,
        }
    }
}

impl<T> tokio_util::codec::Decoder for ProtobufLenPrefixCodec<T>
where
    T: prost::Message + Default,
{
    type Item = T;
    type Error = StreamBodyError;

    fn decode(&mut self, buf: &mut BytesMut) -> Result<Option<T>, StreamBodyError> {
        let buf_len = buf.len();
        if buf_len == 0 {
            return Ok(None);
        }

        if self.cursor.current_obj_len == 0 {
            let bytes = buf.chunk();
            let byte = bytes[0];
            if byte < 0x80 {
                buf.advance(1);
                self.cursor.current_obj_len = u64::from(byte) as usize;
            } else if buf_len > 10 || bytes[buf_len - 1] < 0x80 {
                let (value, advance) = decode_varint_slice(bytes)?;
                buf.advance(advance);
                self.cursor.current_obj_len = value as usize;
            }
            Ok(None)
        } else if self.cursor.current_obj_len > self.max_length {
            Err(StreamBodyError::new(
                StreamBodyKind::MaxLenReachedError,
                None,
                Some("Max object length reached".into()),
            ))
        } else if buf_len >= self.cursor.current_obj_len {
            let obj_bytes = buf.copy_to_bytes(self.cursor.current_obj_len);
            let result: Result<Option<T>, StreamBodyError> =
                prost::Message::decode(obj_bytes).map(Some).map_err(|err| {
                    StreamBodyError::new(StreamBodyKind::CodecError, Some(Box::new(err)), None)
                });
            self.cursor.current_obj_len = 0;
            result
        } else {
            Ok(None)
        }
    }

    fn decode_eof(&mut self, buf: &mut BytesMut) -> Result<Option<T>, StreamBodyError> {
        self.decode(buf)
    }
}

/// Decodes a LEB128-encoded variable length integer from the slice, returning the value and the
/// number of bytes read.
///
/// ## Safety
///
/// The caller must ensure that `bytes` is non-empty and either `bytes.len() >= 10` or the last
/// element in bytes is < `0x80`.
#[inline]
fn decode_varint_slice(bytes: &[u8]) -> Result<(u64, usize), StreamBodyError> {
    assert!(!bytes.is_empty());
    assert!(bytes.len() > 10 || bytes[bytes.len() - 1] < 0x80);

    let mut b: u8 = bytes[0];
    let mut part0: u32 = u32::from(b);
    if b < 0x80 {
        return Ok((u64::from(part0), 1));
    }
    part0 -= 0x80;
    b = bytes[1];
    part0 += u32::from(b) << 7;
    if b < 0x80 {
        return Ok((u64::from(part0), 2));
    }
    part0 -= 0x80 << 7;
    b = bytes[2];
    part0 += u32::from(b) << 14;
    if b < 0x80 {
        return Ok((u64::from(part0), 3));
    }
    part0 -= 0x80 << 14;
    b = bytes[3];
    part0 += u32::from(b) << 21;
    if b < 0x80 {
        return Ok((u64::from(part0), 4));
    }
    part0 -= 0x80 << 21;
    let value = u64::from(part0);

    b = bytes[4];
    let mut part1: u32 = u32::from(b);
    if b < 0x80 {
        return Ok((value + (u64::from(part1) << 28), 5));
    }
    part1 -= 0x80;
    b = bytes[5];
    part1 += u32::from(b) << 7;
    if b < 0x80 {
        return Ok((value + (u64::from(part1) << 28), 6));
    }
    part1 -= 0x80 << 7;
    b = bytes[6];
    part1 += u32::from(b) << 14;
    if b < 0x80 {
        return Ok((value + (u64::from(part1) << 28), 7));
    }
    part1 -= 0x80 << 14;
    b = bytes[7];
    part1 += u32::from(b) << 21;
    if b < 0x80 {
        return Ok((value + (u64::from(part1) << 28), 8));
    }
    part1 -= 0x80 << 21;
    let value = value + ((u64::from(part1)) << 28);

    b = bytes[8];
    let mut part2: u32 = u32::from(b);
    if b < 0x80 {
        return Ok((value + (u64::from(part2) << 56), 9));
    }
    part2 -= 0x80;
    b = bytes[9];
    part2 += u32::from(b) << 7;
    if b < 0x02 {
        return Ok((value + (u64::from(part2) << 56), 10));
    }

    Err(StreamBodyError::new(
        StreamBodyKind::CodecError,
        None,
        Some("invalid varint".into()),
    ))
}

#[cfg(test)]
mod tests {
    use prost::Message as _;
    use tokio_util::codec::Decoder;

    use super::*;

    #[derive(Clone, PartialEq, prost::Message)]
    struct TestMsg {
        #[prost(string, tag = "1")]
        name: String,
        #[prost(uint32, tag = "2")]
        value: u32,
    }

    fn encode_len_prefixed(msg: &TestMsg) -> Vec<u8> {
        let mut buf = Vec::new();
        let mut encoded = Vec::new();
        msg.encode(&mut encoded).unwrap();
        // write varint length
        let mut len = encoded.len();
        while len >= 0x80 {
            buf.push((len as u8) | 0x80);
            len >>= 7;
        }
        buf.push(len as u8);
        buf.extend_from_slice(&encoded);
        buf
    }

    #[test]
    fn normal_parse_single_message() {
        let msg = TestMsg {
            name: "alice".into(),
            value: 42,
        };
        let data = encode_len_prefixed(&msg);
        let mut codec = ProtobufLenPrefixCodec::<TestMsg>::new_with_max_length(1024);
        let mut buf = BytesMut::from(&data[..]);
        // First decode reads varint, returns None
        assert!(matches!(codec.decode(&mut buf), Ok(None)));
        // Second decode reads the message body
        let result = codec.decode(&mut buf).unwrap();
        assert_eq!(result, Some(msg));
    }

    #[test]
    fn empty_input_returns_none() {
        let mut codec = ProtobufLenPrefixCodec::<TestMsg>::new_with_max_length(1024);
        let mut buf = BytesMut::new();
        assert!(matches!(codec.decode(&mut buf), Ok(None)));
    }

    #[test]
    fn truncated_payload_returns_none() {
        let msg = TestMsg {
            name: "alice".into(),
            value: 42,
        };
        let data = encode_len_prefixed(&msg);
        let mut codec = ProtobufLenPrefixCodec::<TestMsg>::new_with_max_length(1024);
        let mut buf = BytesMut::from(&data[..3]);
        assert!(matches!(codec.decode(&mut buf), Ok(None)));
    }

    #[test]
    fn max_length_exceeded_returns_error() {
        let msg = TestMsg {
            name: "alice".into(),
            value: 42,
        };
        let data = encode_len_prefixed(&msg);
        let mut codec = ProtobufLenPrefixCodec::<TestMsg>::new_with_max_length(5);
        let mut buf = BytesMut::from(&data[..]);
        // First decode reads varint, returns None
        let _ = codec.decode(&mut buf);
        // Second decode: message body (10 bytes) exceeds max_length (5)
        assert!(codec.decode(&mut buf).is_err());
    }

    #[test]
    fn multiple_messages_in_sequence() {
        let m1 = TestMsg {
            name: "a".into(),
            value: 1,
        };
        let m2 = TestMsg {
            name: "b".into(),
            value: 2,
        };
        let mut data = encode_len_prefixed(&m1);
        data.extend_from_slice(&encode_len_prefixed(&m2));

        let mut codec = ProtobufLenPrefixCodec::<TestMsg>::new_with_max_length(1024);
        let mut buf = BytesMut::from(&data[..]);
        // First message: varint + body
        assert!(matches!(codec.decode(&mut buf), Ok(None)));
        assert_eq!(codec.decode(&mut buf).unwrap(), Some(m1));
        // Second message: varint + body
        assert!(matches!(codec.decode(&mut buf), Ok(None)));
        assert_eq!(codec.decode(&mut buf).unwrap(), Some(m2));
    }
}
