//! Callback context.

use crate::error::*;
use crate::Stage;
use libc::c_char;
use milter_sys as sys;
use std::cell::{Ref, RefCell, RefMut};
use std::convert::TryInto;
use std::ffi::{CStr, CString};
use std::ptr::{self, NonNull};
use std::result;

/// Milter context supplied to milter callback functions.
///
/// The context also provides access to connection or message data.
///
/// The type parameter must always be the same for the callbacks in a milter.
///
/// TODO reconsider -- this is a problem because data handle fetches T eagerly,
/// must not be mistyped
pub struct Context<T> {
    base: ContextBase,

    /// A handle on managed user data.
    pub data: DataHandle<T>,
}

impl<T> Context<T> {
    /// Constructs a new context from the milter library-supplied raw context
    /// pointer.
    ///
    /// You do not normally need to use this method; a `Context` is already
    /// supplied to the callbacks.
    ///
    /// # Panics
    ///
    /// Panics if `ptr` is null.
    pub fn new(ptr: *mut sys::SMFICTX) -> Self {
        assert!(!ptr.is_null());

        Context {
            base: ContextBase::new(ptr),
            data: DataHandle::new(ptr),
        }
    }

    /// Returns the value for the given macro, if present.
    ///
    /// # Errors
    ///
    /// If conversion of arguments or return values at the boundary to the
    /// milter library fails, an error variant is returned.
    ///
    /// # Examples
    ///
    /// ```no_run
    /// # let context = milter::Context::<i32>::new(std::ptr::null_mut());
    /// let ip_address = context.macro_value("{daemon_addr}")?;
    /// # Ok::<(), milter::Error>(())
    /// ```
    pub fn macro_value(&self, name: &str) -> Result<Option<&str>> {
        self.base.macro_value(name)
    }

    /// Requests the given macros to be made available in the given stage.
    ///
    /// Macro requirements should be handled during negotiation.
    pub fn set_requested_macros(&self, stage: Stage, macros: &str) -> Result<()> {
        self.base.set_requested_macros(stage, macros)
    }

    /// Sets the default SMTP (and ESMTP) error reply code and message.
    pub fn set_error_reply(&self, code: &str, ext_code: Option<&str>, msg_lines: Vec<&str>) -> Result<()> {
        self.base.set_error_reply(code, ext_code, msg_lines)
    }

    /// Instructs the milter library event loop to exit, thereby shutting down
    /// this milter. This ultimately causes [`Milter::run`] to return.
    ///
    /// [`Milter::run`]: struct.Milter.html#method.run
    pub fn shut_down(&self) {
        self.base.shut_down()
    }
}

/// Context of the end of message, ie actions stage.
pub struct ActionContext<T = ()> {
    base: ContextBase,

    /// A handle on managed data in this context.
    pub data: DataHandle<T>,
}

impl<T> ActionContext<T> {
    /// Constructs a new action context from the milter library-supplied raw
    /// context pointer.
    ///
    /// You do not normally need to use this method; an `ActionContext` is
    /// already supplied to the `on_eom` callback.
    ///
    /// # Panics
    ///
    /// Panics if `ptr` is null.
    pub fn new(ptr: *mut sys::SMFICTX) -> Self {
        assert!(!ptr.is_null());

        ActionContext {
            base: ContextBase::new(ptr),
            data: DataHandle::new(ptr),
        }
    }

    /// Same as
    /// [`Context::macro_value`](struct.Context.html#method.macro_value).
    pub fn macro_value(&self, name: &str) -> Result<Option<&str>> {
        self.base.macro_value(name)
    }

    /// Same as
    /// [`Context::set_error_reply`](struct.Context.html#method.set_error_reply).
    pub fn set_error_reply(&self, code: &str, ext_code: Option<&str>, msg_lines: Vec<&str>) -> Result<()> {
        self.base.set_error_reply(code, ext_code, msg_lines)
    }

    /// Same as [`Context::shut_down`](struct.Context.html#method.shut_down).
    pub fn shut_down(&self) {
        self.base.shut_down()
    }

