use std::{future::Future, time::Duration};

use actix_web::web::Bytes;
use russh::{ChannelId, server::Handle};
use tokio::{
    io::{AsyncRead, AsyncReadExt},
    time::sleep,
};

pub async fn forward<'a, R, Fut, Fwd>(
    session_handle: &'a Handle,
    chan_id: ChannelId,
    r: &mut R,
    mut fwd: Fwd,
) -> Result<(), russh::Error>
where
    R: AsyncRead + Send + Unpin,
    Fut: Future<Output = Result<(), Bytes>> + 'a,
    Fwd: FnMut(&'a Handle, ChannelId, Bytes) -> Fut,
{
    const BUF_SIZE: usize = 1024 * 32;
    const MAX_RETRIES: usize = 5;
    const RETRY_DELAY: u64 = 10;

    let mut buf = [0u8; BUF_SIZE];
    loop {
        let read = r.read(&mut buf).await.map_err(russh::Error::IO)?;

        if read == 0 {
            break;
        }

        let mut chunk = Bytes::copy_from_slice(&buf[..read]);
        let mut retries = 0;
        loop {
            match fwd(session_handle, chan_id, chunk).await {
                Ok(()) => break,
                Err(unsent) => {
                    retries += 1;
                    if retries >= MAX_RETRIES {
                        return Ok(());
                    }
                    chunk = unsent;
                    sleep(Duration::from_millis(RETRY_DELAY)).await;
                }
            }
        }
    }

    Ok(())
}