arcly-stream 0.1.3

//! The live stream handle — a lock-free, zero-copy broadcast fan-out bus.

use crate::observe::{NoopObserver, Observer};
use crate::{frame::FrameFlags, AppName, MediaFrame, StreamId, StreamKey};
use arc_swap::ArcSwap;
use std::net::SocketAddr;
use std::sync::atomic::{AtomicU64, Ordering};
use std::sync::{Arc, Mutex as StdMutex};
use std::time::{SystemTime, UNIX_EPOCH};
use tokio::sync::broadcast::error::RecvError;
use tokio::sync::{broadcast, RwLock};

/// Current wall-clock time in Unix milliseconds (saturating to 0 pre-epoch).
pub(crate) fn now_ms() -> u64 {
    SystemTime::now()
        .duration_since(UNIX_EPOCH)
        .map(|d| d.as_millis() as u64)
        .unwrap_or(0)
}

/// Current lifecycle state of a stream.
#[derive(Debug, Clone, PartialEq, Eq)]
pub enum StreamState {
    /// No publisher yet.
    Idle,
    /// A publisher is connected and sending data.
    Publishing,
    /// The stream is being transcoded into one or more renditions.
    Transcoding,
    /// The stream is being recorded.
    Recording,
    /// The publisher has disconnected; the stream has ended.
    Ended,
}

impl std::fmt::Display for StreamState {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        f.write_str(match self {
            StreamState::Idle => "idle",
            StreamState::Publishing => "publishing",
            StreamState::Transcoding => "transcoding",
            StreamState::Recording => "recording",
            StreamState::Ended => "ended",
        })
    }
}

/// Runtime metadata about a stream, updated continuously while publishing.
///
/// Resolution and the ingest protocol are set by the protocol handler (e.g. via
/// [`StreamHandle::update_metadata`] after parsing the codec config), while the
/// `fps` and `*_bitrate_bps` fields are overlaid live from measured throughput
/// by [`StreamHandle::metadata_snapshot`].
#[derive(Debug, Clone)]
pub struct StreamMetadata {
    /// The `(app, stream_id)` this metadata describes.
    pub key: StreamKey,
    /// Publisher remote address.
    pub publisher_addr: Option<SocketAddr>,
    /// Video width in pixels (0 = unknown).
    pub width: u32,
    /// Video height in pixels (0 = unknown).
    pub height: u32,
    /// Video frames per second (0 = unknown). Overlaid from measured throughput.
    pub fps: f64,
    /// Measured video ingest bitrate in bits-per-second.
    pub video_bitrate_bps: u64,
    /// Measured audio ingest bitrate in bits-per-second.
    pub audio_bitrate_bps: u64,
    /// Timestamp of the first frame received (Unix ms).
    pub started_at_ms: u64,
    /// Protocol used for ingest (e.g. `"rtmp"`).
    pub ingest_protocol: String,
}

impl StreamMetadata {
    /// Create zeroed metadata for `(app, stream_id)`.
    pub fn new(app: AppName, stream_id: StreamId) -> Self {
        Self {
            key: StreamKey::new(app, stream_id),
            publisher_addr: None,
            width: 0,
            height: 0,
            fps: 0.0,
            video_bitrate_bps: 0,
            audio_bitrate_bps: 0,
            started_at_ms: 0,
            ingest_protocol: String::new(),
        }
    }
}

/// A point-in-time snapshot of a stream's measured quality of service.
#[derive(Debug, Clone, Copy, Default)]
pub struct Qos {
    /// Video bitrate over the last ~1s window (bits/sec).
    pub video_bitrate_bps: u64,
    /// Audio bitrate over the last ~1s window (bits/sec).
    pub audio_bitrate_bps: u64,
    /// Video frames per second over the last ~1s window.
    pub fps: f64,
    /// Cumulative frames published on this stream.
    pub total_frames: u64,
    /// Cumulative payload bytes published on this stream.
    pub total_bytes: u64,
}

