wayvr/wlx-overlay-s/src/overlays/screen/backend.rs

use std::{
    sync::{Arc, LazyLock, atomic::AtomicU64},
    time::Instant,
};

use glam::{Affine2, Vec2, vec2};
use wlx_capture::{WlxCapture, frame::Transform};

use crate::{
    backend::{
        XrBackend,
        input::{HoverResult, PointerHit, PointerMode},
    },
    graphics::ExtentExt,
    state::AppState,
    subsystem::hid::{MOUSE_LEFT, MOUSE_MIDDLE, MOUSE_RIGHT, WheelDelta},
    windowing::backend::{
        FrameMeta, OverlayBackend, OverlayEventData, RenderResources, ShouldRender, ui_transform,
    },
};
use wlx_common::overlays::{BackendAttrib, BackendAttribValue, MouseTransform, StereoMode};

use super::capture::{ScreenPipeline, WlxCaptureIn, WlxCaptureOut, receive_callback};

const CURSOR_SIZE: f32 = 16. / 1440.;

static START: LazyLock<Instant> = LazyLock::new(Instant::now);
static NEXT_MOVE: AtomicU64 = AtomicU64::new(0);

fn can_move() -> bool {
    START.elapsed().as_millis() as u64 > NEXT_MOVE.load(std::sync::atomic::Ordering::Relaxed)
}

fn set_next_move(millis_from_now: u64) {
    NEXT_MOVE.store(
        START.elapsed().as_millis() as u64 + millis_from_now,
        std::sync::atomic::Ordering::Relaxed,
    );
}

pub struct ScreenBackend {
    name: Arc<str>,
    capture: Box<dyn WlxCapture<WlxCaptureIn, WlxCaptureOut>>,
    pipeline: Option<ScreenPipeline>,
    cur_frame: Option<WlxCaptureOut>,
    meta: Option<FrameMeta>,
    mouse_transform: Affine2,
    interaction_transform: Option<Affine2>,
    stereo: Option<StereoMode>,
    pub(super) logical_pos: Vec2,
    pub(super) logical_size: Vec2,
    pub(super) mouse_transform_original: Transform,
    mouse_transform_override: MouseTransform,
    just_resumed: bool,
}

impl ScreenBackend {
    pub(super) fn new_raw(
        name: Arc<str>,
        xr_backend: XrBackend,
        capture: Box<dyn WlxCapture<WlxCaptureIn, WlxCaptureOut>>,
    ) -> Self {
        Self {
            name,
            capture,
            pipeline: None,
            cur_frame: None,
            meta: None,
            mouse_transform: Affine2::ZERO,
            interaction_transform: None,
            stereo: if matches!(xr_backend, XrBackend::OpenXR) {
                Some(StereoMode::None)
            } else {
                None
            },
            logical_pos: Vec2::ZERO,
            logical_size: Vec2::ZERO,
            mouse_transform_original: Transform::Undefined,
            mouse_transform_override: MouseTransform::Default,
            just_resumed: false,
        }
    }

    pub(super) fn apply_mouse_transform_with_override(&mut self, override_transform: Transform) {
        let size = self.logical_size;
        let pos = self.logical_pos;

        let transform = match override_transform {
            Transform::Undefined => self.mouse_transform_original,
            other => other,
        };

        self.mouse_transform = match transform {
            Transform::Normal | Transform::Undefined => {
                Affine2::from_cols(vec2(size.x, 0.), vec2(0., size.y), pos)
            }
            Transform::Rotated90 => Affine2::from_cols(
                vec2(0., size.y),
                vec2(-size.x, 0.),
                vec2(pos.x + size.x, pos.y),
            ),
            Transform::Rotated180 => Affine2::from_cols(
                vec2(-size.x, 0.),
                vec2(0., -size.y),
                vec2(pos.x + size.x, pos.y + size.y),
            ),
            Transform::Rotated270 => Affine2::from_cols(
                vec2(0., -size.y),
                vec2(size.x, 0.),
                vec2(pos.x, pos.y + size.y),
            ),
            Transform::Flipped => Affine2::from_cols(
                vec2(-size.x, 0.),
                vec2(0., size.y),
                vec2(pos.x + size.x, pos.y),
            ),
            Transform::Flipped90 => {
                Affine2::from_cols(vec2(0., size.y), vec2(size.x, 0.), vec2(pos.x, pos.y))
            }
            Transform::Flipped180 => Affine2::from_cols(
                vec2(size.x, 0.),
                vec2(0., -size.y),
                vec2(pos.x, pos.y + size.y),
            ),
            Transform::Flipped270 => Affine2::from_cols(
                vec2(0., -size.y),
                vec2(-size.x, 0.),
                vec2(pos.x + size.x, pos.y + size.y),
            ),
        };
    }
}

