use std::time::{Duration, Instant};

use tokio::{
    io::{AsyncReadExt, AsyncWriteExt},
    net::{TcpListener, TcpStream},
    process::Command,
};

#[tokio::main(flavor = "current_thread")]
async fn main() {
    eprintln!(
        "USAGE:
- Set env var ADDR=addr:port
- Set env var INTERVAL=seconds
- Set env vars ADDR1-ADDRn for n proxies
- Create scripts ./proxy1 - ./proxy(n-1) to check if proxy target is online (exit status 0)
- If all scripts fail, proxy n is used"
    );
    let interval: u64 = std::env::var("INTERVAL")
        .expect("expected env var INTERVAL to be set")
        .trim()
        .parse()
        .expect(
            "expected env var INTERVAL to be set to an integer value greater than or equal to zero",
        );
    let proxies = (1..)
        .map_while(|i| std::env::var(format!("ADDR{i}")).ok())
        .collect::<Vec<_>>();
    if proxies.is_empty() {
        panic!("expected env vars ADDR1-ADDRn");
    }

    let listener =
        TcpListener::bind(std::env::var("ADDR").expect("expected env var ADDR to be set"))
            .await
            .unwrap();

    let mut next_check = Instant::now();
    let mut selected_proxy = 0;
    loop {
        let (con, _) = listener.accept().await.unwrap();
        let now = Instant::now();
        if now >= next_check {
            next_check = now + Duration::from_secs(interval);
            let pproxy = selected_proxy;
            selected_proxy = proxies.len();
            for i in 1..proxies.len() {
                if Command::new(format!("./proxy{i}"))
                    .status()
                    .await
                    .is_ok_and(|status| status.success())
                {
                    selected_proxy = i;
                }
            }
            if selected_proxy != pproxy {
                eprintln!(
                    "Using proxy {selected_proxy}: {}",
                    proxies[selected_proxy - 1]
                );
            }
        }
        if selected_proxy > 0 {
            if let Some(proxy) = proxies.get(selected_proxy - 1).cloned() {
                if let Ok(prox) = TcpStream::connect(&proxy).await {
                    let (mut con_read, mut con_write) = con.into_split();
                    let (mut prox_read, mut prox_write) = prox.into_split();
                    tokio::task::spawn(async move {
                        let mut buf = [0u8; 256];
                        while let Ok(n) = con_read.read(&mut buf).await {
                            if n > 0 && prox_write.write_all(&buf[0..n]).await.is_err() {
                                break;
                            }
                        }
                    });
                    tokio::task::spawn(async move {
                        let mut buf = [0u8; 256];
                        while let Ok(n) = prox_read.read(&mut buf).await {
                            if n > 0 && con_write.write_all(&buf[0..n]).await.is_err() {
                                break;
                            }
                        }
                    });
                }
            }
        }
    }
}