use std::{
    marker::PhantomData,
    net::{IpAddr, Ipv4Addr, Ipv6Addr, SocketAddr, SocketAddrV4, SocketAddrV6},
};

use buffers::{ByteBuf, ByteBufOwned};
use bytes::BytesMut;
use clone_to_owned::CloneToOwned;
use serde::{Deserialize, Serialize};

mod small_slice {
    pub struct SmallSlice<const N: usize> {
        data: [u8; N],
        len: usize,
    }

    impl<const N: usize> AsRef<[u8]> for SmallSlice<N> {
        fn as_ref(&self) -> &[u8] {
            self.as_slice()
        }
    }

    impl<const N: usize> SmallSlice<N> {
        #[allow(clippy::new_without_default)]
        pub fn new() -> Self {
            Self {
                data: [0u8; N],
                len: 0,
            }
        }

        pub fn new_from_buf(buf: &[u8]) -> Self {
            let mut s = Self::new();
            s.extend(buf);
            s
        }

        pub fn extend(&mut self, buf: &[u8]) {
            self.data[self.len..self.len + buf.len()].copy_from_slice(buf);
            self.len += buf.len();
        }

        pub fn as_slice(&self) -> &[u8] {
            &self.data[..self.len]
        }
    }
}

pub use small_slice::SmallSlice;

#[derive(Clone, Copy, PartialEq, Eq)]
pub struct Compact<T>(pub T);

impl<T> From<T> for Compact<T> {
    fn from(value: T) -> Self {
        Self(value)
    }
}

impl<T: core::fmt::Debug> core::fmt::Debug for Compact<T> {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        self.0.fmt(f)
    }
}

pub type CompactIpV4 = Compact<Ipv4Addr>;
pub type CompactIpV6 = Compact<Ipv6Addr>;
pub type CompactIpAddr = Compact<IpAddr>;
pub type CompactSocketAddr4 = Compact<SocketAddrV4>;
pub type CompactSocketAddr6 = Compact<SocketAddrV6>;
pub type CompactSocketAddr = Compact<SocketAddr>;

pub trait CompactSerialize: Sized {
    type Slice: AsRef<[u8]>;

    fn expecting() -> &'static str;
    fn as_slice(&self) -> Self::Slice;
    fn from_slice_unchecked_len(buf: &[u8]) -> Self {
        Self::from_slice(buf).unwrap()
    }
    fn from_slice(buf: &[u8]) -> Option<Self>;
}

pub trait CompactSerializeFixedLen {
    fn fixed_len() -> usize;
}

impl CompactSerialize for Ipv4Addr {
    type Slice = [u8; 4];

    fn as_slice(&self) -> Self::Slice {
        self.octets()
    }

    fn from_slice(buf: &[u8]) -> Option<Self> {
        Some(Self::from(
            TryInto::<[u8; 4]>::try_into(buf.get(..4)?).unwrap(),
        ))
    }

    fn from_slice_unchecked_len(buf: &[u8]) -> Self {
        Self::from(TryInto::<[u8; 4]>::try_into(buf).unwrap())
    }

    fn expecting() -> &'static str {
        "4 bytes for IPv4"
    }
}

impl CompactSerializeFixedLen for Ipv4Addr {
    fn fixed_len() -> usize {
        4
    }
}

impl CompactSerialize for Ipv6Addr {
    type Slice = [u8; 16];

    fn as_slice(&self) -> Self::Slice {
        self.octets()
    }

    fn from_slice(buf: &[u8]) -> Option<Self> {
        Some(Self::from(
            TryInto::<[u8; 16]>::try_into(buf.get(..16)?).unwrap(),
        ))
    }

    fn from_slice_unchecked_len(buf: &[u8]) -> Self {
        Self::from(TryInto::<[u8; 16]>::try_into(buf).unwrap())
    }

    fn expecting() -> &'static str {
        "16 bytes for IPv6"
    }
}

