//! Socket management functions
use core::{net::SocketAddr, time::Duration};

use abi::{
    error::Error,
    io::RawFd,
    net::{options, MessageHeader, MessageHeaderMut, SocketInterfaceQuery, SocketType},
    option::{OptionSizeHint, OptionValue},
};
use abi_serde::wire;

/// Short-hand macro for [get_socket_option].
///
/// Automatically sets up the `buffer` argument for socket options which implement
/// [SocketOptionSizeHint].
pub macro get_socket_option($fd:expr, $variant_ty:ty) {{
    let mut buffer = [0; <$variant_ty as $crate::io::OptionSizeHint>::SIZE_HINT];
    $crate::net::get_socket_option::<$variant_ty>($fd, &mut buffer)
}}

/// Retrieves the value of a socket's option
pub fn get_socket_option<'de, T: OptionValue<'de>>(
    fd: RawFd,
    buffer: &'de mut [u8],
) -> Result<T::Value, Error> {
    let len = unsafe { crate::sys::get_socket_option(fd, T::VARIANT.into(), buffer) }?;
    T::load(&buffer[..len])
}

/// Sets the value of a socket's option, using the buffer provided to marshall the payload
pub fn set_socket_option_with<'de, T: OptionValue<'de>>(
    fd: RawFd,
    buffer: &mut [u8],
    value: &T::Value,
) -> Result<(), Error> {
    let len = T::store(value, buffer)?;
    unsafe { crate::sys::set_socket_option(fd, T::VARIANT.into(), &buffer[..len]) }
}

/// Sets the value of a socket's option, using a buffer set up automatically based on
/// [SocketOptionSizeHint]. See [set_socket_option_with].
pub fn set_socket_option<'de, T: OptionValue<'de> + OptionSizeHint>(
    fd: RawFd,
    value: &T::Value,
) -> Result<(), Error>
where
    [u8; T::SIZE_HINT]: Sized,
{
    let mut buffer = [0; T::SIZE_HINT];
    set_socket_option_with::<T>(fd, &mut buffer, value)
}

/// Returns the socket's local address or an error if it isn't bound yet
pub fn local_address(fd: RawFd) -> Result<SocketAddr, Error> {
    let mut buffer = [0; options::LocalAddress::SIZE_HINT];
    get_socket_option::<options::LocalAddress>(fd, &mut buffer)?.ok_or(Error::InvalidOperation)
}

/// Returns the socket's remote address or an error if the socket is not connected
pub fn peer_address(fd: RawFd) -> Result<SocketAddr, Error> {
    let mut buffer = [0; options::PeerAddress::SIZE_HINT];
    get_socket_option::<options::PeerAddress>(fd, &mut buffer)?.ok_or(Error::NotConnected)
}

fn bind_ip(fd: RawFd, local: &SocketAddr, listen: bool) -> Result<(), Error> {
    let mut address = [0; 32];
    let address_len = wire::to_slice(local, &mut address)?;
    unsafe { crate::sys::bind(fd, &address[..address_len]) }?;
    if listen {
        unsafe { crate::sys::listen(fd) }?;
    }
    Ok(())
}

fn connect_ip(fd: RawFd, remote: &SocketAddr, timeout: Option<Duration>) -> Result<(), Error> {
    let mut address = [0; 32];
    let address_len = wire::to_slice(remote, &mut address)?;
    set_socket_option::<options::ConnectTimeout>(fd, &timeout)?;
    unsafe { crate::sys::connect(fd, &address[..address_len]) }
}

/// Creates a new socket and binds it to a local address
pub fn create_and_bind_ip(
    ty: SocketType,
    local: &SocketAddr,
    listen: bool,
) -> Result<RawFd, Error> {
    let fd = unsafe { crate::sys::create_socket(ty) }?;
    match bind_ip(fd, local, listen) {
        Ok(()) => Ok(fd),
        Err(error) => {
            unsafe { crate::sys::close(fd) }.ok();
            Err(error)
        }
    }
}

