visual_cortex/

watcher.rs

use std::collections::HashSet;
use std::sync::Arc;
use std::time::SystemTime;

use arc_swap::ArcSwapOption;
use tokio::sync::broadcast;
use tokio::time::MissedTickBehavior;

use crate::error::{Error, Result};
use crate::event::{Event, EventKind, ThresholdDirection};
use crate::session::{recompute_capture_rate, Session, WatcherEntry};
use crate::stream::EventStream;
use visual_cortex_capture::{Frame, FrameView, Rate, Region};
use visual_cortex_vision::{
    Detector, DetectorError, DetectorOutput, FrameDiff, OcrDetector, OcrEngine, TemplateMatcher,
};

const DEGRADED_THRESHOLD: u32 = 5;

pub(crate) enum Mode {
    Change,
    EveryMatch,
    Threshold(Box<dyn Fn(f64) -> bool + Send + 'static>),
    Template { threshold: f32 },
}

/// Configures one watcher. Created via `Session::watch`; consumed by `subscribe`.
pub struct WatcherBuilder<'s> {
    session: &'s Session,
    name: String,
    region: Region,
    rate: Rate,
    frame_diff: bool,
    gates: Vec<Box<dyn Detector>>,
    detector: Option<Box<dyn Detector>>,
    mode: Mode,
    template_error: Option<Error>,
}

impl<'s> WatcherBuilder<'s> {
    pub(crate) fn new(session: &'s Session, name: &str) -> Self {
        Self {
            session,
            name: name.to_string(),
            region: Region::Full,
            rate: Rate::hz(1.0),
            frame_diff: true,
            gates: Vec::new(),
            detector: None,
            mode: Mode::Change,
            template_error: None,
        }
    }

    pub fn region(mut self, region: Region) -> Self {
        self.region = region;
        self
    }

    pub fn rate(mut self, rate: Rate) -> Self {
        self.rate = rate;
        self
    }

    /// The implicit pixels-changed gate. On by default; disable for detectors
    /// that must run even when the region is static.
    pub fn frame_diff(mut self, enabled: bool) -> Self {
        self.frame_diff = enabled;
        self
    }

    /// Add a gate: a boolean detector that must report `Bool(true)` for the
    /// main detector to run this tick. Gates run in registration order.
    pub fn gate(mut self, gate: impl Detector) -> Self {
        self.gates.push(Box::new(gate));
        self
    }

    pub fn detector(mut self, detector: impl Detector) -> Self {
        self.detector = Some(Box::new(detector));
        self
    }

    /// Default mode: emit `Changed` when the detector output transitions.
    pub fn on_change(mut self) -> Self {
        self.mode = Mode::Change;
        self
    }

    /// Emit `Matched` on every tick where the detector matches.
    pub fn on_every_match(mut self) -> Self {
        self.mode = Mode::EveryMatch;
        self
    }

    /// Emit `ThresholdCrossed` when `predicate(numeric_output)` flips.
    pub fn on_threshold(mut self, predicate: impl Fn(f64) -> bool + Send + 'static) -> Self {
        self.mode = Mode::Threshold(Box::new(predicate));
        self
    }

    /// Watch for a PNG template. Emits `TemplateAppeared` when the best match
    /// score reaches `threshold` (0.0..=1.0) and `TemplateVanished` when it
    /// falls back below. Sets both the detector and the emission mode; like
    /// all builder configuration, the last call wins.
    pub fn template(mut self, png_bytes: &[u8], threshold: f32) -> Self {
        match TemplateMatcher::from_png_bytes(png_bytes) {
            Ok(matcher) => {
                self.detector = Some(Box::new(matcher));
                self.mode = Mode::Template { threshold };
            }
            Err(e) => {
                self.template_error =
                    Some(Error::InvalidTemplate(self.name.clone(), e.to_string()));
            }
        }
        self
    }

    /// OCR the region and emit text transitions (`Changed` with
    /// `DetectorOutput::Text`/`None`). "Text appeared" is the
    /// `None -> Text(..)` transition. Sets the detector only; the emission
    /// mode stays chainable.
    pub fn ocr_text(mut self, engine: impl OcrEngine) -> Self {
        self.detector = Some(Box::new(OcrDetector::text(engine)));
        self
    }

