use anyhow::{bail, Context, Result};
use base64::{engine::general_purpose::STANDARD as BASE64, Engine};
use nym_sdk::mixnet::{MixnetClientBuilder, StoragePaths};
use serde::Deserialize;
use std::env;
use std::path::PathBuf;
use std::time::Duration;
use tokio::io::{AsyncBufReadExt, AsyncRead, AsyncWrite, AsyncWriteExt, BufReader};
use tokio_socks::tcp::Socks5Stream;

const MAX_REQUEST_BYTES: usize = 128 * 1024;
const MAX_PAYLOAD_BYTES: usize = 96 * 1024;
const SMTP_TIMEOUT: Duration = Duration::from_secs(60);
const TOR_CONNECT_ATTEMPTS: u32 = 3;

#[derive(Debug, Deserialize)]
struct Envelope {
    version: u8,
    entry_address: String,
    payload: String,
}

fn storage_paths() -> Result<StoragePaths> {
    let path = env::var_os("YAMN_NYM_INGRESS_STORAGE")
        .map(PathBuf::from)
        .unwrap_or_else(|| PathBuf::from("/var/lib/yamnweb/nym-ingress"));
    StoragePaths::new_from_dir(&path).context("prepare Nym ingress storage")
}

fn allowed_entry(address: &str) -> bool {
    env::var("YAMN_ALLOWED_ENTRY_ADDRESSES")
        .ok()
        .map(|list| list.split(',').any(|item| item.trim() == address))
        .unwrap_or(false)
}

fn dry_run() -> bool {
    matches!(
        env::var("YAMN_DRY_RUN").as_deref(),
        Ok("1") | Ok("true") | Ok("yes")
    )
}

fn split_address(address: &str) -> Result<(&str, &str)> {
    if address.len() > 320
        || address
            .chars()
            .any(|c| matches!(c, '\r' | '\n' | '\0' | ' '))
    {
        bail!("invalid entry address");
    }
    let (local, domain) = address
        .rsplit_once('@')
        .context("entry address has no domain")?;
    if local.is_empty() || domain.is_empty() || domain.contains('@') {
        bail!("invalid entry address");
    }
    Ok((local, domain))
}

fn validate_envelope(envelope: &Envelope) -> Result<(String, Vec<u8>)> {
    if envelope.version != 1 {
        bail!("unsupported envelope version");
    }
    split_address(&envelope.entry_address)?;
    if !allowed_entry(&envelope.entry_address) {
        bail!("entry address is not allowlisted");
    }
    let payload = BASE64
        .decode(&envelope.payload)
        .context("invalid base64 YAMN payload")?;
    if payload.is_empty() || payload.len() > MAX_PAYLOAD_BYTES || payload.contains(&0) {
        bail!("invalid YAMN payload size or contents");
    }
    Ok((envelope.entry_address.clone(), payload))
}

async fn smtp_read_response<R: AsyncRead + Unpin>(reader: &mut BufReader<R>) -> Result<u16> {
    let mut line = String::new();
    loop {
        line.clear();
        let read = reader
            .read_line(&mut line)
            .await
            .context("read SMTP response")?;
        if read == 0 {
            bail!("SMTP connection closed");
        }
        if line.len() < 3 {
            bail!("malformed SMTP response");
        }
        let parsed: u16 = line[..3].parse().context("invalid SMTP response code")?;
        if line.as_bytes().get(3) != Some(&b'-') {
            return Ok(parsed);
        }
    }
}

async fn smtp_command<S: AsyncRead + AsyncWrite + Unpin>(
    reader: &mut BufReader<S>,
    command: &[u8],
    expected_class: u16,
) -> Result<()> {
    reader
        .get_mut()
        .write_all(command)
        .await
        .context("write SMTP command")?;
    reader
        .get_mut()
        .flush()
        .await
        .context("flush SMTP command")?;
    let code = smtp_read_response(reader).await?;
    if code / 100 != expected_class {
        bail!("SMTP command rejected with code {code}");
    }
    Ok(())
}

