Layers/src/wayland_shell.rs

use crate::ffi::ViewportHandle;
use raw_window_handle::{
    RawDisplayHandle, RawWindowHandle, WaylandDisplayHandle, WaylandWindowHandle,
};
use smithay_client_toolkit::{
    compositor::{CompositorHandler, CompositorState},
    delegate_compositor, delegate_keyboard, delegate_layer, delegate_output, delegate_pointer,
    delegate_registry, delegate_seat,
    output::{OutputHandler, OutputState},
    registry::{ProvidesRegistryState, RegistryState},
    registry_handlers,
    seat::{
        keyboard::{KeyEvent, KeyboardHandler, Keysym, Modifiers, RawModifiers},
        pointer::{PointerEvent, PointerEventKind, PointerHandler},
        Capability, SeatHandler, SeatState,
    },
    shell::{
        wlr_layer::{
            Anchor, KeyboardInteractivity, Layer, LayerShell, LayerShellHandler, LayerSurface,
            LayerSurfaceConfigure,
        },
        WaylandSurface,
    },
};
use std::ptr::NonNull;
use wayland_client::{
    globals::{registry_queue_init, GlobalList},
    protocol::{wl_keyboard, wl_output, wl_pointer, wl_seat, wl_surface},
    Connection, Proxy, QueueHandle,
};

const DEFAULT_LOGICAL_SIZE: (u32, u32) = (480, 640);
const MIN_LOGICAL_SIZE: (u32, u32) = (380, 220);
const ANCHOR_MARGIN: i32 = 24;

/// Run the Wayland (sctk + wlr-layer-shell) front end. Returns `true` if it took
/// over the session, `false` if init failed (wayland socket unreachable, no
/// layer-shell, etc.) so the caller can fall through to the winit path. Hard
/// exiting here would silently kill the plugin process inside KiCad — which the
/// user sees as "didn't load."
pub fn try_run() -> bool {
    tracing::info!("wayland: try_run entered");

    let conn = match probe_connection() {
        Some(c) => c,
        None => {
            tracing::warn!(
                "wayland: no reachable compositor socket; falling through to winit"
            );
            return false;
        }
    };

    let (globals, mut event_queue) = match registry_queue_init::<State>(&conn) {
        Ok(pair) => {
            tracing::info!("wayland: registry_queue_init OK");
            pair
        }
        Err(e) => {
            tracing::warn!("wayland: registry init failed ({e}); falling through to winit");
            return false;
        }
    };
    let qh = event_queue.handle();

    log_globals(&globals);

    let compositor = match CompositorState::bind(&globals, &qh) {
        Ok(c) => {
            tracing::info!("wayland: bound wl_compositor");
            c
        }
        Err(e) => {
            tracing::warn!("wayland: wl_compositor missing ({e}); falling through to winit");
            return false;
        }
    };

    let layer_shell = match LayerShell::bind(&globals, &qh) {
        Ok(l) => {
            tracing::info!("wayland: bound zwlr_layer_shell_v1");
            l
        }
        Err(e) => {
            tracing::warn!(
                "wayland: zwlr_layer_shell_v1 not advertised ({e}); falling through to winit"
            );
            return false;
        }
    };

    let surface = compositor.create_surface(&qh);
    let layer = layer_shell.create_layer_surface(
        &qh,
        surface,
        Layer::Top,
        Some("layers"),
        None,
    );
    layer.set_anchor(Anchor::TOP | Anchor::RIGHT);
    layer.set_size(DEFAULT_LOGICAL_SIZE.0, DEFAULT_LOGICAL_SIZE.1);
    layer.set_margin(ANCHOR_MARGIN, ANCHOR_MARGIN, 0, 0);
    layer.set_keyboard_interactivity(KeyboardInteractivity::OnDemand);
    layer.commit();
    tracing::info!(
        anchor = "TOP|RIGHT",
        size = ?(DEFAULT_LOGICAL_SIZE.0, DEFAULT_LOGICAL_SIZE.1),
        margin_top = ANCHOR_MARGIN,
        margin_right = ANCHOR_MARGIN,
        "wayland: layer surface committed (waiting for first configure)"
    );

    let raw_display = RawDisplayHandle::Wayland(WaylandDisplayHandle::new(
        match NonNull::new(conn.backend().display_ptr() as *mut _) {
            Some(p) => p,
            None => {
                tracing::warn!("wayland: null display ptr; falling through to winit");
                return false;
            }
        },
    ));
    let raw_window = RawWindowHandle::Wayland(WaylandWindowHandle::new(
        match NonNull::new(layer.wl_surface().id().as_ptr() as *mut _) {
            Some(p) => p,
            None => {
                tracing::warn!("wayland: null surface ptr; falling through to winit");
                return false;
            }
        },
    ));

    let mut state = State {
        registry_state: RegistryState::new(&globals),
        seat_state: SeatState::new(&globals, &qh),
        output_state: OutputState::new(&globals, &qh),
        viewport: None,
        layer,
        keyboard: None,
        pointer: None,
        modifiers: Modifiers::default(),
        last_pointer: (0.0, 0.0),
        width: DEFAULT_LOGICAL_SIZE.0,
        height: DEFAULT_LOGICAL_SIZE.1,
        scale: 1,
        first_configure: true,
        raw_display,
        raw_window,
        exit: false,
    };

    tracing::info!("wayland: layer-shell up, entering dispatch loop");
    while !state.exit {
        if let Err(e) = event_queue.blocking_dispatch(&mut state) {
            tracing::error!("wayland dispatch: {e}");
            break;
        }
    }
    true
}