    /// OCR the region and parse a number out of the text (see
    /// [`patterns::number`](visual_cortex_vision::patterns::number)). Pair with
    /// `.on_threshold(..)` for HP-bar-style watchers.
    pub fn ocr(
        mut self,
        engine: impl OcrEngine,
        pattern: impl Fn(&str) -> Option<f64> + Send + 'static,
    ) -> Self {
        self.detector = Some(Box::new(OcrDetector::number(engine, pattern)));
        self
    }

    /// Register the watcher and return its event stream.
    pub fn subscribe(self) -> Result<EventStream> {
        if let Some(e) = self.template_error {
            return Err(e);
        }

        let detector = self
            .detector
            .ok_or_else(|| Error::MissingDetector(self.name.clone()))?;
        // Atomically reserve the name so duplicate watchers fail fast, storing
        // the rate now so a concurrent recompute already sees this watcher.
        {
            let mut watchers = self.session.inner.watchers.lock().unwrap();
            if watchers.contains_key(&self.name) {
                return Err(Error::DuplicateWatcher(self.name));
            }
            watchers.insert(
                self.name.clone(),
                WatcherEntry {
                    rate: self.rate,
                    task: None,
                },
            );
        }

        let rx = self.session.inner.events.subscribe();
        let runtime = WatcherRuntime {
            name: Arc::from(self.name.as_str()),
            region: self.region,
            rate: self.rate,
            frame_diff: self.frame_diff,
            gates: self.gates,
            detector: Some(detector),
            mode: self.mode,
        };
        let task = tokio::spawn(watcher_loop(
            runtime,
            self.session.inner.slot.clone(),
            self.session.inner.events.clone(),
        ));
        self.session
            .inner
            .watchers
            .lock()
            .unwrap()
            .get_mut(&self.name)
            .expect("entry reserved above")
            .task = Some(task);
        recompute_capture_rate(&self.session.inner);
        Ok(EventStream::new(rx, HashSet::from([self.name])))
    }
}

struct WatcherRuntime {
    name: Arc<str>,
    region: Region,
    rate: Rate,
    frame_diff: bool,
    gates: Vec<Box<dyn Detector>>,
    detector: Option<Box<dyn Detector>>,
    mode: Mode,
}

async fn watcher_loop(
    mut w: WatcherRuntime,
    slot: Arc<ArcSwapOption<Frame>>,
    events: broadcast::Sender<Event>,
) {
    let mut interval = tokio::time::interval(w.rate.period());
    interval.set_missed_tick_behavior(MissedTickBehavior::Skip);
    let mut diff_gate = w.frame_diff.then(FrameDiff::new);
    let mut last_output: Option<DetectorOutput> = None;
    let mut inside_threshold: Option<bool> = None;
    let mut consecutive_errors: u32 = 0;

    loop {
        interval.tick().await;
        let Some(frame) = slot.load_full() else {
            continue;
        };
        let rect = match w.region.resolve(frame.width(), frame.height()) {
            Ok(rect) => rect,
            Err(e) => {
                tracing::warn!(watcher = %w.name, "region does not fit frame: {e}");
                continue;
            }
        };
        let view = match frame.view(rect) {
            Ok(view) => view,
            Err(e) => {
                tracing::warn!(watcher = %w.name, "crop failed: {e}");
                continue;
            }
        };
        if let Some(gate) = &mut diff_gate {
            if matches!(gate.evaluate(&view), Ok(DetectorOutput::Bool(false))) {
                continue;
            }
        }
        if !gates_pass(&mut w.gates, &view, &w.name) {
            continue;
        }
        let result = if w.detector.as_ref().is_some_and(|d| d.is_heavy()) {
            // Move the detector and an owned frame handle onto the blocking
            // pool; CPU-bound inference must not stall the async scheduler.
            let mut detector = w.detector.take().expect("detector present");
            let frame = frame.clone();
            match tokio::task::spawn_blocking(move || {
                // AssertUnwindSafe: a caught panic may leave the detector
                // logically inconsistent; we accept that — subsequent
                // failures count toward WatcherDegraded.
                let out = std::panic::catch_unwind(std::panic::AssertUnwindSafe(|| {
                    let view = frame.view(rect).expect("rect was validated for this frame");
                    detector.evaluate(&view)
                }))
                .unwrap_or_else(|_| Err(DetectorError::Other("detector panicked".to_string())));
                (detector, out)
            })
            .await
            {
                Ok((detector, out)) => {
                    w.detector = Some(detector);
                    out
                }
                Err(join_err) => {
                    // Unreachable in practice: detector panics are caught above,
                    // and blocking tasks are not cancelled once started. If it
                    // does happen the detector is lost, so the watcher cannot
                    // continue; signal subscribers and end the loop. (The
                    // registry entry remains until unsubscribe/session drop.)
                    tracing::error!(watcher = %w.name, "blocking task failed: {join_err}");
                    send(
                        &events,
                        &w.name,
                        EventKind::WatcherDegraded {
                            consecutive_errors: consecutive_errors + 1,
                        },
                    );
                    return;
                }
            }
        } else {
            w.detector
                .as_mut()
                .expect("detector present")
                .evaluate(&view)
        };
        match result {
            Ok(output) => {
                consecutive_errors = 0;
                emit(
                    &w.mode,
                    &w.name,
                    &events,
                    output,
                    &mut last_output,
                    &mut inside_threshold,
                );
            }
            Err(e) => {
                consecutive_errors += 1;
                tracing::warn!(
                    watcher = %w.name,
                    "detector error ({consecutive_errors} consecutive): {e}"
                );
                if consecutive_errors == DEGRADED_THRESHOLD {
                    send(
                        &events,
                        &w.name,
                        EventKind::WatcherDegraded { consecutive_errors },
                    );
                }
            }
        }
    }
}

