//! Encoding: writing protocol values into a byte stream.
use crate::codec::error::Result;
use crate::context::Context;
use std::{
    collections::{BTreeMap, BTreeSet, BinaryHeap, HashMap, HashSet, LinkedList, VecDeque},
    io::Write,
};

macro_rules! impl_encode {
    (u8) => {
        impl<Data> Encode<Data> for u8 {
            /// write number using big endian
            fn encode(&self, buffer: &mut dyn Write, _: &Context<Data>) -> Result<usize> {
                let size = buffer.write(&self.to_be_bytes())?;
                Ok(size)
            }

            /// write raw bytes
            fn encode_slice(
                slice: &[Self],
                buffer: &mut dyn Write,
                _: &Context<Data>,
            ) -> Result<usize> {
                buffer.write(slice).map_err(Into::into)
            }
        }
    };
    ($t: ty) => {
        impl<Data> Encode<Data> for $t {
            /// write number using big endian
            fn encode(&self, buffer: &mut dyn Write, _: &Context<Data>) -> Result<usize> {
                let size = buffer.write(&self.to_be_bytes())?;
                Ok(size)
            }

            fn encode_slice(
                slice: &[Self],
                buffer: &mut dyn Write,
                _: &Context<Data>,
            ) -> Result<usize> {
                let len = buffer.write(
                    &slice
                        .iter()
                        .flat_map(|n| n.to_be_bytes())
                        .collect::<Box<[u8]>>(),
                )?;
                Ok(len)
            }
        }
    };
}

macro_rules! impl_encode_tuples {
    ($($generic: ident),+) => {
        impl<Data, $($generic),+> Encode<Data> for ($($generic,)+)
        where
            $($generic: Encode<Data>),+
        {
            #[allow(non_snake_case)]
            fn encode(&self, buffer: &mut dyn Write, ctx: &Context<Data>) -> Result<usize> {
                let ($($generic,)+): &($($generic,)+) = self;
                let mut len = 0;
                $(len += $generic.encode(buffer, ctx)?;)+
                Ok(len)
            }
        }
    };
}

/// Write a value into a byte stream
pub trait Encode<Data> {
    fn encode(&self, buffer: &mut dyn Write, ctx: &Context<Data>) -> Result<usize>;
    fn encode_slice(slice: &[Self], buffer: &mut dyn Write, ctx: &Context<Data>) -> Result<usize>
    where
        Self: Sized,
    {
        let mut len = 0;
        for item in slice {
            len += item.encode(buffer, ctx)?;
        }
        Ok(len)
    }
}

// ~ Encode arrays

impl<Data, T: Encode<Data>> Encode<Data> for Vec<T> {
    /// Encode each element of Vec
    /// use `CountPrefix` to prefix the vector with number of elements
    /// use `LenPrefix` to prefix the vector with encoded byte size
    fn encode(&self, buffer: &mut dyn Write, ctx: &Context<Data>) -> Result<usize> {
        T::encode_slice(self, buffer, ctx)
    }
}

impl<Data, T: Encode<Data>> Encode<Data> for VecDeque<T> {
    /// Encode each element of VecDeque
    /// use `CountPrefix` to prefix the VecDeque with number of elements
    /// use `LenPrefix` to prefix the VecDeque with encoded byte size
    fn encode(&self, buffer: &mut dyn Write, ctx: &Context<Data>) -> Result<usize> {
        let (front, _) = self.as_slices();
        T::encode_slice(front, buffer, ctx)
    }
}

impl<Data, T: Encode<Data>> Encode<Data> for LinkedList<T>
where
    T: Clone,
{
    /// Encode each element of LinkedList (use clone..)
    /// use `CountPrefix` to prefix the LinkedList with number of elements
    /// use `LenPrefix` to prefix the LinkedList with encoded byte size
    fn encode(&self, buffer: &mut dyn Write, ctx: &Context<Data>) -> Result<usize> {
        let mut view = Vec::with_capacity(self.len());
        view.extend(self.iter().cloned());
        T::encode_slice(&view, buffer, ctx)
    }
}

// ~ Encode slices

