melodium-engine 0.10.4

use crate::debug::{DataContent, Event, EventKind, TransmissionDebug, TransmissionDetails};
use crate::transmission::{own, Input};
use async_std::channel::{Sender, TrySendError};
use async_std::sync::Mutex as AsyncMutex;
use async_trait::async_trait;
use core::sync::atomic::{AtomicUsize, Ordering};
use futures::stream::{FuturesUnordered, StreamExt};
use melodium_common::descriptor::Flow;
use melodium_common::executive::{
    Output as ExecutiveOutput, SendResult, TrackId, TransmissionError, TransmissionValue, Value,
};
use std::sync::{Arc, Mutex, OnceLock};

const DEFAULT_MAX_BUFFERED_BYTES: usize = 8 * 1024 * 1024;

/// How many bytes an `Output` accumulates (across however many elements that takes) before
/// a `Stream` flow is forced to wait for a receiver instead of continuing to buffer locally.
/// Byte-based rather than element-based so a stream of large values (big strings, buffers,
/// structured data) can't grow unboundedly just because it hasn't reached an element-count
/// threshold yet. Overridable through `MELODIUM_TRANSMISSION_MAX_BUFFERED_BYTES`, mainly for
/// tests that want to exercise the backpressure path without pushing megabytes of data.
fn max_buffered_bytes() -> usize {
    static LIMIT: OnceLock<usize> = OnceLock::new();
    *LIMIT.get_or_init(|| {
        std::env::var("MELODIUM_TRANSMISSION_MAX_BUFFERED_BYTES")
            .ok()
            .and_then(|value| value.parse().ok())
            .unwrap_or(DEFAULT_MAX_BUFFERED_BYTES)
    })
}

/// Buffered batch awaiting transmission, together with its estimated byte size so
/// `check_send` doesn't have to re-walk the whole batch on every call to decide whether
/// the byte watermark has been crossed.
#[derive(Debug, Default)]
struct Buffer {
    data: Option<TransmissionValue>,
    bytes: usize,
}

impl Buffer {
    fn take(&mut self) -> Option<TransmissionValue> {
        self.bytes = 0;
        self.data.take()
    }

    fn replace(&mut self, data: TransmissionValue) {
        self.bytes = data.estimated_size();
        self.data.replace(data);
    }
}

#[derive(Debug)]
pub struct Output {
    senders: Mutex<Arc<Vec<(Sender<Arc<TransmissionValue>>, Option<TransmissionDetails>)>>>,
    count_receivers: AtomicUsize,
    buffer: AsyncMutex<Buffer>,
    flow: Flow,
    track_id: TrackId,
    debug: TransmissionDebug,
}

impl Output {
    pub fn new(flow: Flow, track_id: TrackId, debug: TransmissionDebug) -> Self {
        Self {
            senders: Mutex::new(Arc::new(Vec::new())),
            count_receivers: AtomicUsize::new(0),
            buffer: AsyncMutex::new(Buffer::default()),
            flow,
            track_id,
            debug,
        }
    }

    pub fn flow(&self) -> &Flow {
        &self.flow
    }

    pub fn track_id(&self) -> &TrackId {
        &self.track_id
    }

    pub fn transmission_debug(&self) -> &TransmissionDebug {
        &self.debug
    }

    pub fn add_transmission(&self, inputs: &Vec<Input>) {
        let mut senders = self.senders.lock().unwrap();
        let count = inputs.len();
        // An output is not supposed to have transmission added while it is already in use,
        // so get_mut on Arc is doable.
        if let Some(senders) = Arc::get_mut(&mut senders) {
            for input in inputs {
                senders.push((
                    input.sender().clone(),
                    match input.transmission_debug() {
                        TransmissionDebug::None => None,
                        TransmissionDebug::Basic(_, details)
                        | TransmissionDebug::Detailed(_, details) => Some(details.clone()),
                    },
                ));
            }
            self.count_receivers.fetch_add(count, Ordering::Relaxed);
        }
    }

