mod common;

pub use crate::common::*;
use byte_strings::c_str;
use bytes::Bytes;
use indymilter::{
    message::{
        command::{
            Command, ConnInfoPayload, EnvAddrPayload, HeaderPayload, HeloPayload, OptNegPayload,
        },
        reply::Reply,
        PROTOCOL_VERSION,
    },
    Actions, Callbacks, ProtoOpts, SocketInfo, Status,
};
use std::{
    io::ErrorKind,
    net::Ipv4Addr,
    sync::{Arc, Mutex},
    time::Duration,
};
use tokio::time;

#[tokio::test]
async fn ignore_unexpected_command() {
    init_tracing_subscriber();

    let milter = Milter::spawn(LOCALHOST, Callbacks::<()>::new(), default_config())
        .await
        .unwrap();

    let mut client = Client::connect(milter.addr()).await.unwrap();

    client
        .write_command(Command::Mail(EnvAddrPayload {
            args: vec![c_str!("me@example.com").into()],
        }))
        .await
        .unwrap();

    // A Mail command is actually valid at this stage, because the state logic
    // will retry a failed transition from the helo stage!
    let reply = client.read_reply().await.unwrap();
    assert_eq!(reply, Reply::Continue);

    client.write_command(Command::Eoh).await.unwrap();

    // An Eoh command, too, is valid, because all callbacks between mail and eoh
    // are unavailable and can therefore be skipped.
    let reply = client.read_reply().await.unwrap();
    assert_eq!(reply, Reply::Continue);

    // A ConnInfo command, however, can definitely not be reached and will be
    // ignored.

    client
        .write_command(Command::ConnInfo(ConnInfoPayload {
            hostname: c_str!("mail.example.com").into(),
            socket_info: SocketInfo::Unknown,
        }))
        .await
        .unwrap();

    time::pause();

    let e = time::timeout(Duration::from_secs(5), client.read_reply()).await;

    time::resume();

    assert!(e.is_err());

    client.write_command(Command::Quit).await.unwrap();

    client.disconnect().await.unwrap();

    milter.shutdown().await.unwrap();
}

#[tokio::test]
async fn unknown_command() {
    init_tracing_subscriber();

    let milter = Milter::spawn(LOCALHOST, Callbacks::<()>::new(), default_config())
        .await
        .unwrap();

    let mut client = Client::connect(milter.addr()).await.unwrap();

    client.write_bytes(b"\0\0\0\x02\0\0").await.unwrap();

    let e = client.read_reply().await.unwrap_err();
    assert_eq!(e.kind(), ErrorKind::UnexpectedEof);

    client.disconnect().await.unwrap();

    milter.shutdown().await.unwrap();
}

#[tokio::test]
async fn jumbo_size_command() {
    init_tracing_subscriber();

    let milter = Milter::spawn(LOCALHOST, Callbacks::<()>::new(), default_config())
        .await
        .unwrap();

    let mut client = Client::connect(milter.addr()).await.unwrap();

    let buflen = 10_000_000;
    let mut buf = vec![b'x'; buflen];

    let msglen = (u32::try_from(buflen).unwrap() - 4).to_be_bytes();
    buf[..4].copy_from_slice(&msglen);

    let e = client.write_bytes(&buf).await.unwrap_err();
    assert_eq!(e.kind(), ErrorKind::BrokenPipe);

    assert!(client.disconnect().await.is_err());

    milter.shutdown().await.unwrap();
}