impl<Data, T: Encode<Data>> Encode<Data> for &[T] {
    /// Encode each element of slice (use clone..)
    /// use `CountPrefix` to prefix the slice with number of elements
    /// use `LenPrefix` to prefix the slice with encoded byte size
    fn encode(&self, buffer: &mut dyn Write, ctx: &Context<Data>) -> Result<usize> {
        T::encode_slice(self, buffer, ctx)
    }
}

impl<Data, T: Encode<Data>, const S: usize> Encode<Data> for [T; S] {
    /// Encode each element of slice (use clone..)
    /// use `CountPrefix` to prefix the slice with number of elements
    /// use `LenPrefix` to prefix the slice with encoded byte size
    fn encode(&self, buffer: &mut dyn Write, ctx: &Context<Data>) -> Result<usize> {
        T::encode_slice(self, buffer, ctx)
    }
}

impl<Data, T: Encode<Data>> Encode<Data> for [T] {
    /// Encode each element of slice (use clone..)
    /// use `CountPrefix` to prefix the slice with number of elements
    /// use `LenPrefix` to prefix the slice with encoded byte size
    fn encode(&self, buffer: &mut dyn Write, ctx: &Context<Data>) -> Result<usize> {
        T::encode_slice(self, buffer, ctx)
    }
}

// ~ Encode set

impl<Data, T: Encode<Data>> Encode<Data> for HashSet<T>
where
    T: Clone,
{
    /// Encode each element of HashSet (use clone..)
    /// use `CountPrefix` to prefix the HashSet with number of elements
    /// use `LenPrefix` to prefix the HashSet with encoded byte size
    fn encode(&self, buffer: &mut dyn Write, ctx: &Context<Data>) -> Result<usize> {
        let mut view = Vec::with_capacity(self.len());
        view.extend(self.iter().cloned());
        T::encode_slice(&view, buffer, ctx)
    }
}

impl<Data, T: Encode<Data>> Encode<Data> for BTreeSet<T>
where
    T: Clone,
{
    /// Encode each element of BTreeSet (use clone..)
    /// use `CountPrefix` to prefix the BTreeSet with number of elements
    /// use `LenPrefix` to prefix the BTreeSet with encoded byte size
    fn encode(&self, buffer: &mut dyn Write, ctx: &Context<Data>) -> Result<usize> {
        let mut view = Vec::with_capacity(self.len());
        view.extend(self.iter().cloned());
        T::encode_slice(&view, buffer, ctx)
    }
}

// ~ Encode misc

impl<Data, T> Encode<Data> for BinaryHeap<T>
where
    T: Encode<Data>,
{
    fn encode(&self, buffer: &mut dyn Write, ctx: &Context<Data>) -> Result<usize> {
        T::encode_slice(self.as_slice(), buffer, ctx)
    }
}

// ~ Encode maps

impl<Data, K, V> Encode<Data> for HashMap<K, V>
where
    K: Encode<Data> + Clone,
    V: Encode<Data> + Clone,
{
    /// Encode each element of HashMap (use clone..)
    /// use `CountPrefix` to prefix the HashMap with number of pairs
    /// use `LenPrefix` to prefix the HashMap with encoded byte size
    fn encode(&self, buffer: &mut dyn Write, ctx: &Context<Data>) -> Result<usize> {
        let slice: Vec<(K, V)> = self.clone().into_iter().collect();
        <(K, V) as Encode<Data>>::encode_slice(&slice, buffer, ctx)
    }
}

impl<Data, K, V> Encode<Data> for BTreeMap<K, V>
where
    K: Encode<Data> + Clone,
    V: Encode<Data> + Clone,
{
    /// Encode each element of BTreeMap (use clone..)
    /// use `CountPrefix` to prefix the BTreeMap with number of pairs
    /// use `LenPrefix` to prefix the BTreeMap with encoded byte size
    fn encode(&self, buffer: &mut dyn Write, ctx: &Context<Data>) -> Result<usize> {
        let slice: Vec<(K, V)> = self.clone().into_iter().collect();
        <(K, V) as Encode<Data>>::encode_slice(&slice, buffer, ctx)
    }
}