impl OverlayBackend for ScreenBackend {
    fn init(&mut self, _app: &mut AppState) -> anyhow::Result<()> {
        Ok(())
    }
    fn should_render(&mut self, app: &mut AppState) -> anyhow::Result<ShouldRender> {
        if !self.capture.is_ready() {
            let supports_dmabuf = app
                .gfx
                .device
                .enabled_extensions()
                .ext_external_memory_dma_buf
                && self.capture.supports_dmbuf();

            let allow_dmabuf = &*app.session.config.capture_method != "pw_fallback"
                && &*app.session.config.capture_method != "screencopy";

            let capture_method = app.session.config.capture_method.clone();

            let dmabuf_formats = if !supports_dmabuf {
                log::info!("Capture method does not support DMA-buf");
                if app.gfx_extras.queue_capture.is_none() {
                    log::warn!(
                        "Current GPU does not support multiple queues. Software capture will take place on the main thread. Expect degraded performance."
                    );
                }
                &Vec::new()
            } else if !allow_dmabuf {
                log::info!("Not using DMA-buf capture due to {capture_method}");
                if app.gfx_extras.queue_capture.is_none() {
                    log::warn!(
                        "Current GPU does not support multiple queues. Software capture will take place on the main thread. Expect degraded performance."
                    );
                }
                &Vec::new()
            } else {
                log::warn!(
                    "Using DMA-buf capture. If screens are blank for you, switch to SHM using:"
                );
                log::warn!(
                    "echo 'capture_method: pw_fallback' > ~/.config/wlxoverlay/conf.d/pw_fallback.yaml"
                );

                &app.gfx_extras.drm_formats
            };

            let user_data = WlxCaptureIn::new(self.name.clone(), app);
            self.capture
                .init(dmabuf_formats, user_data, receive_callback);
            self.capture.request_new_frame();
            return Ok(ShouldRender::Unable);
        }

        if let Some(frame) = self.capture.receive() {
            let mut meta = frame.get_frame_meta(&app.session.config);

            if let Some(pipeline) = self.pipeline.as_mut() {
                meta.extent[2] = pipeline.get_depth();
                if self
                    .meta
                    .is_some_and(|old| old.extent[..2] != meta.extent[..2])
                {
                    pipeline.set_extent(
                        app,
                        [meta.extent[0] as _, meta.extent[1] as _],
                        [0., 0.],
                    )?;
                    self.interaction_transform = Some(ui_transform(meta.extent.extent_u32arr()));
                }
            } else {
                let pipeline = ScreenPipeline::new(
                    &meta,
                    app,
                    self.stereo.unwrap_or(StereoMode::None),
                    [0., 0.],
                )?;
                meta.extent[2] = pipeline.get_depth();
                self.pipeline = Some(pipeline);
                self.interaction_transform = Some(ui_transform(meta.extent.extent_u32arr()));
            }

            self.meta = Some(meta);
            self.cur_frame = Some(frame);

            Ok(ShouldRender::Should)
        } else if self.cur_frame.is_some() {
            if self.just_resumed {
                self.just_resumed = false;
                Ok(ShouldRender::Should)
            } else {
                Ok(ShouldRender::Can)
            }
        } else {
            log::trace!("{}: backend ready, but no image received.", self.name);
            Ok(ShouldRender::Unable)
        }
    }
    fn render(&mut self, app: &mut AppState, rdr: &mut RenderResources) -> anyhow::Result<()> {
        // want panic; must be some if should_render was not Unable
        let capture = self.cur_frame.as_ref().unwrap();
        let image = capture.image.clone();

        // want panic; must be Some if cur_frame is also Some
        self.pipeline
            .as_mut()
            .unwrap()
            .render(image, capture.mouse.as_ref(), app, rdr)?;
        self.capture.request_new_frame();
        Ok(())
    }
    fn pause(&mut self, _app: &mut AppState) -> anyhow::Result<()> {
        self.capture.pause();
        Ok(())
    }
    fn resume(&mut self, _app: &mut AppState) -> anyhow::Result<()> {
        self.just_resumed = true;
        self.capture.resume();
        Ok(())
    }
    fn frame_meta(&mut self) -> Option<FrameMeta> {
        self.meta
    }

