use crate::command::{Command, CommandHeader};
use crate::udp::UdpRecvPostcardError::ExtraData;
use crate::udp::UdpSendPostcardError::LengthMismatch;
use serde::{Deserialize, Serialize};
use std::io::ErrorKind;
use std::net::{SocketAddr, ToSocketAddrs, UdpSocket};
use thiserror::Error;

#[derive(Error, Debug)]
pub enum UdpRecvPostcardError {
    #[error("IO Error: {0}")]
    Io(#[from] std::io::Error),
    #[error("Deserialization Error: {0}")]
    Deserialization(#[from] postcard::Error),
    #[error("Extra Data: {amount}")]
    ExtraData { amount: usize },
    #[error("No Data")]
    NoData,
}

#[derive(Error, Debug)]
pub enum UdpSendPostcardError {
    #[error("IO Error: {0}")]
    Io(#[from] std::io::Error),
    #[error("Serialization Error: {0}")]
    Serialization(#[from] postcard::Error),
    #[error("Length Mismatch: {expected} expected. {actual} actual")]
    LengthMismatch { expected: usize, actual: usize },
}

pub trait UdpSocketExt {
    /// Receive a Postcard encoded message from this UDP Socket
    ///
    /// # Errors
    /// An error that could have occurred while trying to receive this message.
    /// If no data was received a `UdpRecvCborError::NoData` error would be returned.
    fn recv_postcard<'de, T: Deserialize<'de>>(
        &self,
        buffer: &'de mut [u8],
    ) -> Result<(T, SocketAddr, usize), UdpRecvPostcardError>;

    /// Send a CBOR encoded message to an address using this socket
    ///
    /// # Errors
    /// An error that could have occurred while trying to send this message
    fn send_postcard<T: Serialize + ?Sized, A: ToSocketAddrs>(
        &self,
        data: &T,
        buffer: &mut [u8],
        addr: A,
    ) -> Result<usize, UdpSendPostcardError>;

    /// Send a command message to an address using this socket
    ///
    /// # Errors
    /// An error that could have occurred while trying to send this message
    fn send_command<T: Command, A: ToSocketAddrs>(
        &self,
        name: &str,
        data: &T,
        addr: A,
    ) -> Result<usize, UdpSendPostcardError>;
}

fn recv_postcard_inner<'de, T: Deserialize<'de>>(
    result: std::io::Result<(usize, SocketAddr)>,
    buffer: &'de mut [u8],
) -> Result<(T, SocketAddr, usize), UdpRecvPostcardError> {
    match result {
        Ok((size, addr)) => match postcard::take_from_bytes::<T>(&buffer[..size]) {
            Ok((res, rem)) => {
                if !rem.is_empty() {
                    return Err(ExtraData { amount: rem.len() });
                }
                Ok((res, addr, size))
            }
            Err(err) => Err(err.into()),
        },
        Err(err) => match err.kind() {
            ErrorKind::WouldBlock | ErrorKind::TimedOut => Err(UdpRecvPostcardError::NoData),
            _ => Err(err.into()),
        },
    }
}

fn send_inner(
    send_result: Result<usize, std::io::Error>,
    expected_size: usize,
) -> Result<usize, UdpSendPostcardError> {
    match send_result {
        Ok(size_sent) => {
            if expected_size != size_sent {
                return Err(LengthMismatch {
                    expected: expected_size,
                    actual: size_sent,
                });
            }
            Ok(size_sent)
        }
        Err(e) => Err(e.into()),
    }
}

impl UdpSocketExt for UdpSocket {
    fn recv_postcard<'de, T: Deserialize<'de>>(
        &self,
        buffer: &'de mut [u8],
    ) -> Result<(T, SocketAddr, usize), UdpRecvPostcardError> {
        recv_postcard_inner(self.recv_from(buffer), buffer)
    }

    fn send_postcard<T: Serialize + ?Sized, A: ToSocketAddrs>(
        &self,
        data: &T,
        buffer: &mut [u8],
        addr: A,
    ) -> Result<usize, UdpSendPostcardError> {
        let result = postcard::to_slice(data, buffer)?;
        let size_encoded = result.len();
        send_inner(self.send_to(result, addr), size_encoded)
    }

    fn send_command<T: Command, A: ToSocketAddrs>(
        &self,
        name: &str,
        data: &T,
        addr: A,
    ) -> Result<usize, UdpSendPostcardError> {
        let mut inner_buffer = [0u8; 512];
        let inner_buffer = postcard::to_slice(data, &mut inner_buffer)?;
        let mut buffer = [0u8; 512];
        self.send_postcard(
            &CommandHeader {
                name,
                data: inner_buffer,
            },
            &mut buffer,
            addr,
        )
    }
}

#[cfg(feature = "async")]
pub mod tokio {
    use super::{
        Command, CommandHeader, Deserialize, Serialize, SocketAddr, UdpRecvPostcardError,
        UdpSendPostcardError, recv_postcard_inner, send_inner,
    };
    use ::tokio::net::ToSocketAddrs;
    use ::tokio::net::UdpSocket;

    pub trait AsyncUdpSocketExt {
        /// Receive a Postcard encoded message from this UDP Socket
        ///
        /// # Errors
        /// An error that could have occurred while trying to receive this message.
        /// If no data was received a `UdpRecvCborError::NoData` error would be returned.
        fn recv_postcard<'de, T: Deserialize<'de>>(
            &self,
            buffer: &'de mut [u8],
        ) -> impl Future<Output = Result<(T, SocketAddr, usize), UdpRecvPostcardError>>;

        /// Send a CBOR encoded message to an address using this socket
        ///
        /// # Errors
        /// An error that could have occurred while trying to send this message
        fn send_postcard<T: Serialize + ?Sized, A: ToSocketAddrs>(
            &self,
            data: &T,
            buffer: &mut [u8],
            addr: A,
        ) -> impl Future<Output = Result<usize, UdpSendPostcardError>>;

        /// Send a command message to an address using this socket
        ///
        /// # Errors
        /// An error that could have occurred while trying to send this message
        fn send_command<T: Command, A: ToSocketAddrs>(
            &self,
            name: &str,
            data: &T,
            addr: A,
        ) -> impl Future<Output = Result<usize, UdpSendPostcardError>>;
    }

    impl AsyncUdpSocketExt for UdpSocket {
        async fn recv_postcard<'de, T: Deserialize<'de>>(
            &self,
            buffer: &'de mut [u8],
        ) -> Result<(T, SocketAddr, usize), UdpRecvPostcardError> {
            recv_postcard_inner(self.recv_from(buffer).await, buffer)
        }

        async fn send_postcard<T: Serialize + ?Sized, A: ToSocketAddrs>(
            &self,
            data: &T,
            buffer: &mut [u8],
            addr: A,
        ) -> Result<usize, UdpSendPostcardError> {
            let result = postcard::to_slice(data, buffer)?;
            let size_encoded = result.len();
            send_inner(self.send_to(result, addr).await, size_encoded)
        }

        async fn send_command<T: Command, A: ToSocketAddrs>(
            &self,
            name: &str,
            data: &T,
            addr: A,
        ) -> Result<usize, UdpSendPostcardError> {
            let mut inner_buffer = [0u8; 512];
            let inner_buffer = postcard::to_slice(data, &mut inner_buffer)?;
            let mut buffer = [0u8; 512];
            self.send_postcard(
                &CommandHeader {
                    name,
                    data: inner_buffer,
                },
                &mut buffer,
                addr,
            )
            .await
        }
    }
}
