wingfoil 6.0.5

use crate::channel::{
    ChannelReceiver, ChannelSender, Message, ReceiverMessageSource, channel_pair,
};
use crate::nodes::channel::ChannelReceiverStream;
use crate::*;

use anyhow::Context as _;
use futures::stream::StreamExt;
use std::future::Future;
use std::pin::Pin;
use std::rc::Rc;

/// Context passed to async producer closures during graph setup.
///
/// This provides the run configuration so producers can adapt their behavior
/// (e.g., derive time ranges for database queries).
#[derive(Clone, Copy, Debug)]
pub struct RunParams {
    pub run_mode: RunMode,
    pub run_for: RunFor,
    pub start_time: NanoTime,
}

impl RunParams {
    /// Compute end time based on run_for.
    ///
    /// Returns `start_time + duration` for `RunFor::Duration`,
    /// `NanoTime::MAX` for `Forever`, or an error for `Cycles`.
    pub fn end_time(&self) -> anyhow::Result<NanoTime> {
        match self.run_for {
            RunFor::Duration(d) => Ok(self.start_time + d),
            RunFor::Forever => Ok(NanoTime::MAX),
            RunFor::Cycles(_) => anyhow::bail!("end_time not available for RunFor::Cycles"),
        }
    }
}

/// A convenience alias for [`futures::Stream`] with items of type `(NanoTime, T)`.
/// used by [StreamOperators::consume_async].
pub trait FutStream<T>: futures::Stream<Item = (NanoTime, T)> + Send {}

impl<STRM, T> FutStream<T> for STRM where STRM: futures::Stream<Item = (NanoTime, T)> + Send {}

type ConsumerFunc<T, FUT> = Box<dyn FnOnce(RunParams, Pin<Box<dyn FutStream<T>>>) -> FUT + Send>;

// Routes the source stream into a user-supplied async consumer via a kanal
// `ChannelSender`. The channel is currently unbounded (see `channel_pair`), so
// a slow consumer grows memory rather than back-pressuring the graph.
//
// Shape-wise this is the channel-backed equivalent of `FixSenderNode` in
// adapters/fix/mod.rs, which writes directly from `cycle` and gets back-
// pressure for free from the kernel TCP buffer. If `channel_pair` grows a
// bounded+blocking mode, the two could share a single "sink node"
// abstraction parameterised by the actual write call.
pub(crate) struct AsyncConsumerNode<T, FUT>
where
    T: Element + Send,
    FUT: Future<Output = anyhow::Result<()>> + Send + 'static,
{
    source: Rc<dyn Stream<T>>,
    sender: ChannelSender<T>,
    func: Option<ConsumerFunc<T, FUT>>,
    handle: Option<tokio::task::JoinHandle<anyhow::Result<()>>>,
    rx: Option<ChannelReceiver<T>>,
}

impl<T, FUT> AsyncConsumerNode<T, FUT>
where
    T: Element + Send,
    FUT: Future<Output = anyhow::Result<()>> + Send + 'static,
{
    pub fn new(source: Rc<dyn Stream<T>>, func: ConsumerFunc<T, FUT>) -> Self {
        let (sender, receiver) = channel_pair(None);
        let rx = Some(receiver);
        let handle = None;

        let func = Some(func);
        Self {
            source,
            sender,
            func,
            handle,
            rx,
        }
    }
}

#[node(active = [source])]
impl<T, FUT> MutableNode for AsyncConsumerNode<T, FUT>
where
    T: Element + Send,
    FUT: Future<Output = anyhow::Result<()>> + Send + 'static,
{
    fn cycle(&mut self, state: &mut GraphState) -> anyhow::Result<bool> {
        // Check if the consumer task has finished early (due to an error).
        // If it has, try to get the actual error instead of "channel closed".
        if self.handle.as_ref().is_some_and(|h| h.is_finished()) {
            // The consumer task has exited unexpectedly. Try to get the actual error.
            let handle = self.handle.take().expect("handle is Some");
            match state.tokio_runtime().block_on(handle) {
                Ok(Ok(())) => {
                    // Task completed successfully, but that's unexpected here
                    anyhow::bail!("consumer task exited early without error");
                }
                Ok(Err(e)) => {
                    // Task returned an error - this is what we want to report
                    return Err(e).context("consumer task failed");
                }
                Err(e) => {
                    // Task panicked or was cancelled
                    anyhow::bail!("consumer task panicked or was cancelled: {e}");
                }
            }
        }
        self.sender.send(state, self.source.peek_value())?;
        Ok(true)
    }

    fn setup(&mut self, state: &mut GraphState) -> anyhow::Result<()> {
        let run_mode = state.run_mode();
        let run_for = state.run_for();
        let rx = self
            .rx
            .take()
            .ok_or_else(|| anyhow::anyhow!("rx is already taken"))?;

        let func = self
            .func
            .take()
            .ok_or_else(|| anyhow::anyhow!("func is already taken"))?;
        let ctx = RunParams {
            run_mode,
            run_for,
            start_time: state.start_time(),
        };

        let f = async move {
            let src = rx
                .to_boxed_message_stream()
                .limit(run_mode, run_for)
                .to_stream();
            let fut = func(ctx, Box::pin(src));
            fut.await
        };

        let handle = state.tokio_runtime().spawn(f);
        self.handle = Some(handle);

        Ok(())
    }

    fn stop(&mut self, _state: &mut GraphState) -> anyhow::Result<()> {
        self.sender.close()?;
        Ok(())
    }

    fn teardown(&mut self, state: &mut GraphState) -> anyhow::Result<()> {
        if let Some(handle) = self.handle.take() {
            state.tokio_runtime().block_on(handle)??;
        }
        Ok(())
    }
}

