super_cereal 0.1.0

use crate::message::{SerialChunk, CHUNK_SIZE};
use crate::transport::{self, ChunkReader, ChunkWriter};
use anyhow::{Context, Result};
use futures::StreamExt;
use tokio::io::{AsyncRead, AsyncReadExt, AsyncWrite, AsyncWriteExt};
use tracing::{debug, warn};

pub async fn pump_read_to_up(
    mut reader: impl AsyncRead + Unpin,
    up_writer: &ChunkWriter,
    label: &str,
) -> Result<()> {
    let mut buf = vec![0u8; CHUNK_SIZE];
    let mut seq = 0u64;

    loop {
        let n = reader
            .read(&mut buf)
            .await
            .with_context(|| format!("{label} read failed"))?;
        if n == 0 {
            continue;
        }

        let chunk = SerialChunk::new(seq, buf[..n].to_vec());
        seq = seq.wrapping_add(1);
        debug!(label, seq = chunk.seq, bytes = n, "forwarding chunk up");
        transport::dds::write_chunk_async(up_writer, chunk).await?;
    }
}

pub async fn pump_read_to_down(
    mut reader: impl AsyncRead + Unpin,
    down_writer: &ChunkWriter,
    label: &str,
) -> Result<()> {
    let mut buf = vec![0u8; CHUNK_SIZE];
    let mut seq = 0u64;

    loop {
        let n = reader
            .read(&mut buf)
            .await
            .with_context(|| format!("{label} read failed"))?;
        if n == 0 {
            continue;
        }

        let chunk = SerialChunk::new(seq, buf[..n].to_vec());
        seq = seq.wrapping_add(1);
        debug!(label, seq = chunk.seq, bytes = n, "forwarding chunk down");
        transport::dds::write_chunk_async(down_writer, chunk).await?;
    }
}

pub async fn pump_down_to_write(
    down_reader: ChunkReader,
    mut writer: impl AsyncWrite + Unpin,
    label: &str,
) -> Result<()> {
    let mut stream = down_reader.async_bare_sample_stream();
    while let Some(sample) = stream.next().await {
        match sample {
            Ok(chunk) => {
                debug!(
                    label,
                    seq = chunk.seq,
                    bytes = chunk.data.len(),
                    "writing chunk from down"
                );
                writer
                    .write_all(&chunk.data)
                    .await
                    .with_context(|| format!("{label} write failed"))?;
                writer
                    .flush()
                    .await
                    .with_context(|| format!("{label} flush failed"))?;
            }
            Err(e) => {
                warn!(label, error = ?e, "DDS read error on down topic");
                return Err(anyhow::anyhow!("DDS read failed on down topic: {e:?}"));
            }
        }
    }
    Ok(())
}

pub async fn pump_up_to_write(
    up_reader: ChunkReader,
    mut writer: impl AsyncWrite + Unpin,
    label: &str,
) -> Result<()> {
    let mut stream = up_reader.async_bare_sample_stream();
    while let Some(sample) = stream.next().await {
        match sample {
            Ok(chunk) => {
                debug!(
                    label,
                    seq = chunk.seq,
                    bytes = chunk.data.len(),
                    "writing chunk from up"
                );
                writer
                    .write_all(&chunk.data)
                    .await
                    .with_context(|| format!("{label} write failed"))?;
                writer
                    .flush()
                    .await
                    .with_context(|| format!("{label} flush failed"))?;
            }
            Err(e) => {
                warn!(label, error = ?e, "DDS read error on up topic");
                return Err(anyhow::anyhow!("DDS read failed on up topic: {e:?}"));
            }
        }
    }
    Ok(())
}

pub async fn run_serve_bridge(
    serial: impl AsyncRead + AsyncWrite + Unpin + Send + 'static,
    channel: transport::DdsChannel,
) -> Result<()> {
    let (mut serial_read, mut serial_write) = tokio::io::split(serial);
    let up_writer = channel.up_writer;
    let down_reader = channel.down_reader;

    let up_task = tokio::spawn(async move {
        pump_read_to_up(&mut serial_read, &up_writer, "serial").await
    });
    let down_task = tokio::spawn(async move {
        pump_down_to_write(down_reader, &mut serial_write, "serial").await
    });

    let result = tokio::select! {
        res = up_task => res,
        res = down_task => res,
    };
    result.context("serve bridge task join failed")??;
    Ok(())
}

pub async fn run_attach_bridge(
    pty: impl AsyncRead + AsyncWrite + Unpin + Send + 'static,
    channel: transport::DdsChannel,
) -> Result<()> {
    let (mut pty_read, mut pty_write) = tokio::io::split(pty);
    let down_writer = channel.down_writer;
    let up_reader = channel.up_reader;

    let down_task = tokio::spawn(async move {
        pump_read_to_down(&mut pty_read, &down_writer, "pty").await
    });
    let up_task = tokio::spawn(async move {
        pump_up_to_write(up_reader, &mut pty_write, "pty").await
    });

    let result = tokio::select! {
        res = down_task => res,
        res = up_task => res,
    };
    result.context("attach bridge task join failed")??;
    Ok(())
}