impl CompactSerializeFixedLen for Ipv6Addr {
    fn fixed_len() -> usize {
        16
    }
}

impl CompactSerialize for IpAddr {
    type Slice = SmallSlice<16>;

    fn as_slice(&self) -> Self::Slice {
        match self {
            IpAddr::V4(a) => SmallSlice::new_from_buf(&a.as_slice()),
            IpAddr::V6(a) => SmallSlice::new_from_buf(&a.as_slice()),
        }
    }

    fn from_slice(buf: &[u8]) -> Option<Self> {
        match buf.len() {
            4 => Some(Ipv4Addr::from_slice_unchecked_len(buf).into()),
            16 => Some(Ipv6Addr::from_slice_unchecked_len(buf).into()),
            _ => None,
        }
    }

    fn expecting() -> &'static str {
        "16 bytes for IPv6 or 4 bytes for IPv4"
    }
}

impl CompactSerialize for SocketAddrV4 {
    type Slice = [u8; 6];

    fn as_slice(&self) -> Self::Slice {
        let mut data = [0u8; 6];
        data[..4].copy_from_slice(&self.ip().octets());
        data[4..6].copy_from_slice(&self.port().to_be_bytes());
        data
    }

    fn from_slice(buf: &[u8]) -> Option<Self> {
        if buf.len() != 6 {
            return None;
        }
        Some(Self::from_slice_unchecked_len(buf))
    }

    fn from_slice_unchecked_len(buf: &[u8]) -> Self {
        let ip = Ipv4Addr::from_slice_unchecked_len(&buf[..4]);
        let port = u16::from_be_bytes([buf[4], buf[5]]);
        SocketAddrV4::new(ip, port)
    }

    fn expecting() -> &'static str {
        "6 bytes for SocketAddrV4"
    }
}

impl CompactSerializeFixedLen for SocketAddrV4 {
    fn fixed_len() -> usize {
        6
    }
}

impl CompactSerialize for SocketAddrV6 {
    type Slice = [u8; 18];

    fn as_slice(&self) -> Self::Slice {
        let mut data = [0u8; 18];
        data[..16].copy_from_slice(&self.ip().octets());
        data[16..18].copy_from_slice(&self.port().to_be_bytes());
        data
    }

    fn from_slice(buf: &[u8]) -> Option<Self> {
        if buf.len() != 18 {
            return None;
        }
        Some(Self::from_slice_unchecked_len(buf))
    }

    fn from_slice_unchecked_len(buf: &[u8]) -> Self {
        let ip = Ipv6Addr::from_slice_unchecked_len(&buf[..16]);
        let port = u16::from_be_bytes([buf[16], buf[17]]);
        SocketAddrV6::new(ip, port, 0, 0)
    }

    fn expecting() -> &'static str {
        "18 bytes for SocketAddrV6"
    }
}

impl CompactSerializeFixedLen for SocketAddrV6 {
    fn fixed_len() -> usize {
        18
    }
}

impl CompactSerialize for SocketAddr {
    type Slice = SmallSlice<18>;

    fn as_slice(&self) -> Self::Slice {
        let mut s = SmallSlice::new_from_buf(self.ip().as_slice().as_ref());
        s.extend(&self.port().to_be_bytes());
        s
    }

    fn from_slice(buf: &[u8]) -> Option<Self> {
        match buf.len() {
            6 => Some(SocketAddrV4::from_slice_unchecked_len(buf).into()),
            18 => Some(SocketAddrV6::from_slice_unchecked_len(buf).into()),
            _ => None,
        }
    }

    fn expecting() -> &'static str {
        "18 bytes for SocketAddrV6 or 6 bytes for SocketAddrV4"
    }
}

impl<T: CompactSerialize> Serialize for Compact<T> {
    fn serialize<S>(&self, serializer: S) -> Result<S::Ok, S::Error>
    where
        S: serde::Serializer,
    {
        serializer.serialize_bytes(self.0.as_slice().as_ref())
    }
}