/// Lock-free throughput counters folded into [`Qos`] / [`StreamMetadata`].
///
/// A stream has a single publisher (enforced by `start_publish`), so the
/// read-modify-write here is effectively single-writer and `Relaxed` is sound.
#[derive(Default)]
struct QosCounters {
    total_frames: AtomicU64,
    total_bytes: AtomicU64,
    window_start_ms: AtomicU64,
    window_video_bytes: AtomicU64,
    window_audio_bytes: AtomicU64,
    window_video_frames: AtomicU64,
    cur_video_bitrate: AtomicU64,
    cur_audio_bitrate: AtomicU64,
    cur_fps_milli: AtomicU64, // fps × 1000, integer-encoded
    last_frame_ms: AtomicU64,
}

/// Keyframe-anchored replay buffer for instant playback start.
struct GopBuffer {
    /// Frames since (and including) the most recent keyframe.
    frames: Vec<Arc<MediaFrame>>,
    /// Hard cap on buffered frames (memory bound between keyframes).
    capacity: usize,
}

/// A live handle to a single active stream.
///
/// Multiple subscribers (HLS packager, DASH packager, WebRTC SFU, recorders …)
/// call [`StreamHandle::subscribe_resilient`] to receive every [`MediaFrame`]
/// cheaply via a `broadcast` channel (zero-copy `Bytes` cloning).
///
/// Each broadcast slot holds one `Arc<MediaFrame>` pointer (8 bytes), so e.g.
/// 4096 slots ≈ 32 KB per stream.
#[derive(Clone)]
pub struct StreamHandle {
    metadata: Arc<RwLock<StreamMetadata>>,
    state: Arc<RwLock<StreamState>>,
    key: StreamKey,
    tx: broadcast::Sender<Arc<MediaFrame>>,
    /// Latest video CONFIG (AVCDecoderConfigurationRecord) frame, if seen.
    /// Uses `ArcSwap` for lock-free reads from multiple subscriber tasks.
    video_config: Arc<ArcSwap<Option<Arc<MediaFrame>>>>,
    /// Latest audio CONFIG (AudioSpecificConfig) frame, if seen.
    audio_config: Arc<ArcSwap<Option<Arc<MediaFrame>>>>,
    /// Rolling GOP buffer for instant-start (empty when `gop_capacity == 0`).
    gop: Arc<StdMutex<GopBuffer>>,
    gop_capacity: usize,
    /// Live throughput counters.
    qos: Arc<QosCounters>,
    /// Injected telemetry hook (no-op by default).
    observer: Arc<dyn Observer>,
}

impl StreamHandle {
    /// Create a handle with the no-op observer and no GOP cache.
    pub fn new(app: AppName, stream_id: StreamId, capacity: usize) -> Self {
        Self::with_observer(app, stream_id, capacity, 0, Arc::new(NoopObserver))
    }

    /// Create a handle wired to a host-supplied observer.
    ///
    /// `gop_capacity` bounds the keyframe-anchored replay buffer (0 disables it).
    pub fn with_observer(
        app: AppName,
        stream_id: StreamId,
        capacity: usize,
        gop_capacity: usize,
        observer: Arc<dyn Observer>,
    ) -> Self {
        let (tx, _) = broadcast::channel(capacity);
        let qos = QosCounters::default();
        // Treat creation as the last activity so a just-claimed stream is not
        // instantly considered idle before its first frame arrives.
        qos.last_frame_ms.store(now_ms(), Ordering::Relaxed);
        Self {
            metadata: Arc::new(RwLock::new(StreamMetadata::new(
                app.clone(),
                stream_id.clone(),
            ))),
            state: Arc::new(RwLock::new(StreamState::Idle)),
            key: StreamKey::new(app, stream_id),
            tx,
            video_config: Arc::new(ArcSwap::new(Arc::new(None))),
            audio_config: Arc::new(ArcSwap::new(Arc::new(None))),
            gop: Arc::new(StdMutex::new(GopBuffer {
                frames: Vec::new(),
                capacity: gop_capacity,
            })),
            gop_capacity,
            qos: Arc::new(qos),
            observer,
        }
    }

