//! Variable length integer and related types
//! used by the Bitmessage protocol.

use std::{
    fmt,
    io::{self, Read, Write},
    mem::size_of,
};

use serde::{
    de::{self, Visitor},
    Deserialize, Deserializer, Serialize, Serializer,
};

use crate::io::{LenBm, LimitedReadFrom, ReadFrom, TooLongError, WriteTo};

/// A variable length integer type used by the Bitmessage protocol.
#[derive(Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Hash, Debug)]
pub struct VarInt(u64);

impl VarInt {
    /// Constructs a var_int from a value.
    pub fn new(value: u64) -> Self {
        Self(value)
    }

    /// Returns the value as `u64`.
    pub fn as_u64(self) -> u64 {
        self.0
    }
}

impl fmt::Display for VarInt {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        self.0.fmt(f)
    }
}

impl From<u64> for VarInt {
    fn from(value: u64) -> Self {
        Self(value)
    }
}

impl From<usize> for VarInt {
    fn from(value: usize) -> Self {
        Self(value as u64)
    }
}

impl WriteTo for VarInt {
    fn write_to(&self, w: &mut dyn Write) -> io::Result<()> {
        let v = self.0;
        if v < 0xfd {
            (v as u8).write_to(w)
        } else if v <= 0xffff {
            0xfd_u8.write_to(w)?;
            (v as u16).write_to(w)
        } else if v <= 0xffff_ffff {
            0xfe_u8.write_to(w)?;
            (v as u32).write_to(w)
        } else {
            0xff_u8.write_to(w)?;
            v.write_to(w)
        }
    }
}

/// An error which can be returned when parsing a variable length integer.
///
/// This error is used as a error type wrapped by `std::io::Error`
/// for the [`ReadFrom`] implementation for [`VarInt`].
///
/// [`ReadFrom`]: ../io/trait.ReadFrom.html
/// [`VarInt`]: struct.VarInt.html
#[derive(Clone, PartialEq, Eq, Debug)]
pub enum VarIntError {
    /// The actual value was represented by 3 bytes,
    /// but it could be represented by less bytes.
    /// The minimum value allowed and the actual value are returned
    /// as payloads of this variant.
    U16 {
        /// The minimum value allowed.
        min: u16,
        /// The actual value.
        value: u16,
    },

    /// The actual value was represented by 5 bytes,
    /// but it could be represented by less bytes.
    /// The minimum value allowed and the actual value are returned
    /// as payloads of this variant.
    U32 {
        /// The minimum value allowed.
        min: u32,
        /// The actual value.
        value: u32,
    },

    /// The actual value was represented by 9 bytes,
    /// but it could be represented by less bytes.
    /// The minimum value allowed and the actual value are returned
    /// as payloads of this variant.
    U64 {
        /// The minimum value allowed.
        min: u64,
        /// The actual value.
        value: u64,
    },
}

impl fmt::Display for VarIntError {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        match self {
            Self::U16 { min, value } => write!(f, "invalid VarInt: {:04x} < {:04x}", value, min),
            Self::U32 { min, value } => write!(f, "invalid VarInt: {:08x} < {:08x}", value, min),
            Self::U64 { min, value } => write!(f, "invalid var_int: {:016x} < {:016x}", value, min),
        }
    }
}

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

impl ReadFrom for VarInt {
    fn read_from(r: &mut dyn Read) -> io::Result<Self>
    where
        Self: Sized,
    {
        let head = u8::read_from(r)?;
        if head < 0xfd {
            Ok(Self(head as u64))
        } else if head == 0xfd {
            let value = u16::read_from(r)?;
            if value < 0xfd {
                let err = VarIntError::U16 { min: 0xfd, value };
                return Err(io::Error::new(io::ErrorKind::Other, err));
            }
            Ok(Self(value as u64))
        } else if head == 0xfe {
            let value = u32::read_from(r)?;
            if value < 0x10000 {
                let err = VarIntError::U32 {
                    min: 0x10000,
                    value,
                };
                return Err(io::Error::new(io::ErrorKind::Other, err));
            }
            Ok(Self(value as u64))
        } else {
            let value = u64::read_from(r)?;
            if value < 0x0001_0000_0000 {
                let err = VarIntError::U64 {
                    min: 0x0001_0000_0000,
                    value,
                };
                return Err(io::Error::new(io::ErrorKind::Other, err));
            }
            Ok(Self(value))
        }
    }
}

impl LenBm for VarInt {
    fn len_bm(&self) -> usize {
        let v = self.0;
        if v < 0xfd {
            size_of::<u8>()
        } else if v <= 0xffff {
            1 + size_of::<u16>()
        } else if v <= 0xffff_ffff {
            1 + size_of::<u32>()
        } else {
            1 + size_of::<u64>()
        }
    }
}