// ~ Encode string
impl<Data> Encode<Data> for String {
    /// Encode the string using utf8
    /// use `CountPrefix` to prefix the string with char length
    /// use `LenPrefix` to prefix the string with utf8 bytes length
    fn encode(&self, buffer: &mut dyn Write, _: &Context<Data>) -> Result<usize> {
        let utf8 = self.as_bytes();
        buffer.write_all(utf8)?;
        Ok(utf8.len())
    }
}

impl<Data> Encode<Data> for &str {
    /// Encode the string using utf8
    /// use `CountPrefix` to prefix the string with char length
    /// use `LenPrefix` to prefix the string with utf8 bytes length
    fn encode(&self, buffer: &mut dyn Write, _: &Context<Data>) -> Result<usize> {
        let utf8 = self.as_bytes();
        buffer.write_all(utf8)?;
        Ok(utf8.len())
    }
}

impl<Data, T: Encode<Data>> Encode<Data> for Option<T> {
    /// Encode `T` if Some(T) or do nothing if None
    fn encode(&self, buffer: &mut dyn Write, ctx: &Context<Data>) -> Result<usize> {
        match self {
            Some(val) => {
                let size = val.encode(buffer, ctx)?;
                Ok(size)
            }
            None => Ok(0),
        }
    }
}

// ~ Encode primitive
impl<Data> Encode<Data> for bool {
    fn encode(&self, buffer: &mut dyn Write, _: &Context<Data>) -> Result<usize> {
        let size = buffer.write(&[*self as u8])?;
        Ok(size)
    }

    fn encode_slice(slice: &[Self], buffer: &mut dyn Write, _: &Context<Data>) -> Result<usize>
    where
        Self: Sized,
    {
        let len = buffer.write(&slice.iter().map(|n| *n as u8).collect::<Box<[u8]>>())?;
        Ok(len)
    }
}

impl_encode!(u8);
impl_encode!(u16);
impl_encode!(u32);
impl_encode!(u64);
impl_encode!(u128);
impl_encode!(usize);

impl_encode!(i8);
impl_encode!(i16);
impl_encode!(i32);
impl_encode!(i64);
impl_encode!(i128);
impl_encode!(isize);

#[cfg(feature = "f16")]
impl_encode!(f16);
impl_encode!(f32);
impl_encode!(f64);
#[cfg(feature = "f128")]
impl_encode!(f128);

// ~ Encode tuple

impl<Data> Encode<Data> for () {
    fn encode(&self, _: &mut dyn Write, _: &Context<Data>) -> Result<usize> {
        Ok(0)
    }

    fn encode_slice(_: &[Self], _: &mut dyn Write, _: &Context<Data>) -> Result<usize>
    where
        Self: Sized,
    {
        Ok(0)
    }
}

impl_encode_tuples!(A);
impl_encode_tuples!(A, B);
impl_encode_tuples!(A, B, C);
impl_encode_tuples!(A, B, C, D);
impl_encode_tuples!(A, B, C, D, E);
impl_encode_tuples!(A, B, C, D, E, F);
impl_encode_tuples!(A, B, C, D, E, F, G);
impl_encode_tuples!(A, B, C, D, E, F, G, H);
impl_encode_tuples!(A, B, C, D, E, F, G, H, I);
impl_encode_tuples!(A, B, C, D, E, F, G, H, I, J);
impl_encode_tuples!(A, B, C, D, E, F, G, H, I, J, K);
impl_encode_tuples!(A, B, C, D, E, F, G, H, I, J, K, L);
impl_encode_tuples!(A, B, C, D, E, F, G, H, I, J, K, L, M);
impl_encode_tuples!(A, B, C, D, E, F, G, H, I, J, K, L, M, N);
impl_encode_tuples!(A, B, C, D, E, F, G, H, I, J, K, L, M, N, O);
impl_encode_tuples!(A, B, C, D, E, F, G, H, I, J, K, L, M, N, O, P);
impl_encode_tuples!(A, B, C, D, E, F, G, H, I, J, K, L, M, N, O, P, Q);