struct State {
    registry_state: RegistryState,
    seat_state: SeatState,
    output_state: OutputState,
    /// Held above `layer` so the wgpu surface drops before the wl_surface.
    viewport: Option<ViewportHandle>,
    layer: LayerSurface,
    keyboard: Option<wl_keyboard::WlKeyboard>,
    pointer: Option<wl_pointer::WlPointer>,
    modifiers: Modifiers,
    last_pointer: (f32, f32),
    width: u32,
    height: u32,
    scale: i32,
    first_configure: bool,
    raw_display: RawDisplayHandle,
    raw_window: RawWindowHandle,
    exit: bool,
}

impl State {
    fn ensure_viewport(&mut self) {
        if self.viewport.is_some() {
            return;
        }
        let scale = self.scale.max(1) as f32;
        let logical_w = self.width.max(MIN_LOGICAL_SIZE.0) as f32;
        let logical_h = self.height.max(MIN_LOGICAL_SIZE.1) as f32;
        match ViewportHandle::new_from_raw(self.raw_window, self.raw_display, logical_w, logical_h, scale) {
            Some(h) => {
                tracing::info!(
                    "wayland: viewport ready ({}x{} @ {})",
                    logical_w,
                    logical_h,
                    scale
                );
                self.viewport = Some(h);
            }
            None => tracing::error!("wayland: ViewportHandle::new_from_raw returned None"),
        }
    }

    fn render(&mut self, qh: &QueueHandle<Self>) {
        if let Some(vp) = self.viewport.as_mut() {
            vp.render_frame();
        }
        let surface = self.layer.wl_surface();
        surface.frame(qh, surface.clone());
        self.layer.commit();
    }
}

impl CompositorHandler for State {
    fn scale_factor_changed(
        &mut self,
        _conn: &Connection,
        _qh: &QueueHandle<Self>,
        _surface: &wl_surface::WlSurface,
        new_factor: i32,
    ) {
        if self.scale == new_factor {
            return;
        }
        self.scale = new_factor;
        self.layer.wl_surface().set_buffer_scale(new_factor);
        if let Some(vp) = self.viewport.as_mut() {
            vp.resize_px(self.width as f32, self.height as f32, new_factor as f32);
        }
    }

    fn transform_changed(
        &mut self,
        _: &Connection,
        _: &QueueHandle<Self>,
        _: &wl_surface::WlSurface,
        _: wl_output::Transform,
    ) {
    }

