use std::io::{self, IoSlice, IoSliceMut};
use std::mem::MaybeUninit;
use std::os::fd::{AsFd, OwnedFd};

use lazaret_core::codec::{decode, encode};
use rustix::net::{
    AddressFamily, RecvAncillaryBuffer, RecvAncillaryMessage, RecvFlags, ReturnFlags,
    SendAncillaryBuffer, SendAncillaryMessage, SendFlags, SocketFlags, SocketType, recv, recvmsg,
    send, sendmsg, socketpair,
};

const MAX_PRIVATE_PACKET: usize = 64 * 1024;
const MAX_TRANSFERRED_DESCRIPTORS: usize = 8;

/// # Errors
/// Returns an error when the private sequenced-packet pair cannot be created.
pub fn private_seqpacket_pair() -> std::io::Result<(OwnedFd, OwnedFd)> {
    socketpair(
        AddressFamily::UNIX,
        SocketType::SEQPACKET,
        SocketFlags::CLOEXEC,
        None,
    )
    .map_err(Into::into)
}

/// # Errors
/// Returns an error when the private stream pair cannot be created.
pub fn private_stream_pair() -> std::io::Result<(OwnedFd, OwnedFd)> {
    socketpair(
        AddressFamily::UNIX,
        SocketType::STREAM,
        SocketFlags::CLOEXEC,
        None,
    )
    .map_err(Into::into)
}

/// # Errors
/// Returns an error when a typed private packet is oversized or cannot be sent atomically.
pub fn send_packet<T: serde::Serialize>(descriptor: impl AsFd, message: &T) -> io::Result<()> {
    let encoded = encode(message).map_err(|error| io::Error::other(error.to_string()))?;
    if encoded.is_empty() || encoded.len() > MAX_PRIVATE_PACKET {
        return Err(io::Error::other("private packet has invalid size"));
    }
    let sent = send(descriptor.as_fd(), &encoded, SendFlags::empty())?;
    if sent == encoded.len() {
        Ok(())
    } else {
        Err(io::Error::other("private packet send was incomplete"))
    }
}

/// # Errors
/// Returns an error when a private packet is empty, oversized, malformed, or unknown-version.
pub fn receive_packet<T>(descriptor: impl AsFd) -> io::Result<T>
where
    T: serde::de::DeserializeOwned + serde::Serialize,
{
    let mut encoded = vec![0_u8; MAX_PRIVATE_PACKET + 1];
    let (_, received) = recv(descriptor.as_fd(), &mut encoded, RecvFlags::empty())?;
    if received == 0 || received > MAX_PRIVATE_PACKET {
        return Err(io::Error::other("private packet has invalid size"));
    }
    decode(&encoded[..received]).map_err(|error| io::Error::other(error.to_string()))
}

/// Sends one nonempty bounded descriptor set with a one-byte packet.
///
/// # Errors
/// Refuses empty, oversized, incomplete, or failed descriptor transfer.
pub fn send_descriptors(
    socket: impl AsFd,
    descriptors: &[std::os::fd::BorrowedFd<'_>],
) -> io::Result<()> {
    if descriptors.is_empty() || descriptors.len() > MAX_TRANSFERRED_DESCRIPTORS {
        return Err(io::Error::other("private descriptor set has invalid size"));
    }
    let mut space =
        [MaybeUninit::uninit(); rustix::cmsg_space!(ScmRights(MAX_TRANSFERRED_DESCRIPTORS))];
    let mut control = SendAncillaryBuffer::new(&mut space);
    if !control.push(SendAncillaryMessage::ScmRights(descriptors)) {
        return Err(io::Error::other(
            "private descriptor control buffer is too small",
        ));
    }
    let bytes = [1_u8];
    let sent = sendmsg(
        socket,
        &[IoSlice::new(&bytes)],
        &mut control,
        SendFlags::empty(),
    )?;
    if sent == bytes.len() {
        Ok(())
    } else {
        Err(io::Error::other("private descriptor packet was incomplete"))
    }
}

/// Receives exactly `expected` descriptors with close-on-exec set atomically.
///
/// # Errors
/// Refuses missing, extra, truncated, malformed, or failed descriptor transfer.
pub fn receive_descriptors(socket: impl AsFd, expected: usize) -> io::Result<Vec<OwnedFd>> {
    if expected == 0 || expected > MAX_TRANSFERRED_DESCRIPTORS {
        return Err(io::Error::other("expected descriptor count is invalid"));
    }
    let mut byte = [0_u8];
    let mut slices = [IoSliceMut::new(&mut byte)];
    let mut space =
        [MaybeUninit::uninit(); rustix::cmsg_space!(ScmRights(MAX_TRANSFERRED_DESCRIPTORS))];
    let mut control = RecvAncillaryBuffer::new(&mut space);
    let received = recvmsg(socket, &mut slices, &mut control, RecvFlags::CMSG_CLOEXEC)?;
    if received.bytes != 1
        || byte != [1]
        || received
            .flags
            .intersects(ReturnFlags::TRUNC | ReturnFlags::CTRUNC)
    {
        return Err(io::Error::other("private descriptor packet is malformed"));
    }
    let mut descriptors = Vec::new();
    for message in control.drain() {
        match message {
            RecvAncillaryMessage::ScmRights(rights) => descriptors.extend(rights),
            _ => return Err(io::Error::other("private ancillary message is unexpected")),
        }
    }
    if descriptors.len() == expected {
        Ok(descriptors)
    } else {
        Err(io::Error::other(
            "private descriptor count mismatches request",
        ))
    }
}

#[cfg(test)]
mod tests {
    use std::fs::File;
    use std::os::fd::{AsFd, AsRawFd};

    use rustix::net::{RecvFlags, SendFlags, recv, send};

    use super::{private_seqpacket_pair, receive_descriptors, send_descriptors};

    #[test]
    fn preserves_message_boundaries() {
        let (left, right) = private_seqpacket_pair().unwrap();
        send(left.as_fd(), b"first", SendFlags::empty()).unwrap();
        send(left.as_fd(), b"second", SendFlags::empty()).unwrap();
        let mut buffer = [0_u8; 16];
        let (_, first) = recv(right.as_fd(), &mut buffer, RecvFlags::empty()).unwrap();
        assert_eq!(&buffer[..first], b"first");
        let (_, second) = recv(right.as_fd(), &mut buffer, RecvFlags::empty()).unwrap();
        assert_eq!(&buffer[..second], b"second");
    }

    #[test]
    fn transfers_exact_close_on_exec_descriptors() {
        let (left, right) = private_seqpacket_pair().unwrap();
        let first = File::open("/dev/null").unwrap();
        let second = File::open("/dev/zero").unwrap();
        send_descriptors(left, &[first.as_fd(), second.as_fd()]).unwrap();
        let received = receive_descriptors(right, 2).unwrap();
        assert_ne!(received[0].as_raw_fd(), first.as_raw_fd());
        assert_ne!(received[1].as_raw_fd(), second.as_raw_fd());
        for descriptor in received {
            assert!(
                rustix::io::fcntl_getfd(&descriptor)
                    .unwrap()
                    .contains(rustix::io::FdFlags::CLOEXEC)
            );
        }
    }
}