//! Interface for the `NOTIFY_SOCKET` readiness protocol
//!
//! The entry point of this module is the [`Notifier`] struct.
//!
//! # Example
//!
//! ```no_run
//! # main() -> Result<(), Box<dyn std::error::Error>> {
//! // Do whatever you need for your service to be ready.
//! if let Some(notifier) = Notifier::new(true)? {
//!     notifier.notify(&[NotifyState::Ready])?
//! }
//! # Ok(())
//! # }
//! ```

use core::fmt::Display;
use std::env;
use std::ffi::OsString;
use std::io::IoSlice;
use std::os::unix::net::UnixDatagram;

use rustix::fd::{AsFd, BorrowedFd};

use self::error::{FdNameError, NewNotifierError, NotifyError, SanityCheckError};

pub mod error;

/// A wrapper around the socket specified by `$NOTIFY_SOCKET`
#[derive(Debug)]
pub struct Notifier {
    sock_addr: rustix::net::SocketAddrUnix,
    socket: UnixDatagram,
}

impl Notifier {
    /// Create a new [`Notifier`].
    ///
    /// If `unset_env` is set to `true` this will unset `NOTIFY_SOCKET` resulting in further call to this function to return `Ok(None)`.
    ///
    /// Returns `Ok(None)` if notification isn't supported (`NOTIFY_SOCKET` isn't defined).
    ///
    /// # Errors
    ///
    /// This function error out if the socket couldn't be opened.
    pub fn new(unset_env: bool) -> Result<Option<Self>, NewNotifierError> {
        let env_sock = match env::var_os("NOTIFY_SOCKET") {
            None => return Ok(None),
            Some(v) => v,
        };

        if unset_env {
            env::remove_var("NOTIFY_SOCKET");
        }

        // If the first character of `$NOTIFY_SOCKET` is '@', the string
        // is understood as Linux abstract namespace socket.
        let socket_addr = match env_sock.as_encoded_bytes().strip_prefix(b"@").map(|v| {
            // SAFETY:
            // - Only strip ASCII '@' which is a non-empty UTF-8 substring
            unsafe { OsString::from_encoded_bytes_unchecked(v.to_vec()) }
        }) {
            Some(stripped_addr) => {
                rustix::net::SocketAddrUnix::new_abstract_name(stripped_addr.as_encoded_bytes())
                    .map_err(NewNotifierError::InvalidAbstractSocket)?
            }
            None => rustix::net::SocketAddrUnix::new(env_sock)
                .map_err(NewNotifierError::InvalidSocketPath)?,
        };

        let socket = UnixDatagram::unbound().map_err(NewNotifierError::CouldntOpenSocket)?;
        let ret = Self {
            sock_addr: socket_addr,
            socket,
        };
        Ok(Some(ret))
    }

