nrf-modem 0.10.2

use crate::{
    error::Error,
    socket::{Socket, SocketFamily, SocketProtocol, SocketType, SplitSocketHandle},
    CancellationToken, LteLink,
};
use core::net::SocketAddr;

/// A socket that sends and receives UDP messages
pub struct UdpSocket {
    inner: Socket,
}

macro_rules! impl_receive_from {
    () => {
        /// Try to fill the given buffer with received data.
        /// The part of the buffer that was filled is returned together with the address of the source of the message.
        pub async fn receive_from<'buf>(
            &self,
            buf: &'buf mut [u8],
        ) -> Result<(&'buf mut [u8], SocketAddr), Error> {
            self.receive_from_with_cancellation(buf, &Default::default())
                .await
        }

        /// Try to fill the given buffer with received data.
        /// The part of the buffer that was filled is returned together with the address of the source of the message.
        pub async fn receive_from_with_cancellation<'buf>(
            &self,
            buf: &'buf mut [u8],
            token: &CancellationToken,
        ) -> Result<(&'buf mut [u8], SocketAddr), Error> {
            let (received_len, addr) = self
                .socket()
                .receive_from_with_cancellation(buf, token)
                .await?;
            Ok((&mut buf[..received_len], addr))
        }
    };
}

macro_rules! impl_send_to {
    () => {
        /// Send the given buffer to the given address
        pub async fn send_to(&self, buf: &[u8], addr: SocketAddr) -> Result<(), Error> {
            self.send_to_with_cancellation(buf, addr, &Default::default())
                .await
        }

        /// Send the given buffer to the given address
        pub async fn send_to_with_cancellation(
            &self,
            buf: &[u8],
            addr: SocketAddr,
            token: &CancellationToken,
        ) -> Result<(), Error> {
            self.socket()
                .send_to_with_cancellation(buf, addr, token)
                .await
                .map(|_| ())
        }
    };
}

impl UdpSocket {
    /// Bind a new socket to the given address
    pub async fn bind(addr: SocketAddr) -> Result<Self, Error> {
        Self::bind_with_cancellation(addr, &Default::default()).await
    }

    /// Bind a new socket to the given address
    pub async fn bind_with_cancellation(
        addr: SocketAddr,
        token: &CancellationToken,
    ) -> Result<Self, Error> {
        let lte_link = LteLink::new().await?;

        token.as_result()?;

        let family = match addr {
            SocketAddr::V4(_) => SocketFamily::Ipv4,
            SocketAddr::V6(_) => SocketFamily::Ipv6,
        };

        let socket = Socket::create(family, SocketType::Datagram, SocketProtocol::Udp).await?;

        match unsafe { socket.bind_with_cancellation(addr, token).await } {
            Ok(_) => {
                lte_link.deactivate().await?;
                Ok(UdpSocket { inner: socket })
            }
            Err(e) => {
                lte_link.deactivate().await?;
                socket.deactivate().await?;
                Err(e)
            }
        }
    }

    /// Get the raw underlying file descriptor
    pub fn as_raw_fd(&self) -> i32 {
        self.inner.as_raw_fd()
    }

    fn socket(&self) -> &Socket {
        &self.inner
    }

    /// Split the socket into an owned read and write half
    pub async fn split_owned(self) -> Result<(OwnedUdpReceiveSocket, OwnedUdpSendSocket), Error> {
        let (read_split, write_split) = self.inner.split().await?;

        Ok((
            OwnedUdpReceiveSocket { socket: read_split },
            OwnedUdpSendSocket {
                socket: write_split,
            },
        ))
    }

    /// Split the socket into a borrowed read and write half
    pub fn split(&self) -> (UdpReceiveSocket<'_>, UdpSendSocket<'_>) {
        (
            UdpReceiveSocket { socket: self },
            UdpSendSocket { socket: self },
        )
    }

    impl_receive_from!();
    impl_send_to!();

    /// Deactivates the socket and the LTE link.
    /// A normal drop will do the same thing, but blocking.
    pub async fn deactivate(self) -> Result<(), Error> {
        self.inner.deactivate().await?;
        Ok(())
    }
}

/// A borrowed receive half of a udp socket
pub struct UdpReceiveSocket<'a> {
    socket: &'a UdpSocket,
}

impl UdpReceiveSocket<'_> {
    fn socket(&self) -> &Socket {
        &self.socket.inner
    }

    impl_receive_from!();
}

/// A borrowed send half of a udp socket
pub struct UdpSendSocket<'a> {
    socket: &'a UdpSocket,
}

impl UdpSendSocket<'_> {
    fn socket(&self) -> &Socket {
        &self.socket.inner
    }

    impl_send_to!();
}

/// An owned receive half of a udp socket
pub struct OwnedUdpReceiveSocket {
    socket: SplitSocketHandle,
}

impl OwnedUdpReceiveSocket {
    fn socket(&self) -> &Socket {
        &self.socket
    }

    impl_receive_from!();

    /// Deactivates the socket and the LTE link.
    /// A normal drop will do the same thing, but blocking.
    pub async fn deactivate(self) -> Result<(), Error> {
        self.socket.deactivate().await?;
        Ok(())
    }
}

/// An owned send half of a udp socket
pub struct OwnedUdpSendSocket {
    socket: SplitSocketHandle,
}

impl OwnedUdpSendSocket {
    fn socket(&self) -> &Socket {
        &self.socket
    }

    impl_send_to!();

    /// Deactivates the socket and the LTE link.
    /// A normal drop will do the same thing, but blocking.
    pub async fn deactivate(self) -> Result<(), Error> {
        self.socket.deactivate().await?;
        Ok(())
    }
}