#[test]
fn test_var_int_write_to() {
    let test: VarInt = 0xcd_u64.into();
    let mut bytes = Vec::new();
    test.write_to(&mut bytes).unwrap();
    let expected = [0xcd];
    assert_eq!(bytes, expected);

    let test: VarInt = 0xfe_u64.into();
    let mut bytes = Vec::new();
    test.write_to(&mut bytes).unwrap();
    let expected = [0xfd, 0x00, 0xfe];
    assert_eq!(bytes, expected);

    let test: VarInt = 0xcdef_u64.into();
    let mut bytes = Vec::new();
    test.write_to(&mut bytes).unwrap();
    let expected = [0xfd, 0xcd, 0xef];
    assert_eq!(bytes, expected);

    let test: VarInt = 0x89ab_cdef_u64.into();
    let mut bytes = Vec::new();
    test.write_to(&mut bytes).unwrap();
    let expected = [0xfe, 0x89, 0xab, 0xcd, 0xef];
    assert_eq!(bytes, expected);

    let test: VarInt = 0x0123_4567_89ab_cdef_u64.into();
    let mut bytes = Vec::new();
    test.write_to(&mut bytes).unwrap();
    let expected = [0xff, 0x01, 0x23, 0x45, 0x67, 0x89, 0xab, 0xcd, 0xef];
    assert_eq!(bytes, expected);
}

#[test]
fn test_var_int_read_from() {
    use std::io::Cursor;

    let mut bytes = Cursor::new([0xcd]);
    let test = VarInt::read_from(&mut bytes).unwrap();
    let expected: VarInt = 0xcd_u64.into();
    assert_eq!(test, expected);

    let mut bytes = Cursor::new([0xfd, 0x00, 0xfe]);
    let test = VarInt::read_from(&mut bytes).unwrap();
    let expected: VarInt = 0xfe_u64.into();
    assert_eq!(test, expected);

    let mut bytes = Cursor::new([0xfe, 0x89, 0xab, 0xcd, 0xef]);
    let test = VarInt::read_from(&mut bytes).unwrap();
    let expected: VarInt = 0x89ab_cdef_u64.into();
    assert_eq!(test, expected);

    let mut bytes = Cursor::new([0xff, 0x01, 0x23, 0x45, 0x67, 0x89, 0xab, 0xcd, 0xef]);
    let test = VarInt::read_from(&mut bytes).unwrap();
    let expected: VarInt = 0x0123_4567_89ab_cdef_u64.into();
    assert_eq!(test, expected);
}

#[test]
#[should_panic]
fn test_var_int_error_u16() {
    use std::io::Cursor;

    let mut bytes = Cursor::new([0xfd, 0x00, 0xab]);
    let test = VarInt::read_from(&mut bytes).unwrap();
    let expected: VarInt = 0x00ab_u64.into();
    assert_eq!(test, expected);
}

#[test]
#[should_panic]
fn test_var_int_error_u32() {
    use std::io::Cursor;

    let mut bytes = Cursor::new([0xfe, 0x00, 0x00, 0xcd, 0xef]);
    let test = VarInt::read_from(&mut bytes).unwrap();
    let expected: VarInt = 0x0000_cdef_u64.into();
    assert_eq!(test, expected);
}

#[test]
#[should_panic]
fn test_var_int_error_u64() {
    use std::io::Cursor;

    let mut bytes = Cursor::new([0xff, 0x00, 0x00, 0x00, 0x00, 0x89, 0xab, 0xcd, 0xef]);
    let test = VarInt::read_from(&mut bytes).unwrap();
    let expected: VarInt = 0x0000_0000_89ab_cdef_u64.into();
    assert_eq!(test, expected);
}

/// A variable length byte string type used by the Bitmessage protocol.
#[derive(Clone, PartialEq, Eq, PartialOrd, Ord, Hash, Debug)]
pub struct VarStr(Vec<u8>);

impl VarStr {
    /// Constructs a var_str from a byte string.
    pub fn new(bytes: Vec<u8>) -> Self {
        Self(bytes)
    }
}

impl Serialize for VarStr {
    fn serialize<S>(&self, serializer: S) -> Result<S::Ok, S::Error>
    where
        S: Serializer,
    {
        match String::from_utf8(self.0.clone()) {
            Ok(s) => serializer.serialize_str(&s),
            Err(_) => serializer.serialize_bytes(&self.0),
        }
    }
}

impl<'de> Deserialize<'de> for VarStr {
    fn deserialize<D>(deserializer: D) -> Result<Self, D::Error>
    where
        D: Deserializer<'de>,
    {
        deserializer.deserialize_str(VarStrVisitor)
    }
}

struct VarStrVisitor;

impl<'de> Visitor<'de> for VarStrVisitor {
    type Value = VarStr;

    fn expecting(&self, formatter: &mut fmt::Formatter) -> fmt::Result {
        formatter.write_str("a byte array or a string")
    }

    fn visit_str<E>(self, v: &str) -> Result<Self::Value, E>
    where
        E: de::Error,
    {
        Ok(v.as_bytes().to_vec().into())
    }

