use core::time::Duration;
use hierr::Result;
use hirun::net::{AioFd, Fd, SocketAddr, TcpClient};
use hirun::runtime;

fn main() {
    let opts = hiopt::options!["addr:", "a:", "cache:", "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 cache = 0_usize;
    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(SocketAddr::inet_from(opt_value.unwrap()).unwrap()),
            2 => cache = opt_value.unwrap().parse().unwrap(),
            3..=4 => nth = opt_value.unwrap().parse().unwrap(),
            5..=6 => buflen = opt_value.unwrap().parse().unwrap(),
            7..=8 => return help(),
            _ => unreachable!(),
        }
    }

    for addr in opts.noopt_iter(&args[1..]) {
        server_addr = Some(SocketAddr::inet_from(addr).unwrap());
        break;
    }

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

    let _ = runtime::Builder::new().max_cache(cache).nth(nth).build().unwrap();
    let _ = runtime::block_on_with(
        listening(
            listening_addr.as_ref().unwrap(),
            server_addr.as_ref().unwrap(),
            buflen,
        ),
        runtime::Attr::new().priority(1),
    );
}

fn help() {
    println!(
        "Usage: {} options server_addr",
        hictor::program_invocation_name()
    );
    println!("server_addr: <ipv4:port> | <[ipv6]:port>");
    println!("options:");
    println!("--addr | -a: port | <ipv4:port> | <[ipv6]:port>");
    println!("--cache: max task-cache size");
    println!("--threads | -t: defaut is 0");
    println!("--buffer | -b: buflen, defaut is 1024");
    println!("--help | -h: print help message");
}

async fn do_proxy(down: Fd, up: TcpClient, buflen: usize) -> Result<()> {
    let Ok((up, cookie)) = up.wait_connected().await else {
        println!("connect failed: {}", hierr::Error::last());
        return Ok(());
    };
    let mut buf = vec![0; buflen];
    let mut down = AioFd::new(&down);
    let mut up = AioFd::new_with(&up, cookie);
    let _ = down.set_linger(Some(Duration::new(0, 0)));
    let _ = up.set_linger(Some(Duration::new(0, 0)));
    down.copy_bidirectional(&mut up, &mut buf, Duration::new(5, 0))
        .await;
    Ok(())
}

async fn listening(addr: &SocketAddr, remote: &SocketAddr, buflen: usize) -> Result<()> {
    let fd = Fd::tcp_server(addr)?;
    println!("listening {:?}", addr);
    let mut aio = AioFd::new(&fd);
    while let Ok((down, _)) = aio.accept().await {
        let up = Fd::tcp_connect(None, remote)?;
        let _ = runtime::task::spawn(do_proxy(down, up, buflen)).await;
    }
    Ok(())
}