fn gates_pass(gates: &mut [Box<dyn Detector>], view: &FrameView<'_>, name: &Arc<str>) -> bool {
    for gate in gates.iter_mut() {
        match gate.evaluate(view) {
            Ok(DetectorOutput::Bool(true)) => {}
            Ok(_) => return false, // Bool(false) or any non-boolean output closes the gate
            Err(e) => {
                tracing::warn!(watcher = %name, "gate error, skipping tick: {e}");
                return false;
            }
        }
    }
    true
}

fn send(events: &broadcast::Sender<Event>, name: &Arc<str>, kind: EventKind) {
    let _ = events.send(Event {
        watcher: name.clone(),
        timestamp: SystemTime::now(),
        kind,
    });
}

fn emit(
    mode: &Mode,
    name: &Arc<str>,
    events: &broadcast::Sender<Event>,
    output: DetectorOutput,
    last_output: &mut Option<DetectorOutput>,
    inside_threshold: &mut Option<bool>,
) {
    match mode {
        Mode::EveryMatch => {
            let matched = !matches!(output, DetectorOutput::None | DetectorOutput::Bool(false));
            if matched {
                send(events, name, EventKind::Matched { output });
            }
        }
        Mode::Change => {
            if last_output.as_ref() != Some(&output) {
                send(
                    events,
                    name,
                    EventKind::Changed {
                        old: last_output.clone(),
                        new: output.clone(),
                    },
                );
            }
            *last_output = Some(output);
        }
        Mode::Threshold(predicate) => {
            let Some(value) = output.as_number() else {
                return;
            };
            let inside = predicate(value);
            match *inside_threshold {
                // First observation: fire only if we start inside the threshold —
                // "HP is already low" is worth knowing; "HP is fine" is not.
                None if inside => send(
                    events,
                    name,
                    EventKind::ThresholdCrossed {
                        value,
                        direction: ThresholdDirection::Entered,
                    },
                ),
                Some(prev) if prev != inside => {
                    let direction = if inside {
                        ThresholdDirection::Entered
                    } else {
                        ThresholdDirection::Exited
                    };
                    send(
                        events,
                        name,
                        EventKind::ThresholdCrossed { value, direction },
                    );
                }
                _ => {}
            }
            *inside_threshold = Some(inside);
        }
        Mode::Template { threshold } => {
            let Some(score) = output.as_number() else {
                return;
            };
            let present = score >= *threshold as f64;
            match *inside_threshold {
                // First observation: only an already-visible template is news.
                None if present => send(
                    events,
                    name,
                    EventKind::TemplateAppeared {
                        score: score as f32,
                    },
                ),
                Some(prev) if prev != present => {
                    if present {
                        send(
                            events,
                            name,
                            EventKind::TemplateAppeared {
                                score: score as f32,
                            },
                        );
                    } else {
                        send(events, name, EventKind::TemplateVanished);
                    }
                }
                _ => {}
            }
            *inside_threshold = Some(present);
        }
    }
}