visual-cortex-vision 0.10.0

//! Debug taps for preprocessors: visual insight into what a masking stage
//! is doing, without ever stalling the watcher. Implementations push frames
//! through a bounded channel to a writer thread and DROP on lag (the
//! Recorder discipline); dropped counts are logged and queryable.

use std::path::PathBuf;
use std::sync::atomic::{AtomicU64, Ordering};
use std::sync::mpsc::{sync_channel, SyncSender, TrySendError};
use std::sync::Arc;
use std::thread::JoinHandle;

use visual_cortex_capture::Frame;

/// A tap receiving per-tick stage frames from a preprocessor (see
/// [`StabilityMask::debug_sink`](crate::StabilityMask::debug_sink)).
///
/// `write` is called on the watcher's tick path — implementations must
/// return immediately (hand off to a thread, drop on lag), never block.
pub trait DebugSink: Send + 'static {
    fn write(&mut self, tick: u64, stage: DebugStage, frame: &Frame);

    /// Receives the watcher name at subscribe time (used e.g. for output
    /// subdirectories). Default: ignore.
    fn set_label(&mut self, _label: &str) {}
}

/// Which pipeline artifact a debug frame shows.
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
#[non_exhaustive]
pub enum DebugStage {
    /// The raw cropped input.
    Input,
    /// The per-block baseline estimate rendered as grayscale.
    Baseline,
    /// The human view: input with suppressed blocks dimmed.
    Overlay,
    /// False-color per-block state (idle/kept/holdover/re-seed) — shows WHY
    /// each block did what it did.
    State,
    /// Exactly what the rest of the pipeline sees.
    Output,
}

impl DebugStage {
    fn suffix(self) -> &'static str {
        match self {
            DebugStage::Input => "input",
            DebugStage::Baseline => "baseline",
            DebugStage::Overlay => "overlay",
            DebugStage::State => "state",
            DebugStage::Output => "output",
        }
    }
}

struct Job {
    tick: u64,
    stage: DebugStage,
    label: String,
    width: u32,
    height: u32,
    /// BGRA, cloned off the tick path.
    data: Vec<u8>,
}

/// Push a job without ever blocking: a full channel means the writer is
/// lagging, and debug output loses to the watcher — drop and count.
fn try_dispatch(tx: &SyncSender<Job>, dropped: &AtomicU64, job: Job) {
    match tx.try_send(job) {
        Ok(()) => {}
        Err(TrySendError::Full(_)) => {
            let n = dropped.fetch_add(1, Ordering::Relaxed) + 1;
            tracing::debug!(dropped = n, "PngDump lagging; dropping debug frame");
        }
        Err(TrySendError::Disconnected(_)) => {}
    }
}

/// PNG stills per sampled tick:
/// `<dir>/<watcher>/NNNNNN-{input,baseline,overlay,output}.png`.
///
/// Encoding and disk I/O happen on a dedicated writer thread behind a
/// bounded channel; frames are dropped (and counted) if the writer lags.
pub struct PngDump {
    tx: Option<SyncSender<Job>>,
    writer: Option<JoinHandle<()>>,
    dropped: Arc<AtomicU64>,
    sample_every: u64,
    label: String,
}

impl PngDump {
    /// Write stills under `dir` (created on demand).
    pub fn dir(dir: impl Into<PathBuf>) -> Self {
        let dir = dir.into();
        let (tx, rx) = sync_channel::<Job>(8);
        let writer = std::thread::spawn(move || {
            for job in rx {
                let sub = dir.join(&job.label);
                if let Err(e) = std::fs::create_dir_all(&sub) {
                    tracing::warn!(dir = %sub.display(), error = %e, "PngDump mkdir failed");
                    continue;
                }
                let path = sub.join(format!("{:06}-{}.png", job.tick, job.stage.suffix()));
                // BGRA -> RGBA for the PNG encoder.
                let mut rgba = job.data;
                for px in rgba.chunks_exact_mut(4) {
                    px.swap(0, 2);
                }
                match image::RgbaImage::from_raw(job.width, job.height, rgba) {
                    Some(img) => {
                        if let Err(e) = img.save(&path) {
                            tracing::warn!(path = %path.display(), error = %e, "PngDump write failed");
                        }
                    }
                    None => tracing::warn!("PngDump: frame buffer size mismatch"),
                }
            }
        });
        Self {
            tx: Some(tx),
            writer: Some(writer),
            dropped: Arc::default(),
            sample_every: 1,
            label: "preprocessor".into(),
        }
    }

    /// Only write every Nth tick (default 1 = every tick).
    pub fn sample_every(mut self, n: u64) -> Self {
        self.sample_every = n.max(1);
        self
    }

    /// Frames dropped because the writer thread lagged.
    pub fn dropped(&self) -> u64 {
        self.dropped.load(Ordering::Relaxed)
    }
}

impl DebugSink for PngDump {
    fn write(&mut self, tick: u64, stage: DebugStage, frame: &Frame) {
        if !tick.is_multiple_of(self.sample_every) {
            return;
        }
        let Some(tx) = &self.tx else { return };
        try_dispatch(
            tx,
            &self.dropped,
            Job {
                tick,
                stage,
                label: self.label.clone(),
                width: frame.width(),
                height: frame.height(),
                data: frame.data().to_vec(),
            },
        );
    }

    fn set_label(&mut self, label: &str) {
        self.label = label.to_string();
    }
}

impl Drop for PngDump {
    fn drop(&mut self) {
        drop(self.tx.take()); // close the channel so the writer drains and exits
        if let Some(w) = self.writer.take() {
            let _ = w.join();
        }
    }
}

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

    #[test]
    fn full_channel_drops_and_counts_instead_of_blocking() {
        let (tx, _rx) = sync_channel::<Job>(2); // never drained
        let dropped = AtomicU64::new(0);
        for tick in 0..5 {
            try_dispatch(
                &tx,
                &dropped,
                Job {
                    tick,
                    stage: DebugStage::Output,
                    label: "w".into(),
                    width: 1,
                    height: 1,
                    data: vec![0, 0, 0, 255],
                },
            );
        }
        assert_eq!(dropped.load(Ordering::Relaxed), 3, "2 buffered, 3 dropped");
    }

    #[test]
    fn writes_sampled_stills_under_the_watcher_label() {
        let tmp = tempfile::tempdir().unwrap();
        let mut dump = PngDump::dir(tmp.path()).sample_every(2);
        dump.set_label("tooltip");

        let frame = Frame::solid(2, 2, [10, 20, 30, 255]); // BGRA
        for tick in 0..3 {
            dump.write(tick, DebugStage::Output, &frame);
        }
        drop(dump); // joins the writer: all accepted jobs are on disk

        let dir = tmp.path().join("tooltip");
        assert!(dir.join("000000-output.png").exists());
        assert!(
            !dir.join("000001-output.png").exists(),
            "sample_every(2) must skip odd ticks"
        );
        assert!(dir.join("000002-output.png").exists());

        // BGRA -> RGBA swizzle correctness.
        let img = image::open(dir.join("000000-output.png"))
            .unwrap()
            .to_rgba8();
        assert_eq!(img.get_pixel(0, 0).0, [30, 20, 10, 255]);
    }
}