use anyhow::{bail, Context, Result};
use base64::{engine::general_purpose::STANDARD as BASE64, Engine};
use nym_sdk::mixnet::{MixnetClientBuilder, MixnetMessageSender, Recipient, StoragePaths};
use serde::{Deserialize, Serialize};
use std::env;
use std::io::{self, Read};
use std::path::PathBuf;
use std::time::Duration;

const MAX_PAYLOAD_BYTES: usize = 64 * 1024;
const CONNECT_TIMEOUT: Duration = Duration::from_secs(120);
const SEND_TIMEOUT: Duration = Duration::from_secs(60);
// send_plain_message queues the message for the background mixnet task. Keep
// the client alive long enough for that task to flush the message before the
// process disconnects.
const FLUSH_GRACE: Duration = Duration::from_secs(180);

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

#[derive(Debug, Serialize)]
struct Response<'a> {
    success: bool,
    #[serde(skip_serializing_if = "Option::is_none")]
    error: Option<&'a str>,
}

fn respond(response: Response<'_>) -> Result<()> {
    serde_json::to_writer(io::stdout(), &response).context("encode response")?;
    println!();
    Ok(())
}

fn configured_recipient() -> Result<Recipient> {
    let value = env::var("YAMN_NYM_RECIPIENT")
        .context("YAMN_NYM_RECIPIENT is not configured")?;
    value
        .parse::<Recipient>()
        .map_err(|_| anyhow::anyhow!("invalid configured Nym recipient"))
}

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

async fn send(request: Request) -> Result<()> {
    if request.payload.is_empty() || request.payload.len() > MAX_PAYLOAD_BYTES {
        bail!("invalid envelope size");
    }
    if request.payload.as_bytes().contains(&0)
        || request
            .entry_address
            .chars()
            .any(|c| matches!(c, '\r' | '\n' | '\0'))
    {
        bail!("envelope contains a NUL byte");
    }

    let raw_payload = BASE64
        .encode(request.payload.as_bytes());
    let request_payload = serde_json::json!({
        "version": 1,
        "entry_address": request.entry_address,
        "payload": raw_payload,
    });
    let request_payload = serde_json::to_string(&request_payload).context("encode ingress envelope")?;

    let recipient = configured_recipient()?;
    let paths = storage_paths()?;
    let disconnected = MixnetClientBuilder::new_with_default_storage(paths)
        .await
        .context("prepare Nym client")?
        .build()
        .context("build Nym client")?;
    let mut client = tokio::time::timeout(CONNECT_TIMEOUT, disconnected.connect_to_mixnet())
        .await
        .context("Nym connection timed out")?
        .context("connect to Nym mixnet")?;

    tokio::time::timeout(
        SEND_TIMEOUT,
        client.send_plain_message(recipient, request_payload),
    )
    .await
    .context("Nym send timed out")?
    .context("send envelope through Nym")?;

    tokio::time::sleep(FLUSH_GRACE).await;
    client.disconnect().await;
    Ok(())
}

#[tokio::main]
async fn main() {
    let result = async {
        let mut input = String::new();
        io::stdin()
            .read_to_string(&mut input)
            .context("read request")?;
        let request: Request = serde_json::from_str(&input).context("decode request")?;
        send(request).await
    }
    .await;

    match result {
        Ok(()) => {
            let _ = respond(Response {
                success: true,
                error: None,
            });
        }
        Err(error) => {
            eprintln!("yamn-nym-submit: {error:#}");
            let _ = respond(Response {
                success: false,
                error: Some("Nym submission failed"),
            });
            std::process::exit(1);
        }
    }
}