use std::collections::HashMap;
use std::net::SocketAddr;
use tokio::io::copy;
use tokio::net::{TcpListener, TcpStream};

#[tokio::main]
async fn main() -> anyhow::Result<()> {
    let filename = std::env::args()
        .nth(1)
        .unwrap_or_else(|| "config.toml".to_string());

    let contents = std::fs::read_to_string(filename).unwrap();

    let ports: HashMap<SocketAddr, u16> = toml::from_str(&contents)?;

    let mut handles = Vec::with_capacity(ports.len());

    for (&saddr, &port) in &ports {
        let handle = tokio::spawn(async move {
            let listener = TcpListener::bind(SocketAddr::new("0.0.0.0".parse()?, port)).await?;
            loop {
                let (stream, _) = listener.accept().await?;

                tokio::spawn(async move {
                    let to_stream = TcpStream::connect(saddr).await?;
                    proxy(stream, to_stream).await?;
                    Ok::<(), anyhow::Error>(())
                });
            }

            // We still need this statement in order to specify the return type
            #[allow(unreachable_code)]
            Ok::<(), anyhow::Error>(())
        });

        handles.push(handle);
    }

    for handle in handles {
        handle.await??;
    }

    Ok(())
}

async fn proxy(left: TcpStream, right: TcpStream) -> anyhow::Result<()> {
    let (mut left_read, mut left_write) = left.into_split();
    let (mut right_read, mut right_write) = right.into_split();

    tokio::select! {
        r = copy(&mut left_read, &mut right_write) => r?,
        r = copy(&mut right_read, &mut left_write) => r?
    };

    Ok(())
}