    /// Replaces the envelope sender (`MAIL FROM` address) of the current
    /// message.
    ///
    /// This action is enabled with the flag [`Actions::CHANGE_SENDER`].
    ///
    /// [`Actions::CHANGE_SENDER`]: struct.Actions.html#associatedconstant.CHANGE_SENDER
    pub fn change_sender(&self, mail_from: &str, esmtp_args: Option<&str>) -> Result<()> {
        self.base.change_sender(mail_from, esmtp_args)
    }

    /// Adds an envelope recipient (`RCPT TO` address) for the current message.
    ///
    /// This action is enabled with **two distinct flags**:
    ///
    /// * [`Actions::ADD_RECIPIENT`] if no ESMTP arguments are specified
    /// * [`Actions::ADD_RECIPIENT_EXT`] if additional ESMTP arguments are
    ///   specified
    ///
    /// [`Actions::ADD_RECIPIENT`]: struct.Actions.html#associatedconstant.ADD_RECIPIENT
    /// [`Actions::ADD_RECIPIENT_EXT`]: struct.Actions.html#associatedconstant.ADD_RECIPIENT_EXT
    pub fn add_recipient(&self, rcpt_to: &str, esmtp_args: Option<&str>) -> Result<()> {
        self.base.add_recipient(rcpt_to, esmtp_args)
    }

    /// Removes an envelope recipient (`RCPT TO` address) from the current
    /// message.
    ///
    /// This action is enabled with the flag [`Actions::DELETE_RECIPIENT`].
    ///
    /// [`Actions::DELETE_RECIPIENT`]: struct.Actions.html#associatedconstant.DELETE_RECIPIENT
    pub fn delete_recipient(&self, rcpt_to: &str) -> Result<()> {
        self.base.delete_recipient(rcpt_to)
    }

    /// Appends a header to the list of headers of the current message. If the
    /// header value is to span multiple lines, use `\n` (followed by
    /// whitespace) as the line separator, not `\r\n`.
    ///
    /// This action is enabled with the flag [`Actions::ADD_HEADER`].
    ///
    /// [`Actions::ADD_HEADER`]: struct.Actions.html#associatedconstant.ADD_HEADER
    pub fn add_header(&self, name: &str, value: &str) -> Result<()> {
        self.base.add_header(name, value)
    }

    /// Inserts a header at `index` in the list of headers of the current
    /// message.
    ///
    /// This action is enabled with the flag [`Actions::ADD_HEADER`].
    ///
    /// [`Actions::ADD_HEADER`]: struct.Actions.html#associatedconstant.ADD_HEADER
    pub fn insert_header(&self, index: usize, name: &str, value: &str) -> Result<()> {
        self.base.insert_header(index, name, value)
    }

    /// ‘Changes’, that is, replaces, removes, or appends a header at the
    /// `index`’th occurrence of headers with the given `name` for the current
    /// message. The index is 1-based (starts at 1).
    ///
    /// More precisely,
    ///
    /// * replaces the `index`’th occurrence of headers named `name` with the
    ///   specified value;
    /// * if the value is `None`, that header is instead removed;
    /// * if index is greater than the number of occurrences of headers named
    ///   `name`, a new header is instead appended.
    ///
    /// This action is enabled with the flag [`Actions::CHANGE_HEADER`].
    ///
    /// [`Actions::CHANGE_HEADER`]: struct.Actions.html#associatedconstant.CHANGE_HEADER
    pub fn change_header(&self, name: &str, index: usize, value: Option<&str>) -> Result<()> {
        self.base.change_header(name, index, value)
    }

    /// Appends content to the new message body of the current message.
    ///
    /// This method may be called repeatedly: initially empty, the new body is
    /// augmented with additional content with each call. If this method is not
    /// called, the original message body remains unchanged.
    ///
    /// This action is enabled with the flag [`Actions::CHANGE_BODY`].
    ///
    /// [`Actions::CHANGE_BODY`]: struct.Actions.html#associatedconstant.CHANGE_BODY
    pub fn append_new_body(&self, content: &[u8]) -> Result<()> {
        self.base.append_new_body(content)
    }

    /// Quarantines the current message for the given reason.
    ///
    /// This action is enabled with the flag [`Actions::QUARANTINE`].
    ///
    /// [`Actions::QUARANTINE`]: struct.Actions.html#associatedconstant.QUARANTINE
    pub fn quarantine(&self, reason: &str) -> Result<()> {
        self.base.quarantine(reason)
    }