async fn deliver(entry_address: &str, payload: &[u8]) -> Result<()> {
    let (_, host) = split_address(entry_address)?;
    let proxy = env::var("YAMN_TOR_SOCKS").unwrap_or_else(|_| "127.0.0.1:9050".to_string());
    let (proxy_host, proxy_port) = proxy.rsplit_once(':').context("invalid YAMN_TOR_SOCKS")?;
    let proxy_port: u16 = proxy_port.parse().context("invalid YAMN_TOR_SOCKS port")?;
    let target = format!("{host}:25");
    let mut attempt = 1;
    let stream = loop {
        let connection = tokio::time::timeout(
            SMTP_TIMEOUT,
            Socks5Stream::connect((proxy_host, proxy_port), target.as_str()),
        )
        .await;
        match connection {
            Ok(Ok(stream)) => break stream,
            Ok(Err(_)) if attempt < TOR_CONNECT_ATTEMPTS => {
                eprintln!("Tor SMTP connection attempt {attempt} failed; retrying");
            }
            Err(_) if attempt < TOR_CONNECT_ATTEMPTS => {
                eprintln!("Tor SMTP connection attempt {attempt} timed out; retrying");
            }
            Ok(Err(error)) => {
                return Err(error).context("connect to YAMN SMTP entry through Tor");
            }
            Err(error) => return Err(error).context("Tor SMTP connection timed out"),
        }
        tokio::time::sleep(tor_retry_delay(attempt)).await;
        attempt += 1;
    };
    let mut reader = BufReader::new(stream);
    let greeting = smtp_read_response(&mut reader).await?;
    if greeting / 100 != 2 {
        bail!("YAMN SMTP greeting rejected with code {greeting}");
    }

    let from = env::var("YAMN_INGRESS_FROM").unwrap_or_else(|_| "<>".to_string());
    if from.chars().any(|c| matches!(c, '\r' | '\n' | ' ')) {
        bail!("invalid YAMN_INGRESS_FROM");
    }
    smtp_command(&mut reader, b"EHLO yamn-nym-ingress\r\n", 2).await?;
    smtp_command(&mut reader, format!("MAIL FROM:{from}\r\n").as_bytes(), 2).await?;
    smtp_command(
        &mut reader,
        format!("RCPT TO:<{entry_address}>\r\n").as_bytes(),
        2,
    )
    .await?;
    smtp_command(&mut reader, b"DATA\r\n", 3).await?;

    for line in payload.split_inclusive(|byte| *byte == b'\n') {
        if line.starts_with(b".") {
            reader
                .get_mut()
                .write_all(b".")
                .await
                .context("dot-stuff SMTP payload")?;
        }
        reader
            .get_mut()
            .write_all(line)
            .await
            .context("write SMTP payload")?;
    }
    if !payload.ends_with(b"\n") {
        reader
            .get_mut()
            .write_all(b"\r\n")
            .await
            .context("terminate SMTP payload")?;
    }
    reader
        .get_mut()
        .write_all(b".\r\n")
        .await
        .context("finish SMTP payload")?;
    reader
        .get_mut()
        .flush()
        .await
        .context("flush SMTP payload")?;
    let accepted = smtp_read_response(&mut reader).await?;
    if accepted / 100 != 2 {
        bail!("YAMN SMTP server rejected payload with code {accepted}");
    }
    let _ = smtp_command(&mut reader, b"QUIT\r\n", 2).await;
    Ok(())
}

fn tor_retry_delay(failed_attempt: u32) -> Duration {
    Duration::from_secs(5 * u64::from(failed_attempt))
}

async fn run() -> Result<()> {
    let disconnected = MixnetClientBuilder::new_with_default_storage(storage_paths()?)
        .await
        .context("prepare Nym client")?
        .build()
        .context("build Nym client")?;
    let mut client = disconnected
        .connect_to_mixnet()
        .await
        .context("connect ingress to Nym mixnet")?;
    eprintln!("yamn-nym-ingress address: {}", client.nym_address());

    while let Some(messages) = client.wait_for_messages().await {
        for message in messages {
            if message.message.len() > MAX_REQUEST_BYTES {
                eprintln!("discarded oversized Nym message");
                continue;
            }
            let envelope: Envelope = match serde_json::from_slice(&message.message) {
                Ok(value) => value,
                Err(_) => {
                    eprintln!("discarded malformed Nym envelope");
                    continue;
                }
            };
            let (entry_address, payload) = match validate_envelope(&envelope) {
                Ok(value) => value,
                Err(error) => {
                    eprintln!("discarded invalid YAMN envelope: {error}");
                    continue;
                }
            };
            if dry_run() {
                eprintln!("accepted YAMN envelope in dry-run mode");
                continue;
            }
            if let Err(error) = deliver(&entry_address, &payload).await {
                eprintln!("YAMN delivery failed: {error:#}");
            } else {
                eprintln!("delivered YAMN envelope to allowlisted entry");
            }
        }
    }
    client.disconnect().await;
    Ok(())
}

#[tokio::main]
async fn main() {
    if let Err(error) = run().await {
        eprintln!("yamn-nym-ingress: {error:#}");
        std::process::exit(1);
    }
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn tor_retry_delay_uses_bounded_linear_backoff() {
        assert_eq!(tor_retry_delay(1), Duration::from_secs(5));
        assert_eq!(tor_retry_delay(2), Duration::from_secs(10));
    }
}