    fn frame(
        &mut self,
        _conn: &Connection,
        qh: &QueueHandle<Self>,
        _: &wl_surface::WlSurface,
        _time: u32,
    ) {
        self.render(qh);
    }

    fn surface_enter(
        &mut self,
        _: &Connection,
        _: &QueueHandle<Self>,
        _: &wl_surface::WlSurface,
        _: &wl_output::WlOutput,
    ) {
    }

    fn surface_leave(
        &mut self,
        _: &Connection,
        _: &QueueHandle<Self>,
        _: &wl_surface::WlSurface,
        _: &wl_output::WlOutput,
    ) {
    }
}

impl OutputHandler for State {
    fn output_state(&mut self) -> &mut OutputState {
        &mut self.output_state
    }
    fn new_output(&mut self, _: &Connection, _: &QueueHandle<Self>, _: wl_output::WlOutput) {}
    fn update_output(&mut self, _: &Connection, _: &QueueHandle<Self>, _: wl_output::WlOutput) {}
    fn output_destroyed(&mut self, _: &Connection, _: &QueueHandle<Self>, _: wl_output::WlOutput) {}
}

impl LayerShellHandler for State {
    fn closed(&mut self, _: &Connection, _: &QueueHandle<Self>, _: &LayerSurface) {
        self.exit = true;
    }

    fn configure(
        &mut self,
        _conn: &Connection,
        qh: &QueueHandle<Self>,
        _layer: &LayerSurface,
        configure: LayerSurfaceConfigure,
        serial: u32,
    ) {
        let (w, h) = configure.new_size;
        let new_w = if w == 0 { self.width } else { w };
        let new_h = if h == 0 { self.height } else { h };
        let resized = new_w != self.width || new_h != self.height;
        tracing::info!(
            serial,
            requested_w = w,
            requested_h = h,
            applied_w = new_w,
            applied_h = new_h,
            first = self.first_configure,
            "wayland: configure"
        );
        self.width = new_w;
        self.height = new_h;

        self.ensure_viewport();
        if resized {
            if let Some(vp) = self.viewport.as_mut() {
                vp.resize_px(self.width as f32, self.height as f32, self.scale.max(1) as f32);
            }
        }

        if self.first_configure {
            self.first_configure = false;
            tracing::info!("wayland: first configure → kicking initial render");
            self.render(qh);
        }
    }
}

impl SeatHandler for State {
    fn seat_state(&mut self) -> &mut SeatState {
        &mut self.seat_state
    }
    fn new_seat(&mut self, _: &Connection, _: &QueueHandle<Self>, _: wl_seat::WlSeat) {}

    fn new_capability(
        &mut self,
        _conn: &Connection,
        qh: &QueueHandle<Self>,
        seat: wl_seat::WlSeat,
        capability: Capability,
    ) {
        if capability == Capability::Keyboard && self.keyboard.is_none() {
            if let Ok(kb) = self.seat_state.get_keyboard(qh, &seat, None) {
                self.keyboard = Some(kb);
            }
        }
        if capability == Capability::Pointer && self.pointer.is_none() {
            if let Ok(p) = self.seat_state.get_pointer(qh, &seat) {
                self.pointer = Some(p);
            }
        }
    }

    fn remove_capability(
        &mut self,
        _conn: &Connection,
        _: &QueueHandle<Self>,
        _: wl_seat::WlSeat,
        capability: Capability,
    ) {
        if capability == Capability::Keyboard {
            if let Some(k) = self.keyboard.take() {
                k.release();
            }
        }
        if capability == Capability::Pointer {
            if let Some(p) = self.pointer.take() {
                p.release();
            }
        }
    }

    fn remove_seat(&mut self, _: &Connection, _: &QueueHandle<Self>, _: wl_seat::WlSeat) {}
}