    /// Signals to the milter library that this milter is still alive, causing
    /// it to reset timeouts.
    pub fn keepalive(&self) -> Result<()> {
        self.base.keepalive()
    }
}

/// Shared context base.
struct ContextBase {
    context_ptr: NonNull<sys::SMFICTX>,
}

impl ContextBase {
    fn new(ptr: *mut sys::SMFICTX) -> Self {
        ContextBase { context_ptr: NonNull::new(ptr).unwrap() }
    }

    fn macro_value(&self, name: &str) -> Result<Option<&str>> {
        let name = CString::new(name)?;

        unsafe {
            let value = sys::smfi_getsymval(self.context_ptr.as_ptr(), name.as_ptr() as _);

            Ok(if value.is_null() {
                None
            } else {
                Some(CStr::from_ptr(value).to_str()?)
            })
        }
    }

    fn set_requested_macros(&self, stage: Stage, macros: &str) -> Result<()> {
        let macros = CString::new(macros)?;

        let status_code = unsafe {
            sys::smfi_setsymlist(self.context_ptr.as_ptr(), stage as _, macros.as_ptr() as _)
        };

        StatusCode::from(status_code).into()
    }

    fn set_error_reply(&self, code: &str, ext_code: Option<&str>, msg_lines: Vec<&str>) -> Result<()> {
        let code = CString::new(code)?;
        let ext_code = ext_code.map(CString::new).transpose()?;

        let l: Vec<CString> = msg_lines.into_iter()
            .map(CString::new)
            .collect::<result::Result<_, _>>()?;

        let p = self.context_ptr.as_ptr();
        let c = code.as_ptr() as *mut _;
        let x = ext_code.as_ref().map_or(ptr::null_mut(), |s| s.as_ptr() as *mut _);

        // `smfi_setmlreply` is a varargs function. Fortunately it limits its
        // argument count to 32, so we can actually exhaustively match and
        // expand the arities below.

        macro_rules! set_reply {
            ($line:expr) => {
                sys::smfi_setreply(p, c, x, $line as _)
            };
            ($($line:expr),+) => {
                sys::smfi_setmlreply(
                    p, c, x, $($line.as_ptr() as *mut c_char),+, ptr::null_mut() as *mut c_char
                )
            }
        }

        let status_code = unsafe {
            match l.len() {
                0 => set_reply!(ptr::null_mut()),
                1 => set_reply!(l[0].as_ptr()),
                2 => set_reply!(l[0], l[1]),
                3 => set_reply!(l[0], l[1], l[2]),
                4 => set_reply!(l[0], l[1], l[2], l[3]),
                5 => set_reply!(l[0], l[1], l[2], l[3], l[4]),
                6 => set_reply!(l[0], l[1], l[2], l[3], l[4], l[5]),
                7 => set_reply!(l[0], l[1], l[2], l[3], l[4], l[5], l[6]),
                8 => set_reply!(l[0], l[1], l[2], l[3], l[4], l[5], l[6], l[7]),
                9 => set_reply!(l[0], l[1], l[2], l[3], l[4], l[5], l[6], l[7], l[8]),
                10 => set_reply!(l[0], l[1], l[2], l[3], l[4], l[5], l[6], l[7], l[8], l[9]),
                11 => set_reply!(l[0], l[1], l[2], l[3], l[4], l[5], l[6], l[7], l[8], l[9], l[10]),
                12 => set_reply!(l[0], l[1], l[2], l[3], l[4], l[5], l[6], l[7], l[8], l[9], l[10], l[11]),
                13 => set_reply!(l[0], l[1], l[2], l[3], l[4], l[5], l[6], l[7], l[8], l[9], l[10], l[11], l[12]),
                14 => set_reply!(l[0], l[1], l[2], l[3], l[4], l[5], l[6], l[7], l[8], l[9], l[10], l[11], l[12], l[13]),
                15 => set_reply!(l[0], l[1], l[2], l[3], l[4], l[5], l[6], l[7], l[8], l[9], l[10], l[11], l[12], l[13], l[14]),
                16 => set_reply!(l[0], l[1], l[2], l[3], l[4], l[5], l[6], l[7], l[8], l[9], l[10], l[11], l[12], l[13], l[14], l[15]),
                17 => set_reply!(l[0], l[1], l[2], l[3], l[4], l[5], l[6], l[7], l[8], l[9], l[10], l[11], l[12], l[13], l[14], l[15], l[16]),
                18 => set_reply!(l[0], l[1], l[2], l[3], l[4], l[5], l[6], l[7], l[8], l[9], l[10], l[11], l[12], l[13], l[14], l[15], l[16], l[17]),
                19 => set_reply!(l[0], l[1], l[2], l[3], l[4], l[5], l[6], l[7], l[8], l[9], l[10], l[11], l[12], l[13], l[14], l[15], l[16], l[17], l[18]),
                20 => set_reply!(l[0], l[1], l[2], l[3], l[4], l[5], l[6], l[7], l[8], l[9], l[10], l[11], l[12], l[13], l[14], l[15], l[16], l[17], l[18], l[19]),
                21 => set_reply!(l[0], l[1], l[2], l[3], l[4], l[5], l[6], l[7], l[8], l[9], l[10], l[11], l[12], l[13], l[14], l[15], l[16], l[17], l[18], l[19], l[20]),
                22 => set_reply!(l[0], l[1], l[2], l[3], l[4], l[5], l[6], l[7], l[8], l[9], l[10], l[11], l[12], l[13], l[14], l[15], l[16], l[17], l[18], l[19], l[20], l[21]),
                23 => set_reply!(l[0], l[1], l[2], l[3], l[4], l[5], l[6], l[7], l[8], l[9], l[10], l[11], l[12], l[13], l[14], l[15], l[16], l[17], l[18], l[19], l[20], l[21], l[22]),
                24 => set_reply!(l[0], l[1], l[2], l[3], l[4], l[5], l[6], l[7], l[8], l[9], l[10], l[11], l[12], l[13], l[14], l[15], l[16], l[17], l[18], l[19], l[20], l[21], l[22], l[23]),
                25 => set_reply!(l[0], l[1], l[2], l[3], l[4], l[5], l[6], l[7], l[8], l[9], l[10], l[11], l[12], l[13], l[14], l[15], l[16], l[17], l[18], l[19], l[20], l[21], l[22], l[23], l[24]),
                26 => set_reply!(l[0], l[1], l[2], l[3], l[4], l[5], l[6], l[7], l[8], l[9], l[10], l[11], l[12], l[13], l[14], l[15], l[16], l[17], l[18], l[19], l[20], l[21], l[22], l[23], l[24], l[25]),
                27 => set_reply!(l[0], l[1], l[2], l[3], l[4], l[5], l[6], l[7], l[8], l[9], l[10], l[11], l[12], l[13], l[14], l[15], l[16], l[17], l[18], l[19], l[20], l[21], l[22], l[23], l[24], l[25], l[26]),
                28 => set_reply!(l[0], l[1], l[2], l[3], l[4], l[5], l[6], l[7], l[8], l[9], l[10], l[11], l[12], l[13], l[14], l[15], l[16], l[17], l[18], l[19], l[20], l[21], l[22], l[23], l[24], l[25], l[26], l[27]),
                29 => set_reply!(l[0], l[1], l[2], l[3], l[4], l[5], l[6], l[7], l[8], l[9], l[10], l[11], l[12], l[13], l[14], l[15], l[16], l[17], l[18], l[19], l[20], l[21], l[22], l[23], l[24], l[25], l[26], l[27], l[28]),
                30 => set_reply!(l[0], l[1], l[2], l[3], l[4], l[5], l[6], l[7], l[8], l[9], l[10], l[11], l[12], l[13], l[14], l[15], l[16], l[17], l[18], l[19], l[20], l[21], l[22], l[23], l[24], l[25], l[26], l[27], l[28], l[29]),
                31 => set_reply!(l[0], l[1], l[2], l[3], l[4], l[5], l[6], l[7], l[8], l[9], l[10], l[11], l[12], l[13], l[14], l[15], l[16], l[17], l[18], l[19], l[20], l[21], l[22], l[23], l[24], l[25], l[26], l[27], l[28], l[29], l[30]),
                32 => set_reply!(l[0], l[1], l[2], l[3], l[4], l[5], l[6], l[7], l[8], l[9], l[10], l[11], l[12], l[13], l[14], l[15], l[16], l[17], l[18], l[19], l[20], l[21], l[22], l[23], l[24], l[25], l[26], l[27], l[28], l[29], l[30], l[31]),
                _ => panic!("more than 32 message lines"),
            }
        };

        StatusCode::from(status_code).into()
    }