    fn visit_borrowed_str<E>(self, v: &'de str) -> Result<Self::Value, E>
    where
        E: de::Error,
    {
        Ok(v.as_bytes().to_vec().into())
    }

    fn visit_string<E>(self, v: String) -> Result<Self::Value, E>
    where
        E: de::Error,
    {
        Ok(v.into_bytes().into())
    }

    fn visit_bytes<E>(self, v: &[u8]) -> Result<Self::Value, E>
    where
        E: de::Error,
    {
        Ok(v.to_vec().into())
    }

    fn visit_borrowed_bytes<E>(self, v: &'de [u8]) -> Result<Self::Value, E>
    where
        E: de::Error,
    {
        Ok(v.to_vec().into())
    }

    fn visit_byte_buf<E>(self, v: Vec<u8>) -> Result<Self::Value, E>
    where
        E: de::Error,
    {
        Ok(v.into())
    }
}

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

impl fmt::Display for VarStr {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        String::from_utf8_lossy(&self.0).fmt(f)
    }
}

impl From<Vec<u8>> for VarStr {
    fn from(bytes: Vec<u8>) -> Self {
        Self(bytes)
    }
}

impl WriteTo for VarStr {
    fn write_to(&self, w: &mut dyn Write) -> io::Result<()> {
        VarInt(self.0.len() as u64).write_to(w)?;
        w.write_all(&self.0)
    }
}

impl LimitedReadFrom for VarStr {
    fn limited_read_from(r: &mut dyn Read, max_len: usize) -> io::Result<Self>
    where
        Self: Sized,
    {
        let length = VarInt::read_from(r)?;
        if length.0 > max_len as u64 {
            return Err(io::Error::new(
                io::ErrorKind::Other,
                TooLongError::new(max_len, length.0 as usize),
            ));
        }
        let mut r = r.take(length.0);
        let mut bytes = Vec::with_capacity(length.0 as usize);
        r.read_to_end(&mut bytes)?;
        Ok(Self(bytes))
    }
}

#[test]
fn test_var_str_write_to() {
    let test: VarStr = b"hello".to_vec().into();
    let mut bytes = Vec::new();
    test.write_to(&mut bytes).unwrap();
    let expected = [0x05, b'h', b'e', b'l', b'l', b'o'];
    assert_eq!(bytes, expected);
}

#[test]
fn test_var_str_limited_read_from() {
    use std::io::Cursor;

    let mut bytes = Cursor::new([0x05, b'h', b'e', b'l', b'l', b'o']);
    let test = VarStr::limited_read_from(&mut bytes, 256).unwrap();
    let expected: VarStr = b"hello".to_vec().into();
    assert_eq!(test, expected);
}

/// A variable length list of integers used by the Bitmessage protocol.
#[derive(Clone, PartialEq, Eq, PartialOrd, Ord, Hash, Debug)]
pub struct VarIntList(Vec<VarInt>);

impl AsRef<[VarInt]> for VarIntList {
    fn as_ref(&self) -> &[VarInt] {
        &self.0
    }
}

impl VarIntList {
    /// Constructs var_int_list from a list of
    /// [`VarInt`](struct.VarInt.html)s.
    pub fn new(list: Vec<VarInt>) -> Self {
        Self(list)
    }
}

impl From<Vec<VarInt>> for VarIntList {
    fn from(values: Vec<VarInt>) -> Self {
        Self(values)
    }
}

impl WriteTo for VarIntList {
    fn write_to(&self, w: &mut dyn Write) -> io::Result<()> {
        VarInt(self.0.len() as u64).write_to(w)?;
        for value in &self.0 {
            value.write_to(w)?
        }
        Ok(())
    }
}

impl LimitedReadFrom for VarIntList {
    fn limited_read_from(r: &mut dyn Read, max_len: usize) -> io::Result<Self>
    where
        Self: Sized,
    {
        let length = VarInt::read_from(r)?;
        if length.0 > max_len as u64 {
            return Err(io::Error::new(
                io::ErrorKind::Other,
                TooLongError::new(max_len, length.0 as usize),
            ));
        }
        let mut values = Vec::with_capacity(length.0 as usize);
        for _ in 0..length.0 {
            values.push(VarInt::read_from(r)?);
        }
        Ok(Self(values))
    }
}

#[test]
fn test_var_int_list_write_to() {
    let test: VarIntList = vec![0xcd_u64.into(), 0xfe_u64.into(), 0xcdef_u64.into()].into();
    let mut bytes = Vec::new();
    test.write_to(&mut bytes).unwrap();
    let expected = [0x03, 0xcd, 0xfd, 0x00, 0xfe, 0xfd, 0xcd, 0xef];
    assert_eq!(bytes, expected);
}

#[test]
fn test_var_int_list_limited_read_from() {
    use std::io::Cursor;

    let mut bytes = Cursor::new([0x03, 0xcd, 0xfd, 0x00, 0xfe, 0xfd, 0xcd, 0xef]);
    let test = VarIntList::limited_read_from(&mut bytes, 256).unwrap();
    let expected: VarIntList = vec![0xcd_u64.into(), 0xfe_u64.into(), 0xcdef_u64.into()].into();
    assert_eq!(test, expected);
}