#[tokio::test]
async fn all_stages() {
    init_tracing_subscriber();

    let result = Arc::new(Mutex::new(Vec::new()));

    let callbacks = {
        fn push_stage(data: &mut Option<Vec<&'static str>>, stage: &'static str) -> Status {
            data.as_mut().unwrap().push(stage);
            Status::Continue
        }

        let result = result.clone();

        Callbacks::new()
            .on_negotiate(|cx, _, _| {
                Box::pin(async move {
                    cx.data = Some(vec!["negotiate"]);

                    Status::Continue
                })
            })
            .on_connect(|cx, _, _| Box::pin(async move { push_stage(&mut cx.data, "connect") }))
            .on_helo(|cx, _| Box::pin(async move { push_stage(&mut cx.data, "helo") }))
            .on_mail(|cx, _| Box::pin(async move { push_stage(&mut cx.data, "mail") }))
            .on_rcpt(|cx, _| Box::pin(async move { push_stage(&mut cx.data, "rcpt") }))
            .on_data(|cx| Box::pin(async move { push_stage(&mut cx.data, "data") }))
            .on_header(|cx, _, _| Box::pin(async move { push_stage(&mut cx.data, "header") }))
            .on_eoh(|cx| Box::pin(async move { push_stage(&mut cx.data, "eoh") }))
            .on_body(|cx, _| Box::pin(async move { push_stage(&mut cx.data, "body") }))
            .on_eom(|cx| Box::pin(async move { push_stage(&mut cx.data, "eom") }))
            .on_close(move |cx| {
                let result = result.clone();

                Box::pin(async move {
                    if let Some(mut data) = cx.data.take() {
                        data.push("close");
                        *result.lock().unwrap() = data;
                    }

                    Status::Continue
                })
            })
    };

    let milter = Milter::spawn(LOCALHOST, callbacks, default_config())
        .await
        .unwrap();

    let mut client = Client::connect(milter.addr()).await.unwrap();

    client
        .write_command(Command::OptNeg(OptNegPayload {
            version: PROTOCOL_VERSION,
            actions: Actions::all(),
            opts: ProtoOpts::all(),
        }))
        .await
        .unwrap();

    let reply = client.read_reply().await.unwrap();
    assert!(matches!(reply, Reply::OptNeg { .. }));

    client
        .write_command(Command::ConnInfo(ConnInfoPayload {
            hostname: c_str!("mail.example.com").into(),
            socket_info: SocketInfo::Inet((Ipv4Addr::LOCALHOST, 1234).into()),
        }))
        .await
        .unwrap();

    let reply = client.read_reply().await.unwrap();
    assert_eq!(reply, Reply::Continue);

    client
        .write_command(Command::Helo(HeloPayload {
            hostname: c_str!("mail.example.com").into(),
        }))
        .await
        .unwrap();

    let reply = client.read_reply().await.unwrap();
    assert_eq!(reply, Reply::Continue);

    client
        .write_command(Command::Mail(EnvAddrPayload {
            args: vec![c_str!("sender@example.com").into()],
        }))
        .await
        .unwrap();

    let reply = client.read_reply().await.unwrap();
    assert_eq!(reply, Reply::Continue);

    client
        .write_command(Command::Rcpt(EnvAddrPayload {
            args: vec![c_str!("recipient@example.com").into()],
        }))
        .await
        .unwrap();

    let reply = client.read_reply().await.unwrap();
    assert_eq!(reply, Reply::Continue);

    client.write_command(Command::Data).await.unwrap();

    let reply = client.read_reply().await.unwrap();
    assert_eq!(reply, Reply::Continue);

    client
        .write_command(Command::Header(HeaderPayload {
            name: c_str!("name").into(),
            value: c_str!("value").into(),
        }))
        .await
        .unwrap();

    let reply = client.read_reply().await.unwrap();
    assert_eq!(reply, Reply::Continue);

    client.write_command(Command::Eoh).await.unwrap();

    let reply = client.read_reply().await.unwrap();
    assert_eq!(reply, Reply::Continue);

    client
        .write_command(Command::BodyChunk(Bytes::from_static(b"body")))
        .await
        .unwrap();

    let reply = client.read_reply().await.unwrap();
    assert_eq!(reply, Reply::Continue);

    client
        .write_command(Command::BodyEnd(Bytes::new()))
        .await
        .unwrap();

    let reply = client.read_reply().await.unwrap();
    assert_eq!(reply, Reply::Continue);

    client.write_command(Command::Quit).await.unwrap();

    client.disconnect().await.unwrap();

    milter.shutdown().await.unwrap();

    let result = Arc::try_unwrap(result).unwrap().into_inner().unwrap();

    assert_eq!(
        result,
        [
            "negotiate",
            "connect",
            "helo",
            "mail",
            "rcpt",
            "data",
            "header",
            "eoh",
            "body",
            "eom",
            "close"
        ]
    );
}

#[tokio::test]
async fn session_terminated_on_shutdown() {
    init_tracing_subscriber();

    let callbacks = Callbacks::<()>::new()
        .on_mail(|_, _| {
            Box::pin(async {
                time::sleep(Duration::from_secs(1)).await;

                Status::Continue
            })
        })
        .on_rcpt(|_, _| {
            Box::pin(async {
                panic!();
            })
        });

    let milter = Milter::spawn(LOCALHOST, callbacks, default_config())
        .await
        .unwrap();

    let mut client = Client::connect(milter.addr()).await.unwrap();

    time::pause();

    client
        .write_command(Command::Mail(EnvAddrPayload {
            args: vec![c_str!("me@example.com").into()],
        }))
        .await
        .unwrap();

    time::sleep(Duration::from_millis(100)).await;

    // Milter shutdown waits for all sessions to exit at the beginning of the
    // main processing loop. The milter has still processed `Mail` regularly and
    // has responded with `Continue`.

    milter.shutdown().await.unwrap();

    time::resume();

    let reply = client.read_reply().await.unwrap();
    assert_eq!(reply, Reply::Continue);

    // The connection is gone. Note how due to the stream’s buffering, we can
    // perform one more write (experimentally, up to 512 kB) and flush:
    client
        .write_command(Command::Rcpt(EnvAddrPayload {
            args: vec![c_str!("you@example.com").into()],
        }))
        .await
        .unwrap();

    assert!(client.disconnect().await.is_err());
}
