visual_cortex/

session.rs

use std::collections::HashMap;
use std::sync::{Arc, Mutex};

use arc_swap::ArcSwapOption;
use tokio::sync::{broadcast, watch};
use tokio::task::JoinHandle;
use tokio::time::MissedTickBehavior;

use std::time::SystemTime;

use crate::error::{Error, Result};
use crate::event::{Event, EventKind};
use visual_cortex_capture::{CaptureError, Frame, FrameSource, Rate, Target};

const BROADCAST_CAPACITY: usize = 256;
const INITIAL_CAPTURE_RATE_HZ: f64 = 1.0;
/// Reserved watcher name carried by session-wide events (target lost/reattached).
const SESSION_WATCHER: &str = "$session";

/// A capture session: owns the frame source, the capture loop, and all
/// watcher tasks. Dropping the last `Session` handle stops everything.
#[derive(Clone, Debug)]
pub struct Session {
    pub(crate) inner: Arc<Inner>,
}

#[derive(Debug)]
pub(crate) struct Inner {
    pub(crate) slot: Arc<ArcSwapOption<Frame>>,
    pub(crate) events: broadcast::Sender<Event>,
    pub(crate) rate_tx: watch::Sender<Rate>,
    pub(crate) watchers: Mutex<HashMap<String, WatcherEntry>>,
    pub(crate) capture_task: Mutex<Option<JoinHandle<()>>>,
}

/// One registered watcher: its evaluation rate (for capture-rate math) and the
/// task running its loop. `task` is briefly `None` between name reservation and
/// spawn in `subscribe`.
#[derive(Debug)]
pub(crate) struct WatcherEntry {
    pub(crate) rate: Rate,
    pub(crate) task: Option<JoinHandle<()>>,
}

impl Drop for Inner {
    fn drop(&mut self) {
        if let Some(task) = self.capture_task.lock().unwrap().take() {
            task.abort();
        }
        for (_, mut entry) in self.watchers.lock().unwrap().drain() {
            if let Some(task) = entry.task.take() {
                task.abort();
            }
        }
    }
}

impl Session {
    pub fn builder() -> SessionBuilder {
        SessionBuilder {
            source: None,
            target: None,
            reattach: true,
        }
    }

    /// Begin configuring a watcher. Call `.subscribe()` on the returned
    /// builder to register it and get its event stream.
    pub fn watch(&self, name: &str) -> crate::watcher::WatcherBuilder<'_> {
        crate::watcher::WatcherBuilder::new(self, name)
    }

    /// Remove a watcher by name, aborting its task and recomputing the capture
    /// rate — which may now *decrease* if this was the fastest watcher.
    /// Returns whether a watcher with that name was registered.
    pub fn unsubscribe(&self, name: &str) -> bool {
        let entry = self.inner.watchers.lock().unwrap().remove(name);
        match entry {
            Some(mut entry) => {
                if let Some(task) = entry.task.take() {
                    task.abort();
                }
                recompute_capture_rate(&self.inner);
                true
            }
            None => false,
        }
    }
}

pub struct SessionBuilder {
    source: Option<Box<dyn FrameSource>>,
    target: Option<Target>,
    reattach: bool,
}

impl SessionBuilder {
    /// Provide the frame source directly (used by tests and `FakeSource`).
    /// Mutually exclusive with `target`.
    pub fn source(mut self, source: impl FrameSource) -> Self {
        self.source = Some(Box::new(source));
        self
    }

    /// Capture a real display or window. Mutually exclusive with `source`;
    /// macOS-only in this milestone.
    pub fn target(mut self, target: Target) -> Self {
        self.target = Some(target);
        self
    }

    /// If the capture target is lost, keep trying to reattach (default `true`).
    pub fn reattach(mut self, reattach: bool) -> Self {
        self.reattach = reattach;
        self
    }

    pub async fn build(self) -> Result<Session> {
        let source = self.resolve_source()?;
        let slot = Arc::new(ArcSwapOption::from(None));
        let (events, _) = broadcast::channel(BROADCAST_CAPACITY);
        let (rate_tx, rate_rx) = watch::channel(Rate::hz(INITIAL_CAPTURE_RATE_HZ));
        let capture_task =
            tokio::spawn(capture_loop(source, slot.clone(), rate_rx, events.clone()));
        Ok(Session {
            inner: Arc::new(Inner {
                slot,
                events,
                rate_tx,
                watchers: Mutex::new(HashMap::new()),
                capture_task: Mutex::new(Some(capture_task)),
            }),
        })
    }

    fn resolve_source(self) -> Result<Box<dyn FrameSource>> {
        match (self.source, self.target) {
            (Some(_), Some(_)) => Err(Error::AmbiguousSource),
            (Some(source), None) => Ok(source),
            (None, Some(target)) => build_target_source(&target, self.reattach),
            (None, None) => Err(Error::NoSource),
        }
    }
}

#[cfg(target_os = "macos")]
fn build_target_source(target: &Target, reattach: bool) -> Result<Box<dyn FrameSource>> {
    Ok(Box::new(visual_cortex_capture::ScapSource::new(
        target, reattach,
    )?))
}

#[cfg(not(target_os = "macos"))]
fn build_target_source(_target: &Target, _reattach: bool) -> Result<Box<dyn FrameSource>> {
    Err(Error::UnsupportedPlatform)
}