impl KeyboardHandler for State {
    fn enter(
        &mut self,
        _: &Connection,
        _: &QueueHandle<Self>,
        _: &wl_keyboard::WlKeyboard,
        _surface: &wl_surface::WlSurface,
        _: u32,
        _: &[u32],
        _keysyms: &[Keysym],
    ) {
    }

    fn leave(
        &mut self,
        _: &Connection,
        _: &QueueHandle<Self>,
        _: &wl_keyboard::WlKeyboard,
        _surface: &wl_surface::WlSurface,
        _: u32,
    ) {
    }

    fn press_key(
        &mut self,
        _: &Connection,
        _: &QueueHandle<Self>,
        _: &wl_keyboard::WlKeyboard,
        _: u32,
        event: KeyEvent,
    ) {
        if let Some(vp) = self.viewport.as_mut() {
            let (named, utf8) = map_key_event(&event);
            vp.push_key_event(named, utf8, encode_modifiers(&self.modifiers), true);
        }
    }

    fn release_key(
        &mut self,
        _: &Connection,
        _: &QueueHandle<Self>,
        _: &wl_keyboard::WlKeyboard,
        _: u32,
        event: KeyEvent,
    ) {
        if let Some(vp) = self.viewport.as_mut() {
            let (named, utf8) = map_key_event(&event);
            vp.push_key_event(named, utf8, encode_modifiers(&self.modifiers), false);
        }
    }

    fn repeat_key(
        &mut self,
        _: &Connection,
        _: &QueueHandle<Self>,
        _: &wl_keyboard::WlKeyboard,
        _: u32,
        _event: KeyEvent,
    ) {
    }

    fn update_modifiers(
        &mut self,
        _: &Connection,
        _: &QueueHandle<Self>,
        _: &wl_keyboard::WlKeyboard,
        _serial: u32,
        modifiers: Modifiers,
        _raw: RawModifiers,
        _layout: u32,
    ) {
        self.modifiers = modifiers;
    }
}

impl PointerHandler for State {
    fn pointer_frame(
        &mut self,
        _: &Connection,
        _: &QueueHandle<Self>,
        _: &wl_pointer::WlPointer,
        events: &[PointerEvent],
    ) {
        let Some(vp) = self.viewport.as_mut() else {
            return;
        };
        for event in events {
            if &event.surface != self.layer.wl_surface() {
                continue;
            }
            let (px, py) = event.position;
            let lx = px as f32;
            let ly = py as f32;
            self.last_pointer = (lx, ly);
            match event.kind {
                PointerEventKind::Enter { .. } => {
                    vp.push_mouse_move(lx, ly);
                }
                PointerEventKind::Leave { .. } => {
                    vp.push_mouse_left();
                }
                PointerEventKind::Motion { .. } => {
                    vp.push_mouse_move(lx, ly);
                }
                PointerEventKind::Press { button, .. } => {
                    if let Some(code) = map_button(button) {
                        vp.push_mouse_button(lx, ly, code, true);
                    }
                }
                PointerEventKind::Release { button, .. } => {
                    if let Some(code) = map_button(button) {
                        vp.push_mouse_button(lx, ly, code, false);
                    }
                }
                PointerEventKind::Axis {
                    horizontal,
                    vertical,
                    ..
                } => {
                    let dx = -(horizontal.absolute as f32);
                    let dy = -(vertical.absolute as f32);
                    vp.push_mouse_scroll(lx, ly, dx, dy);
                }
            }
        }
    }
}

