use std::net::SocketAddr;
use tokio::io::copy_bidirectional_with_sizes;
use tokio::net::{TcpListener, TcpStream};
use tokio::runtime::*;

fn main() {
    let opts = hiopt::options!["addr:", "a:", "threads:", "t:", "buffer:", "b:", "help", "h"];

    let args = unsafe { hiopt::raw_args_from_i8(hictor::args()) };

    let mut listening_addr = None;
    let mut server_addr = None;
    let mut nth = 0_usize;
    let mut buflen = 1024_usize;
    // ignore program name
    for opt in opts.opt_iter(&args[1..]) {
        let (opt_idx, opt_value) = opt.unwrap();
        match opt_idx {
            0..=1 => listening_addr = Some(format!("127.0.0.1:{}", opt_value.unwrap())),
            2..=3 => nth = opt_value.unwrap().parse().unwrap(),
            4..=5 => buflen = opt_value.unwrap().parse().unwrap(),
            6..=7 => return help(),
            _ => unreachable!(),
        }
    }

    for addr in opts.noopt_iter(&args[1..]) {
        server_addr = Some(format!("127.0.0.1:{addr}"));
        break;
    }

    if listening_addr.is_none() || server_addr.is_none() {
        return help();
    }

    let rt = if nth == 0 {
        Builder::new_multi_thread().enable_all().build().unwrap()
    } else {
        Builder::new_multi_thread()
            .enable_all()
            .worker_threads(nth)
            .build()
            .unwrap()
    };
    let _ = rt.block_on(listening(
        (&listening_addr.unwrap()).parse().unwrap(),
        (&server_addr.unwrap()).parse().unwrap(),
        buflen,
    ));
}

fn help() {
    println!(
        "Usage: {} options server_addr",
        hictor::program_invocation_name()
    );
    println!("server_addr: port");
    println!("options:");
    println!("--addr | -a: port");
    println!("--threads | -t: defaut is 0");
    println!("--buffer | -b: buflen, defaut is 1024");
    println!("--help | -h: print help message");
}

async fn listening(addr: SocketAddr, remote: SocketAddr, buflen: usize) {
    let Ok(listener) = TcpListener::bind(addr).await else {
        println!("listen fail: {addr}");
        return;
    };

    while let Ok((mut inbound, _)) = listener.accept().await {
        let addr = remote.clone();
        tokio::spawn(async move {
            let Ok(mut outbound) = TcpStream::connect(addr).await else {
                return;
            };
            let _ =
                copy_bidirectional_with_sizes(&mut inbound, &mut outbound, buflen, buflen).await;
        });
    }
}