lightstream 0.5.0

// Copyright Peter G. Bower 2025-2026.
//
// This Source Code Form is subject to the terms of the Mozilla Public
// License, v. 2.0. If a copy of the MPL was not distributed with this
// file, You can obtain one at https://mozilla.org/MPL/2.0/.

//! Async Lightstream protocol writer.
//!
//! Wraps a [`LightstreamCodec`](crate::models::codecs::lightstream::LightstreamCodec) and an [`AsyncWrite`](tokio::io::AsyncWrite) destination, providing
//! methods to send messages and Arrow tables over a single connection.
//!
//! Tables are encoded using the Arrow IPC streaming protocol - schema is
//! sent once per table type, then only record batches after that.
//!
//! With the `protobuf` feature, [`send_proto`] provides typed message
//! sending via prost. With the `msgpack` feature, [`send_msgpack`]
//! provides typed message sending via MessagePack and serde.
//!
//! [`send_proto`]: crate::models::writers::lightstream::LightstreamWriter::send_proto
//! [`send_msgpack`]: crate::models::writers::lightstream::LightstreamWriter::send_msgpack

use std::io;

use minarrow::{Field, TableV};
use tokio::io::{AsyncWrite, AsyncWriteExt};

use crate::models::codecs::lightstream::LightstreamCodec;
use crate::models::frames::lightstream_message::LightstreamMessage;
use crate::traits::stream_buffer::StreamBuffer;

/// Async writer for the Lightstream protocol.
///
/// Encodes messages and Arrow tables via the codec and writes the
/// resulting frames to the underlying `AsyncWrite` destination.
/// Table payloads use the Arrow IPC streaming protocol internally.
///
/// Holds a pooled encode buffer that is reused across `send_table` calls.
/// After warmup no per-batch allocation occurs for table encoding.
pub struct LightstreamWriter<W: AsyncWrite + Unpin + Send, B: StreamBuffer + Unpin = Vec<u8>> {
    codec: LightstreamCodec<B>,
    dest: W,
    /// Pooled buffer for direct table encoding, reused across send_table calls.
    encode_buf: B,
}

impl<W: AsyncWrite + Unpin + Send, B: StreamBuffer + Unpin> LightstreamWriter<W, B> {
    /// Create a new writer over the given destination.
    pub fn new(dest: W) -> Self {
        Self {
            codec: LightstreamCodec::new(None),
            dest,
            encode_buf: B::with_capacity(0),
        }
    }

    /// Register a message type. Returns the assigned type tag.
    pub fn register_message(&mut self, name: impl Into<String>) -> u8 {
        self.codec.register_message(name)
    }

    /// Register a table type with the given schema. Returns the assigned type tag.
    pub fn register_table(&mut self, name: impl Into<String>, schema: Vec<Field>) -> u8 {
        self.codec.register_table(name, schema)
    }

    /// Send an opaque message payload by type name.
    pub async fn send(&mut self, name: &str, payload: &[u8]) -> io::Result<()> {
        let tag = self.codec.tag_by_name(name).ok_or_else(|| {
            io::Error::new(
                io::ErrorKind::InvalidInput,
                format!("unknown type name '{}'", name),
            )
        })?;
        let frame = self.codec.encode_message(tag, payload)?;
        self.dest.write_all(frame.as_ref()).await?;
        Ok(())
    }

    /// Send an Arrow table or table view by type name. A whole table
    /// sends as the full-width view of itself.
    ///
    /// Uses the direct encode path which writes column data from the view's
    /// arrays into a pooled buffer in one pass. After warmup the encode buffer
    /// is reused and no per-batch allocation occurs.
    pub async fn send_table(&mut self, name: &str, table: impl Into<TableV>) -> io::Result<()> {
        let view: TableV = table.into();
        let tag = self.codec.tag_by_name(name).ok_or_else(|| {
            io::Error::new(
                io::ErrorKind::InvalidInput,
                format!("unknown type name '{}'", name),
            )
        })?;
        self.codec.encode_table(tag, &view, &mut self.encode_buf)?;
        self.dest.write_all(self.encode_buf.as_ref()).await?;
        Ok(())
    }

    /// Send a frame - a message payload or an Arrow table - by its type tag.
    /// Mirrors the [`LightstreamMessage`] the reader produces, so a frame can be
    /// forwarded without re-resolving its type name.
    pub async fn send_frame(&mut self, frame: &LightstreamMessage) -> io::Result<()> {
        match frame {
            LightstreamMessage::Message { tag, payload } => {
                let encoded = self.codec.encode_message(*tag, payload)?;
                self.dest.write_all(encoded.as_ref()).await?;
            }
            LightstreamMessage::Table { tag, table } => {
                self.codec.encode_table(*tag, table, &mut self.encode_buf)?;
                self.dest.write_all(self.encode_buf.as_ref()).await?;
            }
        }
        Ok(())
    }

    /// Flush the underlying writer.
    pub async fn flush(&mut self) -> io::Result<()> {
        self.dest.flush().await
    }

    /// Shut down the underlying writer.
    pub async fn shutdown(&mut self) -> io::Result<()> {
        self.dest.shutdown().await
    }

    /// Send a protobuf message by type name.
    ///
    /// Encodes the message via `prost::Message::encode_to_vec` and sends it
    /// as an opaque payload.
    #[cfg(feature = "protobuf")]
    pub async fn send_proto<M: prost::Message>(&mut self, name: &str, msg: &M) -> io::Result<()> {
        let bytes = msg.encode_to_vec();
        self.send(name, &bytes).await
    }

    /// Send a MessagePack-encoded message by type name.
    ///
    /// Serialises the value via `rmp-serde` with `BytesMode::ForceAll` so
    /// that `Vec<u8>` and `&[u8]` fields are stored as binary, not as
    /// arrays of integers.
    #[cfg(feature = "msgpack")]
    pub async fn send_msgpack<M: serde::Serialize>(
        &mut self,
        name: &str,
        msg: &M,
    ) -> io::Result<()> {
        let bytes = encode_msgpack(msg)?;
        self.send(name, &bytes).await
    }

    /// Borrow the codec for inspection.
    pub fn codec(&self) -> &LightstreamCodec<B> {
        &self.codec
    }
}

/// Encode a value to MessagePack bytes with efficient binary handling.
#[cfg(feature = "msgpack")]
fn encode_msgpack<M: serde::Serialize>(msg: &M) -> io::Result<Vec<u8>> {
    let mut buf = Vec::new();
    let mut serializer =
        rmp_serde::Serializer::new(&mut buf).with_bytes(rmp_serde::config::BytesMode::ForceAll);
    msg.serialize(&mut serializer)
        .map_err(|e| io::Error::new(io::ErrorKind::InvalidData, e))?;
    Ok(buf)
}