use crate::listener::Listener;
use rand::seq::SliceRandom;
use std::error::Error;
use tokio::net::{TcpListener, TcpStream};

#[derive(Debug)]
struct TcpHandler {
    stream: TcpStream,
    target: String,
}

impl TcpHandler {
    async fn run(&mut self) -> Result<(), Box<dyn Error>> {
        let mut stream = TcpStream::connect(&self.target).await?;

        tokio::io::copy_bidirectional(&mut self.stream, &mut stream).await?;

        return Ok(());
    }
}

pub(crate) async fn start_tcp_listener(
    mut listener_config: Listener,
) -> Result<(), Box<dyn Error>> {
    println!("start listening on {}", &listener_config.source);
    let listener = TcpListener::bind(&listener_config.source).await?;

    loop {
        let (next_socket, _) = tokio::select! {
           res = listener.accept() => res?,
           _ = listener_config.shutdown.recv() => {
              println!("Exiting listener!");
              return Ok(());
           }
        };

        let targets = listener_config.targets.read().await;
        let mut rng = rand::thread_rng();
        let selected_target = targets.choose(&mut rng).unwrap();

        println!(
            "new connection from {} forwarding to {}",
            next_socket.peer_addr()?,
            &selected_target
        );
        let mut handler = TcpHandler {
            stream: next_socket,
            target: selected_target.clone(),
        };

        tokio::spawn(async move {
            // Process the connection. If an error is encountered, log it.
            if let Err(err) = handler.run().await {
                println!("connection error {}", err);
            }
        });
    }
}