    async fn check_send(&self, force: bool) -> SendResult {
        let (buffer_len, buffer_bytes) = {
            let buffer = self.buffer.lock().await;
            (
                buffer.data.as_ref().map(|buf| buf.len()).unwrap_or(0),
                buffer.bytes,
            )
        };

        if buffer_len > 0 {
            // We can unwrap the `take` because buffer_len must be > 0, so buffer have value.
            let data = self.buffer.lock().await.take().unwrap();
            // Wrapped once here rather than cloned per receiver below: fan-out sends hand
            // out cheap `Arc::clone`s of this single allocation instead of deep-cloning the
            // whole batch once per receiver (see `own`, used on the receiving end).
            let data = Arc::new(data);
            if self.flow == Flow::Block || buffer_bytes >= max_buffered_bytes() || force {
                match self.count_receivers.load(Ordering::Relaxed) {
                    0 => Err(TransmissionError::NoReceiver),
                    1 => {
                        let senders = Arc::clone(&self.senders.lock().unwrap());
                        if let Some((sender, input_transmission_details)) = senders.first() {
                            match sender.send(data).await {
                                Ok(_) => {
                                    match (&self.debug, input_transmission_details) {
                                        (_, None) | (TransmissionDebug::None, _) => {}
                                        (
                                            TransmissionDebug::Basic(world, output_details),
                                            Some(input_details),
                                        )
                                        | (
                                            TransmissionDebug::Detailed(world, output_details),
                                            Some(input_details),
                                        ) => {
                                            world
                                                .send_debug_async(Event::new(
                                                    EventKind::DataTransmitted {
                                                        output: output_details.clone(),
                                                        input: input_details.clone(),
                                                        track_id: self.track_id.clone(),
                                                        data: DataContent::Count {
                                                            count: buffer_len,
                                                        },
                                                    },
                                                ))
                                                .await
                                        }
                                    }
                                    Ok(())
                                }
                                Err(_) => Err(TransmissionError::EverythingClosed),
                            }
                        } else {
                            Err(TransmissionError::NoReceiver)
                        }
                    }
                    _ => {
                        let senders = Arc::clone(&self.senders.lock().unwrap());

                        let transmissions = FuturesUnordered::new();
                        for (sender, input_transmission_details) in senders.iter() {
                            let transmission = {
                                let data = &data;
                                async move {
                                    match sender.send(data.clone()).await {
                                        Ok(_) => {
                                            match (&self.debug, input_transmission_details) {
                                                (_, None) | (TransmissionDebug::None, _) => {}
                                                (
                                                    TransmissionDebug::Basic(world, output_details),
                                                    Some(input_details),
                                                )
                                                | (
                                                    TransmissionDebug::Detailed(
                                                        world,
                                                        output_details,
                                                    ),
                                                    Some(input_details),
                                                ) => {
                                                    world
                                                        .send_debug_async(Event::new(
                                                            EventKind::DataTransmitted {
                                                                output: output_details.clone(),
                                                                input: input_details.clone(),
                                                                track_id: self.track_id.clone(),
                                                                data: DataContent::Count {
                                                                    count: buffer_len,
                                                                },
                                                            },
                                                        ))
                                                        .await
                                                }
                                            }
                                            true
                                        }
                                        Err(_) => false,
                                    }
                                }
                            };
                            transmissions.push(transmission);
                        }

                        let statuses: Vec<_> = transmissions.collect().await;

                        if let Some(_) = statuses.iter().find(|s| **s) {
                            Ok(())
                        } else {
                            Err(TransmissionError::EverythingClosed)
                        }
                    }
                }
            } else {
                match self.count_receivers.load(Ordering::Relaxed) {
                    0 => Err(TransmissionError::NoReceiver),
                    1 => {
                        let senders = Arc::clone(&self.senders.lock().unwrap());
                        if let Some((sender, input_transmission_details)) = senders.first() {
                            match sender.try_send(data) {
                                Ok(_) => {
                                    match (&self.debug, input_transmission_details) {
                                        (_, None) | (TransmissionDebug::None, _) => {}
                                        (
                                            TransmissionDebug::Basic(world, output_details),
                                            Some(input_details),
                                        )
                                        | (
                                            TransmissionDebug::Detailed(world, output_details),
                                            Some(input_details),
                                        ) => {
                                            world
                                                .send_debug_async(Event::new(
                                                    EventKind::DataTransmitted {
                                                        output: output_details.clone(),
                                                        input: input_details.clone(),
                                                        track_id: self.track_id.clone(),
                                                        data: DataContent::Count {
                                                            count: buffer_len,
                                                        },
                                                    },
                                                ))
                                                .await
                                        }
                                    }
                                    Ok(())
                                }
                                Err(TrySendError::Full(data)) => {
                                    self.buffer.lock().await.replace(own(data));
                                    Ok(())
                                }
                                Err(TrySendError::Closed(_)) => {
                                    Err(TransmissionError::EverythingClosed)
                                }
                            }
                        } else {
                            Err(TransmissionError::NoReceiver)
                        }
                    }
                    _ => {
                        let senders = Arc::clone(&self.senders.lock().unwrap());

                        let all_senders_not_full =
                            !senders.iter().any(|(sender, _)| sender.is_full());

                        if all_senders_not_full {
                            let transmissions = FuturesUnordered::new();
                            for (sender, input_transmission_details) in senders.iter() {
                                let transmission = {
                                    let data = &data;
                                    async move {
                                        match sender.try_send(data.clone()) {
                                            Ok(_) => {
                                                match (&self.debug, input_transmission_details) {
                                                    (_, None) | (TransmissionDebug::None, _) => {}
                                                    (
                                                        TransmissionDebug::Basic(
                                                            world,
                                                            output_details,
                                                        ),
                                                        Some(input_details),
                                                    )
                                                    | (
                                                        TransmissionDebug::Detailed(
                                                            world,
                                                            output_details,
                                                        ),
                                                        Some(input_details),
                                                    ) => {
                                                        world
                                                            .send_debug_async(Event::new(
                                                                EventKind::DataTransmitted {
                                                                    output: output_details.clone(),
                                                                    input: input_details.clone(),
                                                                    track_id: self.track_id.clone(),
                                                                    data: DataContent::Count {
                                                                        count: buffer_len,
                                                                    },
                                                                },
                                                            ))
                                                            .await
                                                    }
                                                }
                                                true
                                            }
                                            Err(TrySendError::Full(_)) => unreachable!(),
                                            Err(TrySendError::Closed(_)) => false,
                                        }
                                    }
                                };
                                transmissions.push(transmission);
                            }

                            let statuses: Vec<_> = transmissions.collect().await;

                            if let Some(_) = statuses.iter().find(|s| **s) {
                                Ok(())
                            } else {
                                Err(TransmissionError::EverythingClosed)
                            }
                        } else {
                            self.buffer.lock().await.replace(own(data));
                            Ok(())
                        }
                    }
                }
            }
        } else {
            Ok(())
        }
    }
}