impl<'de, T: CompactSerialize> Deserialize<'de> for Compact<T> {
    fn deserialize<D>(deserializer: D) -> Result<Self, D::Error>
    where
        D: serde::Deserializer<'de>,
    {
        struct Visitor<T> {
            _phantom: PhantomData<T>,
        }
        impl<'de, T: CompactSerialize> serde::de::Visitor<'de> for Visitor<T> {
            type Value = Compact<T>;

            fn expecting(&self, f: &mut std::fmt::Formatter) -> std::fmt::Result {
                f.write_str(T::expecting())
            }

            fn visit_bytes<E>(self, v: &[u8]) -> Result<Self::Value, E>
            where
                E: serde::de::Error,
            {
                match T::from_slice(v) {
                    Some(v) => Ok(Compact(v)),
                    None => Err(E::custom(T::expecting())),
                }
            }
        }
        deserializer.deserialize_bytes(Visitor {
            _phantom: Default::default(),
        })
    }
}

#[derive(Eq, PartialEq)]
pub struct CompactListInBuffer<Buf, T> {
    buf: Buf,
    _phantom: PhantomData<T>,
}

impl<Buf, T> core::fmt::Debug for CompactListInBuffer<Buf, T>
where
    Buf: AsRef<[u8]>,
    T: core::fmt::Debug + CompactSerialize + CompactSerializeFixedLen,
{
    fn fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result {
        struct IterDebug<I>(I);
        impl<I> core::fmt::Debug for IterDebug<I>
        where
            I: Iterator + Clone,
            I::Item: core::fmt::Debug,
        {
            fn fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result {
                f.debug_list().entries(self.0.clone()).finish()
            }
        }
        write!(f, "{:?}", IterDebug(self.iter()))
    }
}

pub type CompactListInBufferOwned<T> = CompactListInBuffer<ByteBufOwned, T>;

impl<T> CompactListInBufferOwned<T>
where
    T: CompactSerialize + CompactSerializeFixedLen,
{
    pub fn new_from_iter(it: impl Iterator<Item = T>) -> Self {
        let mut b = BytesMut::new();
        for item in it {
            b.extend_from_slice(item.as_slice().as_ref());
        }
        Self {
            buf: b.freeze().into(),
            _phantom: Default::default(),
        }
    }
}

impl<Buf, T> CompactListInBuffer<Buf, T>
where
    Buf: AsRef<[u8]>,
    T: CompactSerialize + CompactSerializeFixedLen,
{
    pub fn new_from_buf(buf: Buf) -> Self {
        Self {
            buf,
            _phantom: Default::default(),
        }
    }
    pub fn is_empty(&self) -> bool {
        self.buf.as_ref().is_empty()
    }

    pub fn iter(&self) -> impl Iterator<Item = T> + Clone {
        self.buf
            .as_ref()
            .chunks_exact(T::fixed_len())
            .map(|chunk| T::from_slice_unchecked_len(chunk))
    }

    pub fn get(&self, idx: usize) -> Option<T> {
        let offset = idx * T::fixed_len();
        let end = offset + T::fixed_len();
        let b = self.buf.as_ref().get(offset..end)?;
        Some(T::from_slice_unchecked_len(b))
    }

    pub fn as_borrowed(&self) -> CompactListInBuffer<ByteBuf<'_>, T> {
        CompactListInBuffer {
            buf: self.buf.as_ref().into(),
            _phantom: Default::default(),
        }
    }
}

impl<Buf, T> Serialize for CompactListInBuffer<Buf, T>
where
    Buf: AsRef<[u8]>,
    T: CompactSerialize + CompactSerializeFixedLen,
{
    fn serialize<S>(&self, serializer: S) -> Result<S::Ok, S::Error>
    where
        S: serde::Serializer,
    {
        serializer.serialize_bytes(self.buf.as_ref())
    }
}