    /// The `(app, stream_id)` this handle belongs to.
    pub fn key(&self) -> &StreamKey {
        &self.key
    }

    /// Publish a frame to all current subscribers.  Returns the number of
    /// active receivers; returns `Ok(0)` when there are no subscribers.
    pub fn publish_frame(&self, frame: MediaFrame) -> crate::Result<usize> {
        self.observer.on_frame(&self.key, &frame);
        let len = frame.data.len() as u64;
        let is_audio = frame.is_audio();
        let is_key = frame.is_keyframe();
        let is_config = frame.flags.contains(FrameFlags::CONFIG);
        let arc = Arc::new(frame);

        // Cache the latest CONFIG frame for late-joining subscribers.
        if is_config {
            if is_audio {
                self.audio_config.store(Arc::new(Some(Arc::clone(&arc))));
            } else {
                self.video_config.store(Arc::new(Some(Arc::clone(&arc))));
            }
        }

        // Maintain the keyframe-anchored GOP replay buffer.
        if self.gop_capacity > 0 {
            if let Ok(mut g) = self.gop.lock() {
                if is_key {
                    g.frames.clear();
                    g.frames.push(Arc::clone(&arc));
                } else if !is_config && !g.frames.is_empty() && g.frames.len() < g.capacity {
                    // Only buffer once a keyframe anchors the GOP; CONFIG frames
                    // are replayed separately via `cached_configs`.
                    g.frames.push(Arc::clone(&arc));
                }
            }
        }

        self.record_qos(len, is_audio, is_key);

        let count = self.tx.send(arc).unwrap_or(0);
        Ok(count)
    }

    /// Fold one frame into the rolling throughput window.
    fn record_qos(&self, len: u64, is_audio: bool, _is_key: bool) {
        let q = &self.qos;
        let now = now_ms();
        q.total_frames.fetch_add(1, Ordering::Relaxed);
        q.total_bytes.fetch_add(len, Ordering::Relaxed);
        q.last_frame_ms.store(now, Ordering::Relaxed);
        if is_audio {
            q.window_audio_bytes.fetch_add(len, Ordering::Relaxed);
        } else {
            q.window_video_bytes.fetch_add(len, Ordering::Relaxed);
            q.window_video_frames.fetch_add(1, Ordering::Relaxed);
        }

        let ws = q.window_start_ms.load(Ordering::Relaxed);
        if ws == 0 {
            q.window_start_ms.store(now, Ordering::Relaxed);
        } else if now.saturating_sub(ws) >= 1000 {
            let elapsed = (now - ws) as f64 / 1000.0;
            let vbytes = q.window_video_bytes.swap(0, Ordering::Relaxed);
            let abytes = q.window_audio_bytes.swap(0, Ordering::Relaxed);
            let vframes = q.window_video_frames.swap(0, Ordering::Relaxed);
            q.cur_video_bitrate
                .store((vbytes as f64 * 8.0 / elapsed) as u64, Ordering::Relaxed);
            q.cur_audio_bitrate
                .store((abytes as f64 * 8.0 / elapsed) as u64, Ordering::Relaxed);
            q.cur_fps_milli.store(
                (vframes as f64 / elapsed * 1000.0) as u64,
                Ordering::Relaxed,
            );
            q.window_start_ms.store(now, Ordering::Relaxed);
        }
    }

    /// A snapshot of measured throughput (bitrate, fps, totals).
    pub fn qos(&self) -> Qos {
        let q = &self.qos;
        Qos {
            video_bitrate_bps: q.cur_video_bitrate.load(Ordering::Relaxed),
            audio_bitrate_bps: q.cur_audio_bitrate.load(Ordering::Relaxed),
            fps: q.cur_fps_milli.load(Ordering::Relaxed) as f64 / 1000.0,
            total_frames: q.total_frames.load(Ordering::Relaxed),
            total_bytes: q.total_bytes.load(Ordering::Relaxed),
        }
    }