    fn shut_down(&self) {
        let _ = unsafe { sys::smfi_stop() };
    }

    fn change_sender(&self, mail_from: &str, esmtp_args: Option<&str>) -> Result<()> {
        let mail_from = CString::new(mail_from)?;
        let esmtp_args = esmtp_args.map(CString::new).transpose()?;

        let status_code = unsafe {
            sys::smfi_chgfrom(
                self.context_ptr.as_ptr(),
                mail_from.as_ptr() as _,
                esmtp_args.as_ref().map_or(ptr::null_mut(), |s| s.as_ptr() as _),
            )
        };

        StatusCode::from(status_code).into()
    }

    fn add_recipient(&self, rcpt_to: &str, esmtp_args: Option<&str>) -> Result<()> {
        let rcpt_to = CString::new(rcpt_to)?;

        let status_code = unsafe {
            match esmtp_args {
                None => sys::smfi_addrcpt(self.context_ptr.as_ptr(), rcpt_to.as_ptr() as _),
                Some(esmtp_args) => {
                    let esmtp_args = CString::new(esmtp_args)?;

                    sys::smfi_addrcpt_par(
                        self.context_ptr.as_ptr(),
                        rcpt_to.as_ptr() as _,
                        esmtp_args.as_ptr() as _,
                    )
                }
            }
        };

        StatusCode::from(status_code).into()
    }