struct AsyncProducerStream<T, S, FUT, FUNC>
where
    T: Element + Send,
    S: futures::Stream<Item = anyhow::Result<(NanoTime, T)>> + Send + 'static,
    FUT: Future<Output = anyhow::Result<S>> + Send + 'static,
    FUNC: FnOnce(RunParams) -> FUT + Send + 'static,
{
    func: Option<FUNC>,
    receiver_stream: ChannelReceiverStream<T>,
    handle: Option<tokio::task::JoinHandle<anyhow::Result<()>>>,
    sender: Option<ChannelSender<T>>,
}

impl<T, S, FUT, FUNC> AsyncProducerStream<T, S, FUT, FUNC>
where
    T: Element + Send,
    S: futures::Stream<Item = anyhow::Result<(NanoTime, T)>> + Send + 'static,
    FUT: Future<Output = anyhow::Result<S>> + Send + 'static,
    FUNC: FnOnce(RunParams) -> FUT + Send + 'static,
{
    pub fn new(func: FUNC) -> Self {
        let (sender, receiver) = channel_pair(None);
        let receiver_stream = ChannelReceiverStream::new(receiver, None, None);
        let handle = None;
        let func = Some(func);
        let sender = Some(sender);
        Self {
            receiver_stream,
            func,
            handle,
            sender,
        }
    }
}

impl<T, S, FUT, FUNC> MutableNode for AsyncProducerStream<T, S, FUT, FUNC>
where
    T: Element + Send,
    S: futures::Stream<Item = anyhow::Result<(NanoTime, T)>> + Send + 'static,
    FUT: Future<Output = anyhow::Result<S>> + Send + 'static,
    FUNC: FnOnce(RunParams) -> FUT + Send + 'static,
{
    fn cycle(&mut self, state: &mut GraphState) -> anyhow::Result<bool> {
        self.receiver_stream.cycle(state)
    }

    fn setup(&mut self, state: &mut GraphState) -> anyhow::Result<()> {
        let run_mode = state.run_mode();
        let run_for = state.run_for();
        let mut sender = self
            .sender
            .take()
            .ok_or_else(|| anyhow::anyhow!("sender is already taken"))?;

        match run_mode {
            RunMode::HistoricalFrom(_) => {}
            RunMode::RealTime => sender.set_notifier(state.ready_notifier()),
        };
        let mut sender = sender.into_async();
        let func = self
            .func
            .take()
            .ok_or_else(|| anyhow::anyhow!("func is already taken"))?;
        let ctx = RunParams {
            run_mode,
            run_for,
            start_time: state.start_time(),
        };
        let fut = async move {
            match func(ctx).await {
                Err(e) => {
                    sender
                        .send_message(Message::Error(std::sync::Arc::new(e)))
                        .await;
                }
                Ok(stream) => {
                    let source = stream.to_message_stream(run_mode).limit(run_mode, run_for);
                    let mut source = Box::pin(source);
                    while let Some(message) = source.next().await {
                        sender.send_message(message).await;
                    }
                }
            }
            sender.close().await;
            Ok(())
        };
        let handle = state.tokio_runtime().spawn(fut);
        self.handle = Some(handle);
        self.receiver_stream.setup(state)?;
        Ok(())
    }

    fn teardown(&mut self, state: &mut GraphState) -> anyhow::Result<()> {
        if let Some(handle) = self.handle.take() {
            // Abort first so the sender inside the task is dropped, then await completion
            // so receiver_stream.teardown() can drain the channel without blocking.
            handle.abort();
            let result = state.tokio_runtime().block_on(handle);
            self.receiver_stream.teardown(state)?;
            match result {
                Ok(inner) => inner?,
                Err(e) if e.is_cancelled() => {}
                Err(e) => return Err(anyhow::anyhow!("producer task panicked: {e}")),
            }
        }
        Ok(())
    }
}