    /// Unix-ms timestamp of the most recently published frame (or stream
    /// creation if none yet). Used by the engine's idle reaper.
    pub fn last_frame_ms(&self) -> u64 {
        self.qos.last_frame_ms.load(Ordering::Relaxed)
    }

    /// Returns the most recently seen video and audio CONFIG frames,
    /// for replaying to late-joining subscribers.
    pub fn cached_configs(&self) -> (Option<Arc<MediaFrame>>, Option<Arc<MediaFrame>>) {
        let video = (**self.video_config.load()).clone();
        let audio = (**self.audio_config.load()).clone();
        (video, audio)
    }

    /// The frames a late joiner should be handed before going live: cached
    /// decoder configs followed by the current GOP (keyframe + trailing deltas).
    ///
    /// Replaying these lets a new subscriber start decoding immediately rather
    /// than waiting for the next keyframe — sub-second join times at scale.
    /// Requires the app to have enabled a GOP cache; otherwise only the cached
    /// configs are returned.
    pub fn replay_buffer(&self) -> Vec<Arc<MediaFrame>> {
        let (vcfg, acfg) = self.cached_configs();
        let mut out = Vec::new();
        out.extend(vcfg);
        out.extend(acfg);
        if self.gop_capacity > 0 {
            if let Ok(g) = self.gop.lock() {
                for f in &g.frames {
                    // Avoid duplicating a config frame already pushed above.
                    if !out.iter().any(|c| Arc::ptr_eq(c, f)) {
                        out.push(Arc::clone(f));
                    }
                }
            }
        }
        out
    }

    /// Subscribe to this stream's frame bus.
    ///
    /// The returned raw [`broadcast::Receiver`] surfaces [`RecvError::Lagged`]
    /// when a slow consumer falls behind the channel capacity — callers that
    /// `while let Ok(_) = rx.recv().await` will silently terminate on the first
    /// lag. Prefer [`subscribe_resilient`](Self::subscribe_resilient) unless you
    /// are deliberately handling lag yourself.
    pub fn subscribe(&self) -> broadcast::Receiver<Arc<MediaFrame>> {
        self.tx.subscribe()
    }

    /// Subscribe with a [`Subscription`] that resynchronizes after lag instead
    /// of terminating, reporting each gap to the installed [`Observer`] via
    /// [`Observer::on_subscriber_lagged`].
    pub fn subscribe_resilient(&self) -> Subscription {
        Subscription {
            rx: self.tx.subscribe(),
            key: self.key.clone(),
            observer: Arc::clone(&self.observer),
            max_lag: None,
            skipped: 0,
        }
    }

    /// Number of active subscribers.
    pub fn subscriber_count(&self) -> usize {
        self.tx.receiver_count()
    }

    /// Transition to a new state.
    pub async fn set_state(&self, state: StreamState) {
        let mut guard = self.state.write().await;
        *guard = state;
    }

    /// The current lifecycle state.
    pub async fn current_state(&self) -> StreamState {
        self.state.read().await.clone()
    }

    /// A consistent point-in-time copy of this stream's [`StreamMetadata`], with
    /// the live measured `fps`/bitrate overlaid from [`qos`](Self::qos).
    ///
    /// Cloning the snapshot releases the lock immediately, so callers never hold
    /// the metadata `RwLock` across an `.await`.
    pub async fn metadata_snapshot(&self) -> StreamMetadata {
        let mut m = self.metadata.read().await.clone();
        let q = self.qos();
        m.video_bitrate_bps = q.video_bitrate_bps;
        m.audio_bitrate_bps = q.audio_bitrate_bps;
        if q.fps > 0.0 {
            m.fps = q.fps;
        }
        m
    }