/// Binds a TCP listener socket to some local address
pub fn bind_tcp(local: &SocketAddr) -> Result<RawFd, Error> {
    create_and_bind_ip(SocketType::TcpStream, local, true)
}

/// Binds a raw socket to some network interface
pub fn bind_raw(iface: SocketInterfaceQuery<'_>) -> Result<RawFd, Error> {
    let mut buffer = [0; 128];
    let fd = unsafe { crate::sys::create_socket(SocketType::RawPacket) }?;
    match set_socket_option_with::<options::BoundInterface>(fd, &mut buffer, &iface) {
        Ok(()) => Ok(fd),
        Err(error) => {
            unsafe { crate::sys::close(fd) }.ok();
            Err(error)
        }
    }
}

/// Binds an UDP socket to some local address
pub fn bind_udp(local: &SocketAddr) -> Result<RawFd, Error> {
    create_and_bind_ip(SocketType::UdpPacket, local, false)
}

/// Connect to a TCP listener
pub fn connect_tcp(remote: &SocketAddr, timeout: Option<Duration>) -> Result<RawFd, Error> {
    let fd = unsafe { crate::sys::create_socket(SocketType::TcpStream) }?;
    match connect_ip(fd, remote, timeout) {
        Ok(()) => Ok(fd),
        Err(error) => {
            unsafe { crate::sys::close(fd) }.ok();
            Err(error)
        }
    }
}

/// "Connect" an UDP socket
pub fn connect_udp(socket_fd: RawFd, remote: &SocketAddr) -> Result<(), Error> {
    connect_ip(socket_fd, remote, None)
}

/// Accept an IP connection
#[inline]
pub fn accept_ip(socket_fd: RawFd) -> Result<(RawFd, SocketAddr), Error> {
    let mut source = [0; 32];
    let fd = unsafe { crate::sys::accept(socket_fd, Some(&mut source)) }?;
    match wire::from_slice(&source) {
        Ok(address) => Ok((fd, address)),
        Err(error) => {
            unsafe { crate::sys::close(fd) }.ok();
            Err(error.into())
        }
    }
}

/// Receive from an IP socket (TCP/UDP)
#[inline]
pub fn receive_from_ip(socket_fd: RawFd, buffer: &mut [u8]) -> Result<(usize, SocketAddr), Error> {
    let mut source = [0; 32];
    let mut msg_header = MessageHeaderMut {
        source: Some(&mut source),
        ancillary: None,
        ancillary_len: 0,
        payload: buffer,
    };
    let len = unsafe { crate::sys::receive_message(socket_fd, &mut msg_header) }?;
    let source = wire::from_slice(&source)?;
    Ok((len, source))
}

/// Receive data from a socket
#[inline]
pub fn receive(socket_fd: RawFd, buffer: &mut [u8]) -> Result<usize, Error> {
    let mut msg_header = MessageHeaderMut {
        source: None,
        ancillary: None,
        ancillary_len: 0,
        payload: buffer,
    };
    unsafe { crate::sys::receive_message(socket_fd, &mut msg_header) }
}

/// Send data to an IP address
#[inline]
pub fn send_to_ip(socket_fd: RawFd, data: &[u8], destination: &SocketAddr) -> Result<usize, Error> {
    let mut target = [0; 32];
    let address_len = wire::to_slice(destination, &mut target)?;
    let msg_header = MessageHeader {
        destination: Some(&target[..address_len]),
        ancillary: &[],
        payload: data,
    };
    unsafe { crate::sys::send_message(socket_fd, &msg_header) }
}

/// Send data over a socket
#[inline]
pub fn send(socket_fd: RawFd, data: &[u8]) -> Result<usize, Error> {
    let msg_header = MessageHeader {
        destination: None,
        ancillary: &[],
        payload: data,
    };
    unsafe { crate::sys::send_message(socket_fd, &msg_header) }
}