impl<T, S, FUT, FUNC> StreamPeekRef<Burst<T>> for AsyncProducerStream<T, S, FUT, FUNC>
where
    T: Element + Send,
    S: futures::Stream<Item = anyhow::Result<(NanoTime, T)>> + Send + 'static,
    FUT: Future<Output = anyhow::Result<S>> + Send + 'static,
    FUNC: FnOnce(RunParams) -> FUT + Send + 'static,
{
    fn peek_ref(&self) -> &Burst<T> {
        self.receiver_stream.peek_ref()
    }
}

/// Create a [Stream] from a fallible async function that produces a futures::Stream.
///
/// The closure receives an [`RunParams`] with `run_mode`, `run_for`, and `start_time`,
/// allowing producers to adapt their behavior (e.g., derive time ranges for queries).
///
/// # Example
/// ```ignore
/// produce_async(|ctx| async move {
///     let start = ctx.start_time;
///     let end = ctx.end_time();
///     Ok(async_stream::stream! {
///         // Use start, end in async code...
///     })
/// })
/// ```
#[must_use]
pub fn produce_async<T, S, FUT, FUNC>(func: FUNC) -> Rc<dyn Stream<Burst<T>>>
where
    T: Element + Send,
    S: futures::Stream<Item = anyhow::Result<(NanoTime, T)>> + Send + 'static,
    FUT: Future<Output = anyhow::Result<S>> + Send + 'static,
    FUNC: FnOnce(RunParams) -> FUT + Send + 'static,
{
    AsyncProducerStream::new(func).into_stream()
}

trait StreamMessageSource<T: Element + Send> {
    fn to_message_stream(self, run_mode: RunMode) -> impl futures::Stream<Item = Message<T>>;
}

impl<T, STRM> StreamMessageSource<T> for STRM
where
    STRM: futures::Stream<Item = anyhow::Result<(NanoTime, T)>>,
    T: Element + Send,
{
    fn to_message_stream(self, run_mode: RunMode) -> impl futures::Stream<Item = Message<T>> {
        async_stream::stream! {
            let mut source = Box::pin(self);
            // Group consecutive same-time values into one atomic `HistoricalValue`
            // message so a slow producer can't split a timestamp across the
            // receiver's monotonic clock. The source is time-ordered, so a change
            // in `time` marks the previous group complete.
            let mut pending: Option<ValueAt<Burst<T>>> = None;
            while let Some(result) = source.next().await {
                match result {
                    Ok((time, value)) => match run_mode {
                        RunMode::RealTime => yield Message::RealtimeValue(value),
                        RunMode::HistoricalFrom(_) => match &mut pending {
                            Some(group) if group.time == time => group.value.push(value),
                            _ => {
                                if let Some(group) = pending.take() {
                                    yield Message::HistoricalValue(group);
                                }
                                pending = Some(ValueAt::new(crate::burst![value], time));
                            }
                        },
                    },
                    Err(e) => {
                        if let Some(group) = pending.take() {
                            yield Message::HistoricalValue(group);
                        }
                        yield Message::Error(std::sync::Arc::new(e));
                    }
                }
            }
            if let Some(group) = pending.take() {
                yield Message::HistoricalValue(group);
            }
            yield Message::EndOfStream;
        }
    }
}

trait MessageStream<T: Element + Send> {
    fn limit(self, run_mode: RunMode, run_for: RunFor) -> impl futures::Stream<Item = Message<T>>;

    fn to_stream(self) -> impl futures::Stream<Item = (NanoTime, T)>;
}

