use axum::{
    body::Bytes,
    extract::{DefaultBodyLimit, State},
    http::{HeaderMap, StatusCode},
    routing::post,
    Router,
};
use subtle::ConstantTimeEq;
use tokio::io::{AsyncBufReadExt, AsyncWriteExt, BufReader};
use tokio::net::UnixStream;

struct Config {
    token: &'static str,
    lmtp: &'static str,
    rcpt: Option<&'static str>,
    bind: &'static str,
}

#[tokio::main]
async fn main() -> anyhow::Result<()> {
    tracing_subscriber::fmt::init();

    let cfg: &'static Config = Box::leak(Box::new(Config {
        token: std::env::var("INGEST_TOKEN")?.leak(),
        lmtp: std::env::var("LMTP_SOCK")?.leak(),
        bind: std::env::var("LISTEN_ADDR").unwrap_or_else(|_| "127.0.0.1:8080".into()).leak(),
        rcpt: match std::env::var("LMTP_RCPT") {
            Ok(s) => Some(s.leak()),
            Err(_) => None,
        },
    }));

    let app = Router::new()
        .route("/inject", post(inject))
        .layer(DefaultBodyLimit::max(26 * 1024 * 1024))
        .with_state(cfg);

    let listener = tokio::net::TcpListener::bind(&cfg.bind).await?;
    tracing::info!("listening on {}", listener.local_addr()?);
    axum::serve(listener, app).await?;
    Ok(())
}

async fn inject(
    State(cfg): State<&'static Config>,
    headers: HeaderMap,
    body: Bytes,
) -> StatusCode {
    if !valid_token(&headers, &cfg.token) {
        return StatusCode::UNAUTHORIZED;
    }
    let from = header(&headers, "x-envelope-from");
    let to = header(&headers, "x-envelope-to");

    match lmtp_deliver(cfg, from, to, &body).await {
        Ok(_) => StatusCode::NO_CONTENT,
        Err((true, m)) => {
            tracing::warn!("permanent: {m}");
            StatusCode::UNPROCESSABLE_ENTITY
        }
        Err((false, m)) => {
            tracing::error!("temporary: {m}");
            StatusCode::SERVICE_UNAVAILABLE
        }
    }
}

fn header<'a>(h: &'a HeaderMap, k: &str) -> &'a str {
    h.get(k).and_then(|v| v.to_str().ok()).unwrap_or("")
}

fn valid_token(h: &HeaderMap, expected: &str) -> bool {
    let Some(v) = h.get("authorization").and_then(|v| v.to_str().ok()) else {
        return false;
    };
    let Some(got) = v.strip_prefix("Bearer ") else { return false };
    got.as_bytes().ct_eq(expected.as_bytes()).into()
}

/// Prevent LMTP command executed by the HTTP payload
fn sanitize(v: &str) -> String {
    v.chars().filter(|&c| c != '\r' && c != '\n' && c != '>').collect()
}

async fn lmtp_deliver(
    cfg: &Config,
    from: &str,
    to: &str,
    msg: &[u8],
) -> Result<(), (bool, String)> {
    let stream = UnixStream::connect(cfg.lmtp).await.map_err(|e| (false, e.to_string()))?;
    let (rd, mut wr) = stream.into_split();
    let mut rd = BufReader::new(rd);
    let mut line = String::new();

    macro_rules! io {
        ($e:expr) => {
            $e.await.map_err(|e| (false, e.to_string()))?
        };
    }

    macro_rules! expect {
        ($code:expr) => {{
            loop {
                line.clear();
                if io!(rd.read_line(&mut line)) == 0 {
                    return Err((false, "connection closed".into()));
                }
                if line.as_bytes().get(3) != Some(&b'-') {
                    break;
                }
            }
            if !line.starts_with($code) {
                return Err((line.starts_with('5'), line.trim().to_string()));
            }
        }};
    }

    let rcpt = cfg.rcpt.unwrap_or(to);

    expect!("220");
    io!(wr.write_all(b"LHLO localhost\r\n"));
    expect!("250");
    io!(wr.write_all(format!("MAIL FROM:<{}>\r\n", sanitize(from)).as_bytes()));
    expect!("250");
    io!(wr.write_all(format!("RCPT TO:<{}>\r\n", sanitize(rcpt)).as_bytes()));
    expect!("250");
    io!(wr.write_all(b"DATA\r\n"));
    expect!("354");
    if cfg.rcpt.is_some() {
        io!(wr.write_all(format!("X-Original-To: {}\r\n", sanitize(to)).as_bytes()));
    }

    for chunk in msg.split_inclusive(|&b| b == b'\n') {
        if chunk.first() == Some(&b'.') {
            io!(wr.write_all(b"."));
        }
        io!(wr.write_all(chunk));
    }
    if !msg.ends_with(b"\n") {
        io!(wr.write_all(b"\r\n"));
    }
    io!(wr.write_all(b".\r\n"));
    expect!("250");

    let _ = wr.write_all(b"QUIT\r\n").await;
    Ok(())
}