use anyhow::{Context, Result};
use async_std::{
    io,
    io::prelude::*,
    net::{Ipv4Addr, SocketAddrV4, TcpListener, TcpStream},
    prelude::*,
    task,
};
use std::{env, sync::Arc, time::Duration};

static DURATION: Duration = Duration::from_secs(1);

/// Listens for incoming TCP connections and handles them.
#[async_std::main]
async fn main() -> Result<()> {
    let listener = TcpListener::bind(read_addr()?)
        .await
        .with_context(|| "tarpit: failed to bind TCP listener")?;

    let banner = Arc::new(read_banner().await?);

    let mut incoming = listener.incoming();
    while let Some(stream) = incoming.next().await {
        match stream {
            Err(e) => {
                eprintln!("error\t{}", e);
                continue;
            }
            Ok(s) => {
                let banner = banner.clone();
                task::spawn(handler(s, banner));
            }
        }
    }

    Ok(())
}

/// Handles a single TCP connection, continuously writing banner to the TCP stream sloooooowly.
async fn handler(mut stream: TcpStream, banner: Arc<String>) {
    let peer_addr = stream
        .peer_addr()
        .map(|p| p.to_string())
        .unwrap_or_else(|_| String::from("(unknown)"));

    eprintln!("connected\t{}", peer_addr);

    for word in banner.split_inclusive(char::is_whitespace).cycle() {
        if stream.write_all(word.as_bytes()).await.is_err() {
            eprintln!("disconnected\t{}", peer_addr);
            return;
        }
        task::sleep(DURATION).await;
    }
}

/// Reads address port to bind to from first arguent.
fn read_addr() -> Result<SocketAddrV4> {
    let port = env::args()
        .nth(1)
        .map(|arg| arg.parse())
        .unwrap_or(Ok(22))
        .with_context(|| "tarpit: failed to parse bind port")?;

    Ok(SocketAddrV4::new(Ipv4Addr::new(0, 0, 0, 0), port))
}

/// Reads banner to be written to client from STDIN.
async fn read_banner() -> Result<String> {
    let mut banner = String::new();
    io::stdin()
        .read_to_string(&mut banner)
        .await
        .with_context(|| "tarpit: failed to read banner")?;
    Ok(banner)
}