/// Find a wayland socket and connect, regardless of whether `WAYLAND_DISPLAY` is set.
/// Inside flatpak the env var is often stripped even when the socket itself is bind-
/// mounted into the sandbox at `$XDG_RUNTIME_DIR/wayland-*`. We list candidates,
/// log everything we find, and try each in order.
fn probe_connection() -> Option<Connection> {
    let runtime_dir = std::env::var_os("XDG_RUNTIME_DIR").map(std::path::PathBuf::from);
    tracing::info!(
        xdg_runtime_dir = ?runtime_dir,
        wayland_display_env = ?std::env::var_os("WAYLAND_DISPLAY"),
        "wayland: probing for compositor socket"
    );

    if let Ok(c) = Connection::connect_to_env() {
        tracing::info!("wayland: connected via $WAYLAND_DISPLAY");
        return Some(c);
    }

    let Some(dir) = runtime_dir else {
        tracing::warn!("wayland: $XDG_RUNTIME_DIR is unset; cannot scan for sockets");
        return None;
    };

    let entries = match std::fs::read_dir(&dir) {
        Ok(e) => e,
        Err(e) => {
            tracing::warn!(dir = %dir.display(), "wayland: cannot list runtime dir: {e}");
            return None;
        }
    };

    let mut candidates: Vec<String> = Vec::new();
    for entry in entries.flatten() {
        let name = entry.file_name();
        let Some(name_str) = name.to_str() else { continue; };
        if !name_str.starts_with("wayland-") {
            continue;
        }
        // Skip lock files; we want the socket itself.
        if name_str.ends_with(".lock") {
            continue;
        }
        candidates.push(name_str.to_string());
    }
    candidates.sort();
    tracing::info!(?candidates, "wayland: socket candidates in $XDG_RUNTIME_DIR");

    for name in &candidates {
        // Connection::connect_to_env reads $WAYLAND_DISPLAY; set it for one attempt.
        std::env::set_var("WAYLAND_DISPLAY", name);
        match Connection::connect_to_env() {
            Ok(c) => {
                tracing::info!(socket = name, "wayland: connected via probed socket");
                return Some(c);
            }
            Err(e) => {
                tracing::debug!(socket = name, "wayland: probe failed: {e}");
            }
        }
    }

    None
}

fn log_globals(globals: &GlobalList) {
    let list = globals.contents().clone_list();
    tracing::info!("wayland: registry advertises {} globals", list.len());
    let mut seen_layer_shell = false;
    let mut seen_compositor = false;
    let mut seen_seat = false;
    for g in &list {
        tracing::debug!(
            interface = %g.interface,
            name = g.name,
            version = g.version,
            "wayland global"
        );
        match g.interface.as_str() {
            "zwlr_layer_shell_v1" => seen_layer_shell = true,
            "wl_compositor" => seen_compositor = true,
            "wl_seat" => seen_seat = true,
            _ => {}
        }
    }
    tracing::info!(
        layer_shell = seen_layer_shell,
        compositor = seen_compositor,
        seat = seen_seat,
        "wayland: protocol summary"
    );
}

fn map_button(button: u32) -> Option<u32> {
    match button {
        0x110 => Some(0), // BTN_LEFT
        0x111 => Some(1), // BTN_RIGHT
        0x112 => Some(2), // BTN_MIDDLE
        _ => None,
    }
}

fn map_key_event(event: &KeyEvent) -> (u32, Option<String>) {
    let named = match event.keysym {
        Keysym::Return => 1,
        Keysym::Escape => 2,
        Keysym::BackSpace => 3,
        Keysym::Tab => 4,
        Keysym::Left => 5,
        Keysym::Right => 6,
        Keysym::Up => 7,
        Keysym::Down => 8,
        Keysym::Delete => 9,
        Keysym::Home => 10,
        Keysym::End => 11,
        _ => 0,
    };
    let utf8 = if named == 0 { event.utf8.clone() } else { None };
    (named, utf8)
}

fn encode_modifiers(mods: &Modifiers) -> u32 {
    let mut bits = 0u32;
    if mods.shift {
        bits |= 1;
    }
    if mods.ctrl {
        bits |= 2;
    }
    if mods.alt {
        bits |= 4;
    }
    if mods.logo {
        bits |= 8;
    }
    bits
}

delegate_compositor!(State);
delegate_output!(State);
delegate_seat!(State);
delegate_keyboard!(State);
delegate_pointer!(State);
delegate_layer!(State);
delegate_registry!(State);

impl ProvidesRegistryState for State {
    fn registry(&mut self) -> &mut RegistryState {
        &mut self.registry_state
    }
    registry_handlers![OutputState, SeatState];
}