visual-cortex-capture 0.5.0

use std::sync::Arc;

use scap::capturer::{Capturer, Options};
use scap::frame::{Frame as ScapFrame, FrameType};
use scap::Target as ScapTarget;

use crate::error::CaptureError;
use crate::frame::Frame;
use crate::source::FrameSource;
use crate::target::{select_window, Target, WindowSelector};

/// Frames-per-second requested from `scap`. The session throttles evaluation
/// independently; this only bounds how fresh the backing capture can be.
const CAPTURE_FPS: u32 = 60;

/// Real screen capture on macOS via the `scap` crate (ScreenCaptureKit),
/// behind our `FrameSource` trait.
pub struct ScapSource {
    capturer: Capturer,
    target: Target,
    reattach: bool,
}

impl ScapSource {
    /// Resolve `target`, check the Screen Recording permission, and start a
    /// capturer. If `reattach` is set, a lost target is re-resolved on the fly.
    pub fn new(target: &Target, reattach: bool) -> Result<Self, CaptureError> {
        if !scap::is_supported() {
            return Err(CaptureError::Backend("scap: platform not supported".into()));
        }
        if !scap::has_permission() && !scap::request_permission() {
            return Err(CaptureError::PermissionDenied(
                "grant Screen Recording in System Settings › Privacy & Security".into(),
            ));
        }
        let capturer = build_capturer(target)?;
        Ok(Self {
            capturer,
            target: target.clone(),
            reattach,
        })
    }
}

/// One capturable thing, as reported by the OS.
#[derive(Debug, Clone, PartialEq, Eq)]
pub enum TargetInfo {
    /// `index` matches [`Target::display`] resolution order.
    Display { index: usize, title: String },
    /// `id` matches [`Target::window_id`].
    Window { id: u32, title: String },
}

/// Enumerate capturable displays and windows. Display indices align with the
/// order [`Target::display`] resolves them; window ids feed
/// [`Target::window_id`]. Requires the same Screen Recording permission as
/// capturing.
pub fn list_targets() -> Vec<TargetInfo> {
    let mut out = Vec::new();
    let mut display_index = 0;
    for t in scap::get_all_targets() {
        match t {
            ScapTarget::Display(d) => {
                out.push(TargetInfo::Display {
                    index: display_index,
                    title: d.title,
                });
                display_index += 1;
            }
            ScapTarget::Window(w) => out.push(TargetInfo::Window {
                id: w.id,
                title: w.title,
            }),
        }
    }
    out
}

/// Map our `Target` onto a live `scap::Target` — `TargetNotFound` when
/// nothing matches, `AmbiguousTarget` when a title selector matches several
/// windows (the pure matching rules live in `target::select_window`).
fn resolve_scap_target(target: &Target) -> Result<ScapTarget, CaptureError> {
    let all = scap::get_all_targets();
    let selector = match target {
        Target::Display(index) => {
            return all
                .into_iter()
                .filter(|t| matches!(t, ScapTarget::Display(_)))
                .nth(*index)
                .ok_or_else(|| CaptureError::TargetNotFound(format!("display index {index}")));
        }
        Target::WindowTitled(title) => WindowSelector::Exact(title),
        Target::WindowContaining(sub) => WindowSelector::Containing(sub),
        Target::WindowId(id) => WindowSelector::Id(*id),
    };
    let windows: Vec<(u32, String)> = all
        .iter()
        .filter_map(|t| match t {
            ScapTarget::Window(w) => Some((w.id, w.title.clone())),
            _ => None,
        })
        .collect();
    let id = select_window(&windows, &selector)?;
    all.into_iter()
        .find(|t| matches!(t, ScapTarget::Window(w) if w.id == id))
        .ok_or_else(|| CaptureError::TargetNotFound(format!("window id {id}")))
}

/// Resolve the target and start a fresh BGRA capturer for it.
fn build_capturer(target: &Target) -> Result<Capturer, CaptureError> {
    let scap_target = resolve_scap_target(target)?;
    let options = Options {
        fps: CAPTURE_FPS,
        target: Some(scap_target),
        output_type: FrameType::BGRAFrame,
        show_cursor: true,
        ..Default::default()
    };
    // Defense in depth: the native capture stack has panicked on invalid
    // configurations before (upstream scap 0.0.8); convert any residual
    // panic into an error instead of unwinding through the session.
    std::panic::catch_unwind(std::panic::AssertUnwindSafe(
        move || -> Result<Capturer, CaptureError> {
            let mut capturer = Capturer::build(options)
                .map_err(|e| CaptureError::Backend(format!("scap build: {e:?}")))?;
            capturer
                .start_capture()
                .map_err(|e| CaptureError::Backend(format!("scap start: {e}")))?;
            Ok(capturer)
        },
    ))
    .unwrap_or_else(|_| {
        Err(CaptureError::Backend(
            "capture backend panicked while starting".into(),
        ))
    })
}

/// Convert a `scap` BGRA frame into our BGRA8 `Frame`. Pure and platform-free
/// so it is unit-tested; `Frame::new` validates the byte length.
fn video_bgra_to_frame(width: i32, height: i32, data: Vec<u8>) -> Result<Arc<Frame>, CaptureError> {
    if width <= 0 || height <= 0 {
        return Err(CaptureError::Backend(format!(
            "non-positive frame dimensions {width}x{height}"
        )));
    }
    Frame::new(width as u32, height as u32, data).map(Arc::new)
}

#[async_trait::async_trait]
impl FrameSource for ScapSource {
    async fn capture_frame(&mut self) -> Result<Arc<Frame>, CaptureError> {
        // scap 0.0.8 uses a flat `Frame` enum (no Audio/Video split); a BGRA
        // capturer yields `Frame::BGRA`.
        match self.capturer.get_next_frame() {
            Ok(ScapFrame::BGRA(f)) => video_bgra_to_frame(f.width, f.height, f.data),
            Ok(_) => Err(CaptureError::Backend(
                "scap returned a non-BGRA frame; expected FrameType::BGRAFrame".into(),
            )),
            Err(_) => {
                // The capture channel closed: the target is gone. Optionally try
                // to rebuild the capturer so the next tick can reattach.
                if self.reattach {
                    if let Ok(capturer) = build_capturer(&self.target) {
                        self.capturer = capturer;
                    }
                }
                Err(CaptureError::TargetLost)
            }
        }
    }
}

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

    #[test]
    fn converts_bgra_dimensions_and_rejects_bad_input() {
        let good = video_bgra_to_frame(2, 2, vec![9u8; 16]).unwrap();
        assert_eq!((good.width(), good.height()), (2, 2));
        // Non-positive dimensions and length/stride mismatches must error, not panic.
        assert!(video_bgra_to_frame(-1, 2, vec![]).is_err());
        assert!(video_bgra_to_frame(2, 2, vec![0u8; 15]).is_err());
    }
}