    fn delete_recipient(&self, rcpt_to: &str) -> Result<()> {
        let rcpt_to = CString::new(rcpt_to)?;

        let status_code = unsafe {
            sys::smfi_delrcpt(self.context_ptr.as_ptr(), rcpt_to.as_ptr() as _)
        };

        StatusCode::from(status_code).into()
    }

    fn add_header(&self, name: &str, value: &str) -> Result<()> {
        let name = CString::new(name)?;
        let value = CString::new(value)?;

        let status_code = unsafe {
            sys::smfi_addheader(self.context_ptr.as_ptr(), name.as_ptr() as _, value.as_ptr() as _)
        };

        StatusCode::from(status_code).into()
    }

    fn insert_header(&self, index: usize, name: &str, value: &str) -> Result<()> {
        let index = index.try_into()?;
        let name = CString::new(name)?;
        let value = CString::new(value)?;

        let status_code = unsafe {
            sys::smfi_insheader(
                self.context_ptr.as_ptr(),
                index,
                name.as_ptr() as _,
                value.as_ptr() as _,
            )
        };

        StatusCode::from(status_code).into()
    }

    fn change_header(&self, name: &str, index: usize, value: Option<&str>) -> Result<()> {
        let name = CString::new(name)?;
        let index = index.try_into()?;
        let value = value.map(CString::new).transpose()?;

        let status_code = unsafe {
            sys::smfi_chgheader(
                self.context_ptr.as_ptr(),
                name.as_ptr() as _,
                index,
                value.as_ref().map_or(ptr::null_mut(), |s| s.as_ptr() as _),
            )
        };

        StatusCode::from(status_code).into()
    }

    fn append_new_body(&self, content: &[u8]) -> Result<()> {
        let status_code = unsafe {
            sys::smfi_replacebody(
                self.context_ptr.as_ptr(),
                content.as_ptr() as _,
                content.len() as _,
            )
        };

        StatusCode::from(status_code).into()
    }

    fn quarantine(&self, reason: &str) -> Result<()> {
        let reason = CString::new(reason)?;

        let status_code = unsafe {
            sys::smfi_quarantine(self.context_ptr.as_ptr(), reason.as_ptr() as _)
        };

        StatusCode::from(status_code).into()
    }