    /// Mutate this stream's [`StreamMetadata`] under the write lock.
    ///
    /// Ingest handlers call this as they parse the stream — e.g. on the first
    /// keyframe to record resolution from the codec config, or to set the
    /// publisher address — so the metadata exposed to operators and the control
    /// plane stays live rather than frozen at its zeroed defaults.
    ///
    /// ```no_run
    /// # use arcly_stream::StreamHandle;
    /// # async fn demo(handle: &StreamHandle, addr: std::net::SocketAddr) {
    /// handle
    ///     .update_metadata(|m| {
    ///         m.publisher_addr = Some(addr);
    ///         m.width = 1920;
    ///         m.height = 1080;
    ///         m.ingest_protocol = "rtmp".to_string();
    ///     })
    ///     .await;
    /// # }
    /// ```
    pub async fn update_metadata(&self, f: impl FnOnce(&mut StreamMetadata)) {
        let mut guard = self.metadata.write().await;
        f(&mut guard);
    }
}

impl std::fmt::Debug for StreamHandle {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        f.debug_struct("StreamHandle")
            .field("key", &self.key)
            .field("subscribers", &self.subscriber_count())
            .finish()
    }
}

/// A lag-tolerant subscription to a stream's frame bus.
///
/// Returned by [`StreamHandle::subscribe_resilient`]. Unlike a raw
/// [`broadcast::Receiver`], [`recv`](Self::recv) does not terminate when the
/// consumer falls behind: the dropped span is reported to the [`Observer`] as
/// [`on_subscriber_lagged`](Observer::on_subscriber_lagged) and reception
/// continues from the oldest still-buffered frame. This is the recommended
/// consumer loop for packagers, recorders, and SFUs.
///
/// An optional [`max_lag`](Self::max_lag) bound turns chronic lag into
/// eviction: once cumulative dropped frames exceed the bound, `recv` returns
/// `None` (after an [`on_subscriber_evicted`](Observer::on_subscriber_evicted)
/// notification) so a hopelessly slow consumer is shed rather than wasting
/// buffer churn forever.
pub struct Subscription {
    rx: broadcast::Receiver<Arc<MediaFrame>>,
    key: StreamKey,
    observer: Arc<dyn Observer>,
    max_lag: Option<u64>,
    skipped: u64,
}

impl Subscription {
    /// Evict this subscriber once cumulative dropped frames exceed `max`.
    ///
    /// ```no_run
    /// # use arcly_stream::StreamHandle;
    /// # fn demo(handle: &StreamHandle) {
    /// let sub = handle.subscribe_resilient().max_lag(10_000);
    /// # let _ = sub;
    /// # }
    /// ```
    pub fn max_lag(mut self, max: u64) -> Self {
        self.max_lag = Some(max);
        self
    }

    /// Total frames dropped from this subscriber's view so far.
    pub fn dropped(&self) -> u64 {
        self.skipped
    }

    /// Receive the next frame, resynchronizing past any lag.
    ///
    /// Returns `None` when the stream's sender is dropped (the publisher ended)
    /// or when the `max_lag` eviction threshold is crossed:
    ///
    /// ```no_run
    /// # async fn run(sub: &mut arcly_stream::bus::Subscription) {
    /// while let Some(frame) = sub.recv().await {
    ///     // packetize `frame` …
    /// }
    /// # }
    /// ```
    pub async fn recv(&mut self) -> Option<Arc<MediaFrame>> {
        loop {
            match self.rx.recv().await {
                Ok(frame) => return Some(frame),
                Err(RecvError::Lagged(skipped)) => {
                    self.skipped = self.skipped.saturating_add(skipped);
                    self.observer.on_subscriber_lagged(&self.key, skipped);
                    if let Some(max) = self.max_lag {
                        if self.skipped > max {
                            self.observer.on_subscriber_evicted(&self.key);
                            return None;
                        }
                    }
                    continue;
                }
                Err(RecvError::Closed) => return None,
            }
        }
    }

    /// Borrow the underlying raw receiver, for callers that need
    /// [`broadcast::Receiver`] APIs directly.
    pub fn raw(&mut self) -> &mut broadcast::Receiver<Arc<MediaFrame>> {
        &mut self.rx
    }
}