    /// Notify service manager about status change and send file descriptors.
    ///
    /// Use this together with [`NotifyState::FdStore`]. Otherwise works like [`Notifier::notify()`].
    ///
    /// # Errors
    ///
    /// This function will error out if the passed [`NotifyState`] do not follow the rules set by systemd or if they couldn't be fully sent.
    pub fn notify_with_fds(
        &self,
        state: &[NotifyState<'_>],
        fds: &[BorrowedFd],
    ) -> Result<(), NotifyError> {
        state.iter().try_for_each(NotifyState::sanity_check)?;

        let msg = state
            .iter()
            .fold(String::new(), |acc, state| format!("{acc}{state}\n"))
            .into_bytes();
        let msg_len = msg.len();
        let msg_iov = IoSlice::new(&msg);

        let mut ancillary = if !fds.is_empty() {
            let mut ancillary = rustix::net::SendAncillaryBuffer::default();
            let tmp = rustix::net::SendAncillaryMessage::ScmRights(fds);
            if !ancillary.push(tmp) {
                return Err(NotifyError::PushAncillaryMessage);
            }
            ancillary
        } else {
            rustix::net::SendAncillaryBuffer::default()
        };

        let sent_len = rustix::net::sendmsg_unix(
            self.socket.as_fd(),
            &self.sock_addr,
            &[msg_iov],
            &mut ancillary,
            rustix::net::SendFlags::empty(),
        )
        .map_err(NotifyError::SendMsg)?;

        if sent_len != msg_len {
            return Err(NotifyError::PartialSend);
        }

        Ok(())
    }

    /// Notify service manager about status changes.
    ///
    /// Send a notification to the manager about service status changes. Also see [`Notifier::notify_with_fds()`] to send file descriptors.
    ///
    /// # Errors
    ///
    /// This function will error out if the passed [`NotifyState`] do not follow the rules set by systemd or if they couldn't be fully sent.
    pub fn notify(&self, state: &[NotifyState<'_>]) -> Result<(), NotifyError> {
        self.notify_with_fds(state, &[])
    }
}

/// Status changes, see sd_notify(3).
#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
pub enum NotifyState<'a> {
    /// D-Bus error-style error code.
    BusError(&'a str),
    /// errno-style error code.
    Errno(u8),
    /// A name for the submitted file descriptors.
    FdName(&'a str),
    /// Stores additional file descriptors in the service manager. Use [`Notifier::notify_with_fds()`] with this.
    FdStore,
    /// Remove stored file descriptors. Must be used together with [`NotifyState::FdName`].
    FdStoreRemove,
    /// Tell the service manager to not poll the filedescriptors for errors. This causes
    /// systemd to hold on to broken file descriptors which must be removed manually.
    /// Must be used together with [`NotifyState::FdStore`].
    FdpollDisable,
    /// The main process ID of the service, in case of forking applications.
    Mainpid(libc::pid_t),
    /// Custom state change, as a `KEY=VALUE` string.
    Other(&'a str),
    /// Service startup is finished.
    Ready,
    /// Service is reloading.
    Reloading,
    /// Custom status change.
    Status(&'a str),
    /// Service is beginning to shutdown.
    Stopping,
    /// Tell the service manager to update the watchdog timestamp.
    Watchdog,
    /// Tell the service manager to execute the configured watchdog option.
    WatchdogTrigger,
    /// Reset watchdog timeout value during runtime.
    /// The value is in milliseconds.
    WatchdogUsec(u64),
    /// Tells the service manager to extend the startup, runtime or shutdown service timeout corresponding the current state.
    /// The value is in milliseconds.
    ExtendTimeoutUsec(u64),
}

impl<'a> Display for NotifyState<'a> {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        match *self {
            NotifyState::BusError(s) => write!(f, "BUSERROR={s}"),
            NotifyState::Errno(e) => write!(f, "ERRNO={e}"),
            NotifyState::FdName(name) => write!(f, "FDNAME={name}"),
            NotifyState::FdStore => f.write_str("FDSTORE=1"),
            NotifyState::FdStoreRemove => f.write_str("FDSTOREREMOVE=1"),
            NotifyState::FdpollDisable => f.write_str("FDPOLL=0"),
            NotifyState::Mainpid(pid) => write!(f, "MAINPID={pid}"),
            NotifyState::Other(message) => Display::fmt(message, f),
            NotifyState::Ready => f.write_str("READY=1"),
            NotifyState::Reloading => f.write_str("READY=1"),
            NotifyState::Status(status) => write!(f, "STATUS={status}"),
            NotifyState::Stopping => f.write_str("STOPPING=1"),
            NotifyState::Watchdog => f.write_str("WATCHDOG=1"),
            NotifyState::WatchdogTrigger => f.write_str("WATCHDOG=trigger"),
            NotifyState::WatchdogUsec(milliseconds) => write!(f, "WATCHDOG_USEC={milliseconds}"),
            NotifyState::ExtendTimeoutUsec(milliseconds) => {
                write!(f, "EXTEND_TIMEOUT_USEC={milliseconds}")
            }
        }
    }
}

impl<'a> NotifyState<'a> {
    fn sanity_check(&self) -> Result<(), SanityCheckError> {
        match self {
            NotifyState::FdName(name) => {
                validate_fd_name(name).map_err(SanityCheckError::InvalidFdName)
            }
            _ => Ok(()),
        }
    }
}

fn validate_fd_name(name: &str) -> Result<(), FdNameError> {
    if name.len() > 255 {
        return Err(FdNameError::TooLong {
            length: name.len(),
            name: name.into(),
        });
    }

    for c in name.chars() {
        if !c.is_ascii() || c.is_ascii_control() {
            return Err(FdNameError::NotAsciiNonControl {
                disallowed_char: c,
                name: name.into(),
            });
        }
        if c == ':' {
            return Err(FdNameError::ContainColon(name.into()));
        }
    }
    Ok(())
}