ntr-io 0.0.1

use super as io;
use super::raw::{AsRawSocket, FromRawSocket, IntoRawSocket, RawHandle, RawSocket, ValidRawSocket};
use super::win::*;
use core::fmt;
use core::marker::PhantomData;
use core::mem::{self, ManuallyDrop};

#[derive(Copy, Clone)]
#[repr(transparent)]
pub struct BorrowedSocket<'socket> {
    socket: ValidRawSocket,
    _phantom: PhantomData<&'socket OwnedSocket>,
}

#[repr(transparent)]
pub struct OwnedSocket {
    socket: ValidRawSocket,
}

impl BorrowedSocket<'_> {
    #[inline]
    #[track_caller]
    pub const unsafe fn borrow_raw(socket: RawSocket) -> Self {
        Self {
            socket: ValidRawSocket::new(socket).expect("socket != -1"),
            _phantom: PhantomData,
        }
    }
}

impl OwnedSocket {
    pub fn try_clone(&self) -> io::Result<Self> {
        self.as_socket().try_clone_to_owned()
    }

    #[cfg(not(target_vendor = "uwp"))]
    pub(crate) fn set_no_inherit(&self) -> io::Result<()> {
        unsafe {
            SetHandleInformation(
                (self.as_raw_socket() as RawHandle).as_windows_handle(),
                HANDLE_FLAG_INHERIT.0,
                HANDLE_FLAGS::default(),
            )?;
        }
        Ok(())
    }

    #[cfg(target_vendor = "uwp")]
    pub(crate) fn set_no_inherit(&self) -> io::Result<()> {
        Err(io::Error::Windows(io::WindowsError::new(
            Default::default(),
            "unvailable on UWP",
        )))
    }
}

impl BorrowedSocket<'_> {
    pub fn try_clone_to_owned(&self) -> io::Result<OwnedSocket> {
        let mut info = unsafe { mem::zeroed::<WSAPROTOCOL_INFOW>() };
        let result = unsafe {
            WSADuplicateSocketW(
                self.as_raw_socket().as_windows_socket(),
                GetCurrentProcessId(),
                &mut info,
            )
        };
        net_cvt(result)?;
        let socket = unsafe {
            WSASocketW(
                info.iAddressFamily,
                info.iSocketType,
                info.iProtocol,
                (&info as *const WSAPROTOCOL_INFOW).into(),
                0,
                WSA_FLAG_OVERLAPPED | WSA_FLAG_NO_HANDLE_INHERIT,
            )?
        };

        if socket != INVALID_SOCKET {
            unsafe { Ok(OwnedSocket::from_raw_socket(socket.0 as _)) }
        } else {
            let error = unsafe { WSAGetLastError() };

            if error != WSAEPROTOTYPE && error != WSAEINVAL {
                return Err(io::Error::from_os_code(error.0));
            }

            let socket = unsafe {
                WSASocketW(
                    info.iAddressFamily,
                    info.iSocketType,
                    info.iProtocol,
                    (&info as *const WSAPROTOCOL_INFOW).into(),
                    0,
                    WSA_FLAG_OVERLAPPED | WSA_FLAG_NO_HANDLE_INHERIT,
                )
            }?;

            if socket == INVALID_SOCKET {
                return Err(last_error());
            }

            unsafe {
                let socket = OwnedSocket::from_raw_socket(socket.as_raw_socket());
                socket.set_no_inherit()?;
                Ok(socket)
            }
        }
    }
}

fn last_error() -> io::Error {
    io::Error::from_os_code(unsafe { WSAGetLastError() }.0)
}

impl AsRawSocket for BorrowedSocket<'_> {
    #[inline]
    fn as_raw_socket(&self) -> RawSocket {
        self.socket.as_inner()
    }
}

impl AsRawSocket for OwnedSocket {
    #[inline]
    fn as_raw_socket(&self) -> RawSocket {
        self.socket.as_inner()
    }
}

impl IntoRawSocket for OwnedSocket {
    #[inline]
    fn into_raw_socket(self) -> RawSocket {
        ManuallyDrop::new(self).socket.as_inner()
    }
}

impl FromRawSocket for OwnedSocket {
    #[inline]
    #[track_caller]
    unsafe fn from_raw_socket(socket: RawSocket) -> Self {
        Self {
            socket: ValidRawSocket::new(socket).expect("socket != -1"),
        }
    }
}

impl Drop for OwnedSocket {
    #[inline]
    fn drop(&mut self) {
        unsafe {
            let _ = closesocket(self.socket.as_inner().as_windows_socket());
        }
    }
}

impl fmt::Debug for BorrowedSocket<'_> {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        f.debug_struct("BorrowedSocket")
            .field("socket", &self.socket)
            .finish()
    }
}

impl fmt::Debug for OwnedSocket {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        f.debug_struct("OwnedSocket")
            .field("socket", &self.socket)
            .finish()
    }
}

pub trait AsSocket {
    fn as_socket(&self) -> BorrowedSocket<'_>;
}

impl<T: AsSocket> AsSocket for &T {
    #[inline]
    fn as_socket(&self) -> BorrowedSocket<'_> {
        T::as_socket(self)
    }
}

impl<T: AsSocket> AsSocket for &mut T {
    #[inline]
    fn as_socket(&self) -> BorrowedSocket<'_> {
        T::as_socket(self)
    }
}

#[cfg(feature = "alloc")]
impl<T: AsSocket> AsSocket for alloc::sync::Arc<T> {
    #[inline]
    fn as_socket(&self) -> BorrowedSocket<'_> {
        (**self).as_socket()
    }
}

#[cfg(feature = "alloc")]
impl<T: AsSocket> AsSocket for alloc::rc::Rc<T> {
    #[inline]
    fn as_socket(&self) -> BorrowedSocket<'_> {
        (**self).as_socket()
    }
}

#[cfg(feature = "nightly")]
#[cfg(feature = "alloc")]
impl<T: AsSocket + ?Sized> AsSocket for alloc::rc::UniqueRc<T> {
    #[inline]
    fn as_socket(&self) -> BorrowedSocket<'_> {
        (**self).as_socket()
    }
}

impl<T: AsSocket> AsSocket for Box<T> {
    #[inline]
    fn as_socket(&self) -> BorrowedSocket<'_> {
        (**self).as_socket()
    }
}

impl AsSocket for BorrowedSocket<'_> {
    #[inline]
    fn as_socket(&self) -> BorrowedSocket<'_> {
        *self
    }
}

impl AsSocket for OwnedSocket {
    #[inline]
    fn as_socket(&self) -> BorrowedSocket<'_> {
        unsafe { BorrowedSocket::borrow_raw(self.as_raw_socket()) }
    }
}