/// Set the capture rate to the fastest registered watcher rate, or the idle
/// default when no watchers remain. Called on every subscribe and unsubscribe.
pub(crate) fn recompute_capture_rate(inner: &Inner) {
    let fastest = {
        let watchers = inner.watchers.lock().unwrap();
        watchers
            .values()
            .map(|entry| entry.rate)
            // Fastest = smallest period; `Duration` is `Ord`.
            .min_by_key(|rate| rate.period())
            .unwrap_or_else(|| Rate::hz(INITIAL_CAPTURE_RATE_HZ))
    };
    // Avoid a needless notify (which rearms the capture interval) when unchanged.
    if fastest != *inner.rate_tx.borrow() {
        let _ = inner.rate_tx.send(fastest);
    }
}

/// Emit a session-wide event under the reserved sentinel name. `EventStream`
/// delivers these to every subscriber via `EventKind::is_session_level`.
fn send_session_event(events: &broadcast::Sender<Event>, kind: EventKind) {
    let _ = events.send(Event {
        watcher: Arc::from(SESSION_WATCHER),
        timestamp: SystemTime::now(),
        kind,
    });
}

/// Polls the source at the session capture rate and publishes each frame to the
/// latest-frame slot. Translates target-lost errors into session-wide
/// `TargetLost`/`TargetReattached` events (emitted once per transition).
async fn capture_loop(
    mut source: Box<dyn FrameSource>,
    slot: Arc<ArcSwapOption<Frame>>,
    mut rate_rx: watch::Receiver<Rate>,
    events: broadcast::Sender<Event>,
) {
    let mut interval = tokio::time::interval(rate_rx.borrow().period());
    interval.set_missed_tick_behavior(MissedTickBehavior::Skip);
    // Assume the target is present until a capture proves otherwise.
    let mut target_present = true;
    loop {
        tokio::select! {
            _ = interval.tick() => {
                match source.capture_frame().await {
                    Ok(frame) => {
                        if !target_present {
                            target_present = true;
                            send_session_event(&events, EventKind::TargetReattached);
                        }
                        slot.store(Some(frame));
                    }
                    Err(CaptureError::TargetLost) | Err(CaptureError::TargetNotFound(_)) => {
                        if target_present {
                            target_present = false;
                            send_session_event(&events, EventKind::TargetLost);
                        }
                    }
                    Err(e) => tracing::warn!("capture error: {e}"),
                }
            }
            changed = rate_rx.changed() => {
                if changed.is_err() {
                    break; // session dropped
                }
                // Rearming resets tick phase and fires one immediate tick — acceptable
                // for the capture loop; watcher schedules are independent intervals.
                interval = tokio::time::interval(rate_rx.borrow().period());
                interval.set_missed_tick_behavior(MissedTickBehavior::Skip);
            }
        }
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use visual_cortex_capture::FakeSource;

    #[tokio::test]
    async fn build_without_source_errors() {
        let err = Session::builder().build().await.unwrap_err();
        assert!(matches!(err, Error::NoSource));
    }

    #[tokio::test(start_paused = true)]
    async fn capture_loop_fills_the_frame_slot() {
        let session = Session::builder()
            .source(FakeSource::new([Frame::solid(4, 4, [1, 2, 3, 255])]))
            .build()
            .await
            .unwrap();
        // paused-time test: sleeping auto-advances the clock past capture ticks
        tokio::time::sleep(std::time::Duration::from_secs(3)).await;
        let frame = session
            .inner
            .slot
            .load_full()
            .expect("slot should hold a frame");
        assert_eq!(frame.data()[0], 1);
    }

    use visual_cortex_vision::Brightness;

    fn one_frame_session_source() -> FakeSource {
        FakeSource::new([Frame::solid(4, 4, [0, 0, 0, 255])])
    }

    #[tokio::test(start_paused = true)]
    async fn capture_rate_tracks_fastest_watcher_and_recovers_on_unsubscribe() {
        let session = Session::builder()
            .source(one_frame_session_source())
            .build()
            .await
            .unwrap();

        let _slow = session
            .watch("slow")
            .rate(Rate::hz(2.0))
            .detector(Brightness)
            .subscribe()
            .unwrap();
        assert_eq!(*session.inner.rate_tx.borrow(), Rate::hz(2.0));

        let _fast = session
            .watch("fast")
            .rate(Rate::hz(10.0))
            .detector(Brightness)
            .subscribe()
            .unwrap();
        assert_eq!(*session.inner.rate_tx.borrow(), Rate::hz(10.0));

        // Removing the fastest watcher lets the capture rate fall back.
        assert!(session.unsubscribe("fast"));
        assert_eq!(*session.inner.rate_tx.borrow(), Rate::hz(2.0));

        // Removing the last watcher returns to the idle default.
        assert!(session.unsubscribe("slow"));
        assert_eq!(
            *session.inner.rate_tx.borrow(),
            Rate::hz(INITIAL_CAPTURE_RATE_HZ)
        );
    }

    #[tokio::test]
    async fn build_rejects_source_and_target_together() {
        use visual_cortex_capture::Target;
        let err = Session::builder()
            .source(FakeSource::new([Frame::solid(2, 2, [0, 0, 0, 255])]))
            .target(Target::display(0))
            .build()
            .await
            .unwrap_err();
        assert!(matches!(err, Error::AmbiguousSource));
    }

    #[tokio::test(start_paused = true)]
    async fn unsubscribe_unknown_is_false_and_frees_the_name() {
        let session = Session::builder()
            .source(one_frame_session_source())
            .build()
            .await
            .unwrap();

        assert!(!session.unsubscribe("ghost"));

        let _w = session.watch("w").detector(Brightness).subscribe().unwrap();
        let dup = session.watch("w").detector(Brightness).subscribe();
        assert!(matches!(dup, Err(Error::DuplicateWatcher(n)) if n == "w"));

        // After unsubscribe the name is free to reuse.
        assert!(session.unsubscribe("w"));
        let _w2 = session.watch("w").detector(Brightness).subscribe().unwrap();
    }
}