use crate::codec::error::Result;
use crate::context::Context;
use std::collections::{BTreeMap, BTreeSet, BinaryHeap, HashMap, HashSet, LinkedList, VecDeque};

macro_rules! impl_encode {
    ($t:ty) => {
        impl<Data> Encode<Data> for $t {
            fn encode(&self, buffer: &mut Vec<u8>, _: &Context<Data>) -> Result<usize> {
                buffer.extend_from_slice(&self.to_be_bytes());
                Ok(std::mem::size_of::<$t>())
            }

            fn encode_slice(slice: &[Self], buffer: &mut Vec<u8>, _: &Context<Data>) -> Result<usize>
            where
                Self: Sized,
            {
                let len = slice.len() * std::mem::size_of::<$t>();
                buffer.reserve(len);
                for item in slice {
                    buffer.extend_from_slice(&item.to_be_bytes());
                }
                Ok(len)
            }

            fn encode_iter<'a, I>(iter: I, buffer: &mut Vec<u8>, ctx: &Context<Data>) -> Result<usize>
            where
                Self: 'a + Sized,
                I: IntoIterator<Item = &'a Self>,
            {
                let mut len = 0;
                for item in iter {
                    len += item.encode(buffer, ctx)?;
                }
                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 Vec<u8>, ctx: &Context<Data>) -> Result<usize> {
                let ($($generic,)+): &($($generic,)+) = self;
                let mut len = 0;
                $(len += $generic.encode(buffer, ctx)?;)+
                Ok(len)
            }
        }
    };
}

/// Zero-copy encoding: writes directly into a byte buffer
pub trait Encode<Data> {
    fn encode(&self, buffer: &mut Vec<u8>, ctx: &Context<Data>) -> Result<usize>;

    fn encode_slice(slice: &[Self], buffer: &mut Vec<u8>, ctx: &Context<Data>) -> Result<usize>
    where
        Self: Sized,
    {
        let mut len = 0;
        for item in slice {
            len += item.encode(buffer, ctx)?;
        }
        Ok(len)
    }

    fn encode_iter<'a, I>(iter: I, buffer: &mut Vec<u8>, ctx: &Context<Data>) -> Result<usize>
    where
        Self: 'a + Sized,
        I: IntoIterator<Item = &'a Self>,
    {
        let mut len = 0;
        for item in iter {
            len += item.encode(buffer, ctx)?;
        }
        Ok(len)
    }
}

// ── Encode arrays ───────────────────────────────────────────────────────

impl<Data, T: Encode<Data>> Encode<Data> for Vec<T> {
    fn encode(&self, buffer: &mut Vec<u8>, ctx: &Context<Data>) -> Result<usize> {
        T::encode_slice(self, buffer, ctx)
    }
}

impl<Data, T: Encode<Data>> Encode<Data> for VecDeque<T> {
    fn encode(&self, buffer: &mut Vec<u8>, ctx: &Context<Data>) -> Result<usize> {
        let (front, back) = self.as_slices();
        let mut len = 0;
        len += T::encode_slice(front, buffer, ctx)?;
        len += T::encode_slice(back, buffer, ctx)?;
        Ok(len)
    }
}

impl<Data, T: Encode<Data>> Encode<Data> for LinkedList<T> {
    fn encode(&self, buffer: &mut Vec<u8>, ctx: &Context<Data>) -> Result<usize> {
        let mut len = 0;
        for item in self {
            len += item.encode(buffer, ctx)?;
        }
        Ok(len)
    }
}

// ── Encode slices ───────────────────────────────────────────────────────

impl<Data, T: Encode<Data>> Encode<Data> for &[T] {
    fn encode(&self, buffer: &mut Vec<u8>, ctx: &Context<Data>) -> Result<usize> {
        T::encode_slice(self, buffer, ctx)
    }
}

impl<Data, T: Encode<Data>, const S: usize> Encode<Data> for [T; S] {
    fn encode(&self, buffer: &mut Vec<u8>, ctx: &Context<Data>) -> Result<usize> {
        T::encode_slice(self, buffer, ctx)
    }
}

impl<Data, T: Encode<Data>> Encode<Data> for [T] {
    fn encode(&self, buffer: &mut Vec<u8>, ctx: &Context<Data>) -> Result<usize> {
        T::encode_slice(self, buffer, ctx)
    }
}

// ── Encode references to collections ────────────────────────────────────

impl<Data, T: Encode<Data>> Encode<Data> for &Vec<T> {
    fn encode(&self, buffer: &mut Vec<u8>, ctx: &Context<Data>) -> Result<usize> {
        T::encode_slice(self, buffer, ctx)
    }
}

impl<Data, T: Encode<Data>> Encode<Data> for &VecDeque<T> {
    fn encode(&self, buffer: &mut Vec<u8>, ctx: &Context<Data>) -> Result<usize> {
        let (front, back) = self.as_slices();
        let mut len = 0;
        len += T::encode_slice(front, buffer, ctx)?;
        len += T::encode_slice(back, buffer, ctx)?;
        Ok(len)
    }
}

impl<Data, T: Encode<Data>> Encode<Data> for &LinkedList<T> {
    fn encode(&self, buffer: &mut Vec<u8>, ctx: &Context<Data>) -> Result<usize> {
        let mut len = 0;
        for item in (*self).iter() {
            len += item.encode(buffer, ctx)?;
        }
        Ok(len)
    }
}