#[async_trait]
impl ExecutiveOutput for Output {
    async fn close(&self) {
        let _ = self.check_send(true).await;
        self.senders.lock().unwrap().iter().for_each(|(s, _)| {
            s.close();
        });
        match &self.debug {
            TransmissionDebug::None => {}
            TransmissionDebug::Basic(world, transmission_details)
            | TransmissionDebug::Detailed(world, transmission_details) => {
                world
                    .send_debug_async(Event::new(EventKind::OutputClosed {
                        output: transmission_details.clone(),
                        track_id: self.track_id.clone(),
                    }))
                    .await
            }
        }
    }

    async fn send_many(&self, data: TransmissionValue) -> SendResult {
        match &self.debug {
            TransmissionDebug::None => {}
            TransmissionDebug::Basic(world, transmission_details) => {
                world
                    .send_debug_async(Event::new(EventKind::DataSent {
                        output: transmission_details.clone(),
                        track_id: self.track_id.clone(),
                        data: DataContent::Count { count: data.len() },
                    }))
                    .await
            }
            TransmissionDebug::Detailed(world, transmission_details) => {
                world
                    .send_debug_async(Event::new(EventKind::DataSent {
                        output: transmission_details.clone(),
                        track_id: self.track_id.clone(),
                        data: DataContent::Values {
                            values: data.clone().into(),
                        },
                    }))
                    .await
            }
        }

        {
            let mut lock = self.buffer.lock().await;
            let incoming_bytes = data.estimated_size();
            if let Some(buf) = lock.data.as_mut() {
                buf.append(data);
            } else {
                lock.data = Some(data);
            }
            lock.bytes += incoming_bytes;
        }

        self.check_send(false).await
    }

    async fn send_one(&self, data: Value) -> SendResult {
        match &self.debug {
            TransmissionDebug::None => {}
            TransmissionDebug::Basic(world, transmission_details) => {
                world
                    .send_debug_async(Event::new(EventKind::DataSent {
                        output: transmission_details.clone(),
                        track_id: self.track_id.clone(),
                        data: DataContent::Count { count: 1 },
                    }))
                    .await
            }
            TransmissionDebug::Detailed(world, transmission_details) => {
                world
                    .send_debug_async(Event::new(EventKind::DataSent {
                        output: transmission_details.clone(),
                        track_id: self.track_id.clone(),
                        data: DataContent::Values {
                            values: vec![data.clone()],
                        },
                    }))
                    .await
            }
        }

        {
            let mut lock = self.buffer.lock().await;
            let incoming_bytes = data.estimated_size();
            if let Some(buf) = lock.data.as_mut() {
                buf.push(data);
            } else {
                lock.data = Some(TransmissionValue::new(data));
            }
            lock.bytes += incoming_bytes;
        }
        self.check_send(false).await
    }