    fn keepalive(&self) -> Result<()> {
        let status_code = unsafe {
            sys::smfi_progress(self.context_ptr.as_ptr())
        };

        StatusCode::from(status_code).into()
    }
}

/// A handle on user data managed in the callback context.
///
/// This serves as an accessor to ‘private data’, that is connection-specific
/// and message-specific user data. Care must be taken to manage the data’s
/// lifecycle; specifically, the managed data is not automatically cleaned up
/// but must be re-materialised (and dropped) at some appropriate time like
/// connection close or abort.
///
/// # Data lifecycle
///
/// Every path in the callbacks must have a final call to `DataHandle::take`.
pub struct DataHandle<T> {
    context_ptr: NonNull<sys::SMFICTX>,
    data_ptr: RefCell<Option<NonNull<T>>>,
}

impl<T> DataHandle<T> {
    fn new(ptr: *mut sys::SMFICTX) -> Self {
        assert!(!ptr.is_null());

        let data_ptr = unsafe { sys::smfi_getpriv(ptr) as _ };

        DataHandle {
            context_ptr: NonNull::new(ptr).unwrap(),
            data_ptr: RefCell::new(NonNull::new(data_ptr)),
        }
    }

    /// Hands over data into the context, returning current managed data if
    /// present.
    pub fn replace(&self, data: T) -> Result<Option<T>> {
        unsafe { self.replace_data(Box::into_raw(Box::new(data))) }
    }

    /// Retakes ownership of the managed data if present and removes it from the
    /// context.
    pub fn take(&self) -> Result<Option<T>> {
        unsafe { self.replace_data(ptr::null_mut()) }
    }

    unsafe fn replace_data(&self, ptr: *mut T) -> Result<Option<T>> {
        let status_code = sys::smfi_setpriv(self.context_ptr.as_ptr(), ptr as _).into();

        match status_code {
            StatusCode::Success => {
                Ok(self
                    .data_ptr
                    .replace(NonNull::new(ptr))
                    .map(|x| *Box::from_raw(x.as_ptr())))
            }
            StatusCode::Failure => {
                // On failure, resurrect the data and drop.
                if !ptr.is_null() {
                    Box::from_raw(ptr);
                }
                Err(Error::from(ErrorKind::MilterFailureStatus))
            }
        }
    }

    /// Returns a reference to the managed data if present.
    ///
    /// It is the responsibility of the user only to call this when appropriate:
    /// eg, when the data is already borrowed mutably, `Err` is returned.
    ///
    /// # Errors
    ///
    /// An error is returned when an immutable reference cannot be borrowed.
    ///
    /// # Examples
    ///
    /// ```no_run
    /// # let context = milter::Context::<i32>::new(std::ptr::null_mut());
    /// if let Some(data) = context.data.borrow()? {
    ///     println!("{}", data);
    /// }
    /// # Ok::<(), milter::Error>(())
    /// ```
    pub fn borrow(&self) -> Result<Option<Ref<T>>> {
        let data_ref = self.data_ptr.try_borrow()
            .map_err(|e| Error::new(ErrorKind::DataAccessError, e))?;

        Ok(if data_ref.is_none() {
            None
        } else {
            Some(Ref::map(data_ref, |x| unsafe { x.as_ref().unwrap().as_ref() }))
        })
    }

    /// Returns a mutable reference to the managed data if present.
    ///
    /// It is the responsibility of the user only to call this when appropriate:
    /// only one exclusive mutable borrow is allowed at any time.
    ///
    /// # Examples
    ///
    /// ```no_run
    /// # let context = milter::Context::<usize>::new(std::ptr::null_mut());
    /// if let Some(mut msg_count) = context.data.borrow_mut()? {
    ///     *msg_count += 1;
    /// }
    /// # Ok::<(), milter::Error>(())
    /// ```
    pub fn borrow_mut(&self) -> Result<Option<RefMut<T>>> {
        let data_ref = self.data_ptr.try_borrow_mut()
            .map_err(|e| Error::new(ErrorKind::DataAccessError, e))?;

        Ok(if data_ref.is_none() {
            None
        } else {
            Some(RefMut::map(data_ref, |x| unsafe { x.as_mut().unwrap().as_mut() }))
        })
    }
}