impl<Data, T: Encode<Data>> Encode<Data> for &HashSet<T> {
    fn encode(&self, buffer: &mut Vec<u8>, ctx: &Context<Data>) -> Result<usize> {
        T::encode_iter(self.iter(), buffer, ctx)
    }
}

impl<Data, T: Encode<Data>> Encode<Data> for &BTreeSet<T> {
    fn encode(&self, buffer: &mut Vec<u8>, ctx: &Context<Data>) -> Result<usize> {
        T::encode_iter(self.iter(), buffer, ctx)
    }
}

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

impl<Data, K, V> Encode<Data> for &HashMap<K, V>
where
    K: Encode<Data>,
    V: Encode<Data>,
{
    fn encode(&self, buffer: &mut Vec<u8>, ctx: &Context<Data>) -> Result<usize> {
        let mut len = 0;
        for (k, v) in (*self).iter() {
            len += k.encode(buffer, ctx)?;
            len += v.encode(buffer, ctx)?;
        }
        Ok(len)
    }
}

impl<Data, K, V> Encode<Data> for &BTreeMap<K, V>
where
    K: Encode<Data>,
    V: Encode<Data>,
{
    fn encode(&self, buffer: &mut Vec<u8>, ctx: &Context<Data>) -> Result<usize> {
        let mut len = 0;
        for (k, v) in (*self).iter() {
            len += k.encode(buffer, ctx)?;
            len += v.encode(buffer, ctx)?;
        }
        Ok(len)
    }
}

// ── Encode set (owned) ──────────────────────────────────────────────────

impl<Data, T: Encode<Data>> Encode<Data> for HashSet<T> {
    fn encode(&self, buffer: &mut Vec<u8>, ctx: &Context<Data>) -> Result<usize> {
        T::encode_iter(self.iter(), buffer, ctx)
    }
}

impl<Data, T: Encode<Data>> Encode<Data> for BTreeSet<T> {
    fn encode(&self, buffer: &mut Vec<u8>, ctx: &Context<Data>) -> Result<usize> {
        T::encode_iter(self.iter(), buffer, ctx)
    }
}

// ── Encode misc ─────────────────────────────────────────────────────────

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

// ── Encode maps (owned) ─────────────────────────────────────────────────

impl<Data, K, V> Encode<Data> for HashMap<K, V>
where
    K: Encode<Data>,
    V: Encode<Data>,
{
    fn encode(&self, buffer: &mut Vec<u8>, ctx: &Context<Data>) -> Result<usize> {
        let mut len = 0;
        for (k, v) in self {
            len += k.encode(buffer, ctx)?;
            len += v.encode(buffer, ctx)?;
        }
        Ok(len)
    }
}

impl<Data, K, V> Encode<Data> for BTreeMap<K, V>
where
    K: Encode<Data>,
    V: Encode<Data>,
{
    fn encode(&self, buffer: &mut Vec<u8>, ctx: &Context<Data>) -> Result<usize> {
        let mut len = 0;
        for (k, v) in self {
            len += k.encode(buffer, ctx)?;
            len += v.encode(buffer, ctx)?;
        }
        Ok(len)
    }
}

// ── Encode string ───────────────────────────────────────────────────────

impl<Data> Encode<Data> for String {
    fn encode(&self, buffer: &mut Vec<u8>, _: &Context<Data>) -> Result<usize> {
        let utf8 = self.as_bytes();
        buffer.extend_from_slice(utf8);
        Ok(utf8.len())
    }
}

impl<Data> Encode<Data> for &str {
    fn encode(&self, buffer: &mut Vec<u8>, _: &Context<Data>) -> Result<usize> {
        let utf8 = self.as_bytes();
        buffer.extend_from_slice(utf8);
        Ok(utf8.len())
    }
}

impl<Data, T: Encode<Data>> Encode<Data> for Option<T> {
    fn encode(&self, buffer: &mut Vec<u8>, ctx: &Context<Data>) -> Result<usize> {
        match self {
            Some(val) => val.encode(buffer, ctx),
            None => Ok(0),
        }
    }
}

// ── Encode primitive ────────────────────────────────────────────────────

impl<Data> Encode<Data> for bool {
    fn encode(&self, buffer: &mut Vec<u8>, _: &Context<Data>) -> Result<usize> {
        buffer.push(*self as u8);
        Ok(1)
    }

    fn encode_slice(slice: &[Self], buffer: &mut Vec<u8>, _: &Context<Data>) -> Result<usize>
    where
        Self: Sized,
    {
        let len = slice.len();
        buffer.reserve(len);
        for &b in slice {
            buffer.push(b as u8);
        }
        Ok(len)
    }
}

// ── Encode pointer ──────────────────────────────────────────────────────

impl<Data, T> Encode<Data> for std::sync::Arc<T>
where
    T: Encode<Data>,
{
    fn encode(&self, buffer: &mut Vec<u8>, ctx: &Context<Data>) -> Result<usize> {
        self.as_ref().encode(buffer, ctx)
    }
}

impl<Data, T> Encode<Data> for std::sync::Arc<[T]>
where
    T: Encode<Data>,
{
    fn encode(&self, buffer: &mut Vec<u8>, ctx: &Context<Data>) -> Result<usize> {
        self.as_ref().encode(buffer, ctx)
    }
}

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 Vec<u8>, _: &Context<Data>) -> Result<usize> {
        Ok(0)
    }

    fn encode_slice(_: &[Self], _: &mut Vec<u8>, _: &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);