    async fn force_send(&self) {
        let _ = self.check_send(true).await;
    }
}

impl From<Input> for Output {
    fn from(value: Input) -> Self {
        let o = Output::new(*value.flow(), *value.track_id(), TransmissionDebug::None);
        o.add_transmission(&vec![value]);
        o
    }
}

#[cfg(test)]
mod fan_out_tests {
    use super::*;
    use melodium_common::executive::{Input as ExecutiveInput, Output as ExecutiveOutput};

    #[test]
    fn single_receiver_gets_sent_value() {
        async_std::task::block_on(async {
            let input = Input::new(Flow::Stream, 0, TransmissionDebug::None);
            let output = Output::new(Flow::Stream, 0, TransmissionDebug::None);
            output.add_transmission(&vec![input.clone()]);

            output.send_one(Value::U8(7)).await.unwrap();

            assert_eq!(input.recv_one().await.unwrap(), Value::U8(7));
        });
    }

    // Every receiver must see the full, correct batch even though the underlying
    // TransmissionValue is now shared (Arc-cloned) across receivers rather than
    // deep-cloned once per receiver — this is the correctness guard for that change.
    #[test]
    fn every_fan_out_receiver_gets_the_full_correct_batch() {
        async_std::task::block_on(async {
            let input_a = Input::new(Flow::Stream, 0, TransmissionDebug::None);
            let input_b = Input::new(Flow::Stream, 0, TransmissionDebug::None);
            let output = Output::new(Flow::Stream, 0, TransmissionDebug::None);
            output.add_transmission(&vec![input_a.clone(), input_b.clone()]);

            // Sends are interleaved with receives: Output only retries a locally
            // buffered batch on the next send/close call, there's no background
            // flush, so a second send before both receivers drain the first would
            // sit buffered forever in this single-task test.
            output.send_one(Value::U8(1)).await.unwrap();
            assert_eq!(input_a.recv_one().await.unwrap(), Value::U8(1));
            assert_eq!(input_b.recv_one().await.unwrap(), Value::U8(1));

            output.send_one(Value::U8(2)).await.unwrap();
            assert_eq!(input_a.recv_one().await.unwrap(), Value::U8(2));
            assert_eq!(input_b.recv_one().await.unwrap(), Value::U8(2));
        });
    }
}

#[cfg(test)]
mod backpressure_tests {
    use super::*;
    use async_std::future::timeout;
    use melodium_common::executive::{Input as ExecutiveInput, Output as ExecutiveOutput};
    use std::collections::VecDeque;
    use std::time::Duration;

    fn wire_pair() -> (Output, Input) {
        let input = Input::new(Flow::Stream, 0, TransmissionDebug::None);
        let output = Output::new(Flow::Stream, 0, TransmissionDebug::None);
        output.add_transmission(&vec![input.clone()]);
        (output, input)
    }

    // Below the byte watermark, a full receiver slot must not block the producer: the
    // batch stays in Output's local buffer and is retried opportunistically instead.
    #[test]
    fn small_batch_does_not_block_once_receiver_slot_is_full() {
        async_std::task::block_on(async {
            let (output, _input) = wire_pair();

            // Occupy the receiver's single buffered slot.
            output.send_one(Value::Byte(1)).await.unwrap();

            let result = timeout(Duration::from_millis(200), output.send_one(Value::Byte(2))).await;
            assert!(
                result.is_ok(),
                "a send small enough to stay under the byte watermark must not block \
                 waiting for a receiver, even if the receiver hasn't drained yet"
            );
        });
    }

    // Once buffered bytes cross the watermark, Output must switch from opportunistic
    // buffering to actually waiting for the receiver — this is the core backpressure
    // guarantee that bounds memory regardless of how large individual values are.
    #[test]
    fn crossing_byte_watermark_blocks_until_receiver_drains() {
        async_std::task::block_on(async {
            let (output, input) = wire_pair();

            // Occupy the receiver's single buffered slot.
            output.send_one(Value::Byte(1)).await.unwrap();

            let big =
                TransmissionValue::Byte(VecDeque::from(vec![0u8; DEFAULT_MAX_BUFFERED_BYTES + 1]));
            let result = timeout(Duration::from_millis(200), output.send_many(big)).await;
            assert!(
                result.is_err(),
                "a send crossing the byte watermark must block waiting for the receiver \
                 to drain, instead of buffering unbounded bytes locally"
            );

            // Drain so the test doesn't leave a task hanging; not asserted further since
            // the send above was already cancelled by the timeout.
            let _ = input.recv_one().await;
        });
    }
}