    fn notify(&mut self, _app: &mut AppState, _event_data: OverlayEventData) -> anyhow::Result<()> {
        Ok(())
    }

    fn on_hover(&mut self, app: &mut AppState, hit: &PointerHit) -> HoverResult {
        #[cfg(debug_assertions)]
        log::trace!("Hover: {:?}", hit.uv);
        if can_move()
            && (!app.session.config.focus_follows_mouse_mode
                || app.input_state.pointers[hit.pointer].now.move_mouse)
        {
            let pos = self.mouse_transform.transform_point2(hit.uv);
            app.hid_provider.inner.mouse_move(pos);
            set_next_move(u64::from(app.session.config.mouse_move_interval_ms));
        }
        HoverResult {
            consume: true,
            ..HoverResult::default()
        }
    }
    fn on_pointer(&mut self, app: &mut AppState, hit: &PointerHit, pressed: bool) {
        let mut btn = match hit.mode {
            PointerMode::Right => MOUSE_RIGHT,
            PointerMode::Middle => MOUSE_MIDDLE,
            _ => MOUSE_LEFT,
        };

        // Swap left and right buttons if left-handed mode is enabled
        if app.session.config.left_handed_mouse {
            btn = match btn {
                MOUSE_LEFT => MOUSE_RIGHT,
                MOUSE_RIGHT => MOUSE_LEFT,
                other => other,
            };
        }

        if pressed {
            set_next_move(u64::from(app.session.config.click_freeze_time_ms));
        }

        app.hid_provider.inner.send_button(btn, pressed);

        if !pressed {
            return;
        }
        let pos = self.mouse_transform.transform_point2(hit.uv);
        app.hid_provider.inner.mouse_move(pos);
    }

    fn on_scroll(&mut self, app: &mut AppState, _hit: &PointerHit, delta: WheelDelta) {
        app.hid_provider.inner.wheel(delta);
    }

    fn on_left(&mut self, _app: &mut AppState, _hand: usize) {}

    fn get_interaction_transform(&mut self) -> Option<Affine2> {
        self.interaction_transform
    }
    #[allow(unreachable_patterns)]
    fn get_attrib(&self, attrib: BackendAttrib) -> Option<BackendAttribValue> {
        match attrib {
            BackendAttrib::Stereo => self.stereo.map(BackendAttribValue::Stereo),
            BackendAttrib::MouseTransform => Some(BackendAttribValue::MouseTransform(
                self.mouse_transform_override,
            )),
            _ => None,
        }
    }
    #[allow(unreachable_patterns)]
    fn set_attrib(&mut self, app: &mut AppState, value: BackendAttribValue) -> bool {
        match value {
            BackendAttribValue::Stereo(new) => {
                if let Some(stereo) = self.stereo.as_mut() {
                    log::debug!("{}: stereo: {stereo:?} → {new:?}", self.name);
                    *stereo = new;
                    if let Some(pipeline) = self.pipeline.as_mut() {
                        pipeline.set_stereo(app, new).unwrap(); // only panics if gfx is dead
                    }
                    true
                } else {
                    false
                }
            }
            BackendAttribValue::MouseTransform(new) => {
                self.mouse_transform_override = new;
                let frame_transform = match new {
                    MouseTransform::Default => Transform::Undefined,
                    MouseTransform::Normal => Transform::Normal,
                    MouseTransform::Rotated90 => Transform::Rotated90,
                    MouseTransform::Rotated180 => Transform::Rotated180,
                    MouseTransform::Rotated270 => Transform::Rotated270,
                    MouseTransform::Flipped => Transform::Flipped,
                    MouseTransform::Flipped90 => Transform::Flipped90,
                    MouseTransform::Flipped180 => Transform::Flipped180,
                    MouseTransform::Flipped270 => Transform::Flipped270,
                };
                self.apply_mouse_transform_with_override(frame_transform);
                true
            }
            _ => false,
        }
    }
}