impl<Buf, T> CloneToOwned for CompactListInBuffer<Buf, T>
where
    Buf: CloneToOwned,
{
    type Target = CompactListInBuffer<Buf::Target, T>;

    fn clone_to_owned(&self, within_buffer: Option<&bytes::Bytes>) -> Self::Target {
        CompactListInBuffer {
            buf: self.buf.clone_to_owned(within_buffer),
            _phantom: Default::default(),
        }
    }
}

impl<'de, Buf, T> Deserialize<'de> for CompactListInBuffer<Buf, T>
where
    Buf: Deserialize<'de> + AsRef<[u8]>,
    T: CompactSerialize + CompactSerializeFixedLen,
{
    fn deserialize<D>(deserializer: D) -> Result<Self, D::Error>
    where
        D: serde::Deserializer<'de>,
    {
        let buf = Buf::deserialize(deserializer)?;
        let len = buf.as_ref().len();
        if len % T::fixed_len() != 0 {
            struct Exp(&'static str);
            impl serde::de::Expected for Exp {
                fn fmt(&self, f: &mut std::fmt::Formatter) -> std::fmt::Result {
                    write!(f, "an exact multiple of [{}]", self.0)
                }
            }
            return Err(<D::Error as serde::de::Error>::invalid_length(
                len,
                &Exp(T::expecting()),
            ));
        }
        Ok(Self {
            buf,
            _phantom: Default::default(),
        })
    }
}

#[cfg(test)]
mod tests {
    use std::{
        fmt::Debug,
        net::{IpAddr, Ipv4Addr, Ipv6Addr, SocketAddr, SocketAddrV4, SocketAddrV6},
    };

    use bencode::bencode_serialize_to_writer;
    use buffers::ByteBuf;

    use crate::compact_ip::{
        Compact, CompactListInBuffer, CompactListInBufferOwned, CompactSerialize,
        CompactSerializeFixedLen,
    };

    fn test_same_impl<T: CompactSerialize + Copy + Debug + PartialEq + Eq>(input: T) {
        println!("{input:?}");
        let input = Compact(input);
        let mut w = Vec::new();
        bencode_serialize_to_writer(input, &mut w).unwrap();
        println!("input={input:?}, serialized={w:?}");
        let deserialized: Compact<T> = bencode::from_bytes(&w).unwrap();
        assert_eq!(deserialized, input);
    }

    fn test_same_list_impl<
        T: CompactSerialize + CompactSerializeFixedLen + Copy + Debug + PartialEq + Eq,
        const N: usize,
    >(
        input: [T; N],
    ) {
        let l = CompactListInBufferOwned::new_from_iter(input.into_iter());
        let mut w = Vec::new();
        bencode_serialize_to_writer(&l, &mut w).unwrap();
        let deserialized: CompactListInBuffer<ByteBuf, T> = bencode::from_bytes(&w).unwrap();
        assert_eq!(deserialized.iter().collect::<Vec<_>>(), input);
    }

    #[test]
    fn test_invalid_compact_list() {
        let err = bencode::from_bytes::<CompactListInBuffer<ByteBuf, SocketAddrV4>>(b"5:aaaaa")
            .unwrap_err()
            .to_string();
        assert!(err.contains("invalid length"), "{}", err)
    }

    fn test_compact_one<A: CompactSerialize + Copy + PartialEq + Debug>(addr: A, data: &[u8]) {
        let v = bencode::from_bytes::<Compact<A>>(data).unwrap();
        assert_eq!(addr, v.0);
    }

    #[test]
    fn test_compact() {
        test_compact_one(Ipv4Addr::new(1, 2, 3, 4), b"4:\x01\x02\x03\x04");
        test_compact_one(
            Ipv6Addr::new(1, 2, 3, 4, 5, 6, 7, 8),
            b"16:\x00\x01\x00\x02\x00\x03\x00\x04\x00\x05\x00\x06\x00\x07\x00\x08",
        );
        test_compact_one(IpAddr::V4(Ipv4Addr::new(1, 2, 3, 4)), b"4:\x01\x02\x03\x04");
        test_compact_one(
            IpAddr::V6(Ipv6Addr::new(1, 2, 3, 4, 5, 6, 7, 8)),
            b"16:\x00\x01\x00\x02\x00\x03\x00\x04\x00\x05\x00\x06\x00\x07\x00\x08",
        );

        test_compact_one(
            SocketAddrV4::new(Ipv4Addr::new(1, 2, 3, 4), 1),
            b"6:\x01\x02\x03\x04\x00\x01",
        );
        test_compact_one(
            SocketAddrV6::new(Ipv6Addr::new(1, 2, 3, 4, 5, 6, 7, 8), 1, 0, 0),
            b"18:\x00\x01\x00\x02\x00\x03\x00\x04\x00\x05\x00\x06\x00\x07\x00\x08\x00\x01",
        );
        test_compact_one(
            SocketAddr::V4(SocketAddrV4::new(Ipv4Addr::new(1, 2, 3, 4), 1)),
            b"6:\x01\x02\x03\x04\x00\x01",
        );
        test_compact_one(
            SocketAddr::V6(SocketAddrV6::new(
                Ipv6Addr::new(1, 2, 3, 4, 5, 6, 7, 8),
                1,
                0,
                0,
            )),
            b"18:\x00\x01\x00\x02\x00\x03\x00\x04\x00\x05\x00\x06\x00\x07\x00\x08\x00\x01",
        );
    }

    #[test]
    fn test_compact_list() {
        let addrs = bencode::from_bytes::<CompactListInBuffer<ByteBuf, SocketAddrV4>>(
            b"12:\x00\x00\x00\x00\x00\xff\x01\x02\x03\x04\xaa\xaa",
        )
        .unwrap();
        const FIRST: SocketAddrV4 = SocketAddrV4::new(Ipv4Addr::UNSPECIFIED, 255);
        const SECOND: SocketAddrV4 = SocketAddrV4::new(Ipv4Addr::new(1, 2, 3, 4), 43690);
        assert_eq!(addrs.get(0), Some(FIRST));
        assert_eq!(addrs.get(1), Some(SECOND));
        assert_eq!(addrs.get(2), None);

        assert_eq!(addrs.iter().next(), Some(FIRST));
        assert_eq!(addrs.iter().nth(1), Some(SECOND));
        assert_eq!(addrs.iter().nth(2), None);
    }

    #[test]
    fn test_same() {
        test_same_list_impl([Ipv4Addr::LOCALHOST, Ipv4Addr::BROADCAST]);
        test_same_list_impl([Ipv6Addr::LOCALHOST, Ipv6Addr::UNSPECIFIED]);
        test_same_list_impl([
            SocketAddrV4::new(Ipv4Addr::LOCALHOST, 10),
            SocketAddrV4::new(Ipv4Addr::BROADCAST, 11),
        ]);
        test_same_list_impl([
            SocketAddrV6::new(Ipv6Addr::LOCALHOST, 10, 0, 0),
            SocketAddrV6::new(Ipv6Addr::UNSPECIFIED, 11, 0, 0),
        ]);

        test_same_impl(Ipv4Addr::LOCALHOST);
        test_same_impl(Ipv6Addr::LOCALHOST);
        test_same_impl(IpAddr::V4(Ipv4Addr::LOCALHOST));
        test_same_impl(IpAddr::V6(Ipv6Addr::LOCALHOST));
        test_same_impl(SocketAddrV4::new(Ipv4Addr::LOCALHOST, 100));
        test_same_impl(SocketAddrV6::new(Ipv6Addr::LOCALHOST, 100, 0, 0));
        test_same_impl(SocketAddr::new(IpAddr::V4(Ipv4Addr::LOCALHOST), 100));
        test_same_impl(SocketAddr::new(IpAddr::V6(Ipv6Addr::LOCALHOST), 100));
    }
}