impl<T, STRM> MessageStream<T> for STRM
where
    STRM: futures::Stream<Item = Message<T>>,
    T: Element + Send,
{
    fn limit(self, run_mode: RunMode, run_for: RunFor) -> impl futures::Stream<Item = Message<T>> {
        async_stream::stream! {
            let time0 = run_mode.start_time();
            let mut time = time0;
            let mut elapsed: NanoTime;
            let mut cycle = 0;
            let mut source = Box::pin(self);
            let mut finished = false;
            while let Some(message) = source.next().await
            {
                match &message {
                    Message::RealtimeValue(_) => {
                        time = NanoTime::now();
                    }
                    Message::HistoricalValue(value_at) => {
                        time = value_at.time;
                    },
                    Message::CheckPoint(t) => {
                        time = *t;
                    },
                    Message::EndOfStream => {
                        finished = true;
                    },
                    Message::Error(_) => {
                        // Error messages are passed through
                    },
                }
                elapsed = time - time0;
                yield (message);
                if run_for.done(cycle, elapsed) || finished {
                    break;
                }
                cycle +=1;
            }
        }
    }

    fn to_stream(self) -> impl futures::Stream<Item = (NanoTime, T)> {
        async_stream::stream! {
            let mut source = Box::pin(self);
            while let Some(message) = source.next().await {
                match message {
                    Message::RealtimeValue(value) => {
                        yield (NanoTime::now(), value)
                    }
                    Message::HistoricalValue(value_at) => {
                        let time = value_at.time;
                        for value in value_at.value {
                            yield (time, value)
                        }
                    },
                    Message::CheckPoint(_) => {},
                    Message::EndOfStream => {},
                    Message::Error(err) => {
                        // Log the error and stop the stream.
                        // This error likely came from an async producer or consumer task.
                        log::error!("Error in async stream: {err:#?}");
                        break;
                    },
                }
            }
        }
    }
}

#[cfg(test)]
mod tests {

    use crate::*;
    use futures::StreamExt;
    use std::pin::Pin;
    use std::time::Duration;

    /// An async producer that emits same-time values with latency between them
    /// must not drive engine time backwards. The `sleep` makes the producer lag
    /// the receiver, so without producer-side grouping the second value lands
    /// behind engine time and aborts with "time less than graph time".
    #[test]
    fn async_same_time_burst_breaks_monotonic_engine_time() {
        let _ = env_logger::try_init();

        let producer = move |_ctx: RunParams| async move {
            Ok(async_stream::stream! {
                yield Ok((NanoTime::new(100), 1u32));
                // Lands after the receiver has drained the first value and let
                // engine time advance past t=100.
                tokio::time::sleep(Duration::from_millis(50)).await;
                yield Ok((NanoTime::new(100), 2u32));
            })
        };

        let collected = produce_async(producer).collect();

        collected
            .run(RunMode::HistoricalFrom(NanoTime::ZERO), RunFor::Forever)
            .expect("async same-time burst must not drive engine time backwards");

        let delivered: Vec<(u32, NanoTime)> = collected
            .peek_value()
            .iter()
            .flat_map(|burst| burst.value.iter().map(|v| (*v, burst.time)))
            .collect();

        assert_eq!(
            delivered,
            vec![(1, NanoTime::new(100)), (2, NanoTime::new(100))],
            "expected both same-time values delivered at t=100, got {delivered:?}"
        );
    }

    #[test]
    fn produce_async_mid_stream_error() {
        let _ = env_logger::try_init();
        let period = Duration::from_millis(10);

        let producer = move |ctx: RunParams| async move {
            Ok(async_stream::stream! {
                for i in 0..3u32 {
                    let time = ctx.start_time + period * i;
                    yield Ok((time, i));
                }
                yield Err(anyhow::anyhow!("something went wrong"));
            })
        };

        let result = produce_async(producer)
            .collapse()
            .run(RunMode::HistoricalFrom(NanoTime::ZERO), RunFor::Forever);

        let err = result.unwrap_err();
        let root = err.root_cause().to_string();
        assert!(
            root.contains("something went wrong"),
            "expected error message, got: {root}"
        );
    }

    #[test]
    fn async_io_works() {
        let _ = env_logger::try_init();
        let n_runs = 5;
        let period = Duration::from_millis(10);
        let n_periods = 5;
        let run_for = RunFor::Duration(period * n_periods);

        for _ in 0..n_runs {
            for run_mode in [RunMode::RealTime, RunMode::HistoricalFrom(NanoTime::ZERO)] {
                let example_producer = move |ctx: RunParams| async move {
                    Ok(async_stream::stream! {
                        for i in 0.. {
                            let time = match ctx.run_mode {
                                RunMode::HistoricalFrom(_) => {
                                    // wire up historical source here
                                    ctx.start_time + period * i
                                },
                                RunMode::RealTime => {
                                    // wire up real time source here
                                    tokio::time::sleep(period).await; // simulate waiting IO
                                    NanoTime::now()
                                }
                            };
                            yield Ok((time, i * 10));
                        }
                    })
                };

                let example_consumer =
                    async move |_ctx: RunParams, mut source: Pin<Box<dyn FutStream<u32>>>| {
                        while let Some((time, value)) = source.next().await {
                            println!("{time:?}, {value:?}");
                        }
                        Ok(())
                    };

                produce_async(example_producer)
                    .collapse()
                    .logged("on-graph", log::Level::Info)
                    .consume_async(Box::new(example_consumer))
                    .run(run_mode, run_for)
                    .unwrap();
            }
        }
    }
}