gpupaper/src/wayland.rs

//! Wayland state: output enumeration, layer surface creation, event dispatch.

use std::sync::Arc;
use std::time::{Duration, Instant};

use anyhow::{bail, Context, Result};
use smithay_client_toolkit::{
    compositor::{CompositorHandler, CompositorState},
    delegate_compositor, delegate_layer, delegate_output, delegate_registry,
    output::{OutputHandler, OutputState},
    registry::{ProvidesRegistryState, RegistryState},
    registry_handlers,
    shell::{
        wlr_layer::{
            Anchor, KeyboardInteractivity, Layer, LayerShell, LayerShellHandler, LayerSurface,
            LayerSurfaceConfigure,
        },
        WaylandSurface,
    },
};
use wayland_client::{
    globals::registry_queue_init,
    protocol::{wl_output, wl_surface},
    Connection, Proxy, QueueHandle,
};

use crate::renderer::{GpuContext, SurfaceRenderer};
use crate::shader::ShaderSource;

// ---------------------------------------------------------------------------
// Per-output state
// ---------------------------------------------------------------------------

struct OutputSurface {
    layer: LayerSurface,
    renderer: Option<SurfaceRenderer>,
    /// Raw wl_surface pointer passed to SurfaceRenderer.
    surface_ptr: *mut std::ffi::c_void,
    width: u32,
    height: u32,
    draw_ready: bool,
    last_frame: Instant,
}

// Safety: raw pointer is only used from the single Wayland thread.
unsafe impl Send for OutputSurface {}
unsafe impl Sync for OutputSurface {}

// ---------------------------------------------------------------------------
// Application state
// ---------------------------------------------------------------------------

pub struct AppState {
    registry_state: RegistryState,
    output_state: OutputState,
    compositor: CompositorState,
    layer_shell: LayerShell,

    outputs: Vec<OutputSurface>,
    /// Created on first configure; shared by all SurfaceRenderers.
    gpu: Option<Arc<GpuContext>>,

    display_ptr: *mut std::ffi::c_void,
    shader: ShaderSource,
    fps: u32,
    layer_kind: Layer,
    /// "all" / "*" → create a surface per output; otherwise match by name.
    target_output: String,

    frame_duration: Option<Duration>,
    pub exit: bool,
}

unsafe impl Send for AppState {}
unsafe impl Sync for AppState {}

// ---------------------------------------------------------------------------
// Layer creation helper (also called from new_output for hotplug)
// ---------------------------------------------------------------------------

impl AppState {
    fn attach_output(&mut self, qh: &QueueHandle<Self>, wl_out: Option<&wl_output::WlOutput>) {
        let surface = self.compositor.create_surface(qh);
        let surface_ptr = surface.id().as_ptr() as *mut std::ffi::c_void;

        let layer = self.layer_shell.create_layer_surface(
            qh,
            surface,
            self.layer_kind,
            Some("gpupaper"),
            wl_out,
        );
        layer.set_anchor(Anchor::TOP | Anchor::BOTTOM | Anchor::LEFT | Anchor::RIGHT);
        layer.set_exclusive_zone(-1);
        layer.set_size(0, 0);
        layer.set_keyboard_interactivity(KeyboardInteractivity::None);
        layer.commit();

        self.outputs.push(OutputSurface {
            layer,
            renderer: None,
            surface_ptr,
            width: 0,
            height: 0,
            draw_ready: false,
            last_frame: Instant::now(),
        });
    }
}

// ---------------------------------------------------------------------------
// Public entry point
// ---------------------------------------------------------------------------

pub fn parse_layer(s: &str) -> Result<Layer> {
    match s.to_lowercase().as_str() {
        "background" => Ok(Layer::Background),
        "bottom" => Ok(Layer::Bottom),
        "top" => Ok(Layer::Top),
        "overlay" => Ok(Layer::Overlay),
        other => bail!("unknown layer: {}", other),
    }
}

pub fn run(
    target_output: String,
    shader: ShaderSource,
    fps: u32,
    layer_kind: Layer,
) -> Result<()> {
    let conn = Connection::connect_to_env().context("failed to connect to Wayland display")?;
    let display_ptr = conn.backend().display_ptr() as *mut std::ffi::c_void;

    let (globals, mut event_queue) =
        registry_queue_init::<AppState>(&conn).context("registry_queue_init")?;
    let qh = event_queue.handle();

    let compositor =
        CompositorState::bind(&globals, &qh).context("wl_compositor not available")?;
    let layer_shell =
        LayerShell::bind(&globals, &qh).context("zwlr_layer_shell_v1 not available")?;

    let frame_duration = if fps > 0 {
        Some(Duration::from_nanos(1_000_000_000 / fps as u64))
    } else {
        None
    };

    let mut state = AppState {
        registry_state: RegistryState::new(&globals),
        output_state: OutputState::new(&globals, &qh),
        compositor,
        layer_shell,
        outputs: Vec::new(),
        gpu: None,
        display_ptr,
        shader,
        fps,
        layer_kind,
        target_output: target_output.clone(),
        frame_duration,
        exit: false,
    };

    // Discover available outputs.
    event_queue.roundtrip(&mut state).context("initial roundtrip")?;

    let is_all = target_output == "all" || target_output == "*";

    if is_all {
        let wl_outputs: Vec<_> = state.output_state.outputs().collect();
        if wl_outputs.is_empty() {
            // No outputs reported yet; let the compositor assign.
            state.attach_output(&qh, None);
        } else {
            for wl_out in &wl_outputs {
                state.attach_output(&qh, Some(wl_out));
            }
        }
    } else {
        let wl_out = find_output(&state.output_state, &target_output);
        state.attach_output(&qh, wl_out.as_ref());
    }

    // Wait for every surface to receive its configure (and thus create a renderer).
    log::info!("waiting for {} surface(s) to configure…", state.outputs.len());
    while state.outputs.iter().any(|o| o.renderer.is_none()) {
        event_queue
            .blocking_dispatch(&mut state)
            .context("event dispatch failed")?;

        if state.exit {
            return Ok(());
        }
    }
    log::info!("{} surface(s) ready", state.outputs.len());

    // Render the first frame on each surface immediately so the compositor has
    // a buffer and will start sending frame callbacks.
    if let Some(gpu) = state.gpu.as_ref().map(Arc::clone) {
        let elapsed = gpu.start.elapsed().as_secs_f32();
        for output in &mut state.outputs {
            let wl_surface = output.layer.wl_surface();
            wl_surface.frame(&qh, wl_surface.clone());
            if let Some(r) = &mut output.renderer {
                r.render(&gpu, elapsed);
            }
        }
    }
    conn.flush().ok();

    // Main event loop — each output is driven by its own frame callbacks so
    // monitors at different refresh rates pace themselves independently.
    loop {
        event_queue
            .blocking_dispatch(&mut state)
            .context("event dispatch")?;
        conn.flush().ok();

        if state.exit {
            break;
        }

        let gpu = match state.gpu.as_ref().map(Arc::clone) {
            Some(g) => g,
            None => continue,
        };
        let elapsed = gpu.start.elapsed().as_secs_f32();
        let frame_duration = state.frame_duration;

        for output in &mut state.outputs {
            if !output.draw_ready {
                continue;
            }
            output.draw_ready = false;

            if let Some(dur) = frame_duration {
                let since_last = output.last_frame.elapsed();
                if since_last < dur {
                    std::thread::sleep(dur - since_last);
                }
            }
            output.last_frame = Instant::now();

            {
                let wl_surface = output.layer.wl_surface();
                wl_surface.frame(&qh, wl_surface.clone());
            }

            if let Some(r) = &mut output.renderer {
                r.render(&gpu, elapsed);
            }
        }

        conn.flush().ok();
    }

    Ok(())
}

fn find_output(output_state: &OutputState, target: &str) -> Option<wl_output::WlOutput> {
    for output in output_state.outputs() {
        if let Some(info) = output_state.info(&output) {
            if let Some(name) = &info.name {
                if name == target {
                    log::info!("matched output: {}", name);
                    return Some(output);
                }
            }
        }
    }
    log::warn!("output '{}' not found; using default", target);
    None
}

// ---------------------------------------------------------------------------
// Trait implementations
// ---------------------------------------------------------------------------

impl CompositorHandler for AppState {
    fn scale_factor_changed(
        &mut self,
        _: &Connection,
        _: &QueueHandle<Self>,
        _: &wl_surface::WlSurface,
        _: i32,
    ) {
    }

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

    fn frame(
        &mut self,
        _: &Connection,
        _: &QueueHandle<Self>,
        surface: &wl_surface::WlSurface,
        _: u32,
    ) {
        let id = surface.id();
        for output in &mut self.outputs {
            if output.layer.wl_surface().id() == id {
                output.draw_ready = true;
                break;
            }
        }
    }

    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 AppState {
    fn output_state(&mut self) -> &mut OutputState {
        &mut self.output_state
    }

    fn new_output(
        &mut self,
        _conn: &Connection,
        qh: &QueueHandle<Self>,
        output: wl_output::WlOutput,
    ) {
        if self.target_output == "all" || self.target_output == "*" {
            self.attach_output(qh, Some(&output));
        }
    }

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

    fn output_destroyed(
        &mut self,
        _: &Connection,
        _: &QueueHandle<Self>,
        output: wl_output::WlOutput,
    ) {
        // Drop the surface renderer for the disconnected output.
        self.outputs
            .retain(|o| o.layer.wl_surface().id() != output.id());
    }
}

impl LayerShellHandler for AppState {
    fn closed(&mut self, _: &Connection, _: &QueueHandle<Self>, _: &LayerSurface) {
        log::info!("layer surface closed");
        self.exit = true;
    }

    fn configure(
        &mut self,
        _conn: &Connection,
        _qh: &QueueHandle<Self>,
        layer: &LayerSurface,
        configure: LayerSurfaceConfigure,
        _serial: u32,
    ) {
        let layer_id = layer.wl_surface().id();
        let Some(idx) = self.outputs.iter().position(|o| o.layer.wl_surface().id() == layer_id)
        else {
            return;
        };

        let (w, h) = configure.new_size;
        let w = if w == 0 { 1920 } else { w };
        let h = if h == 0 { 1080 } else { h };
        self.outputs[idx].width = w;
        self.outputs[idx].height = h;

        if self.outputs[idx].renderer.is_some() {
            // Subsequent configure — just resize.
            let gpu = self.gpu.as_ref().map(Arc::clone);
            if let Some(gpu) = gpu {
                if let Some(r) = self.outputs[idx].renderer.as_mut() {
                    r.resize(&gpu, w, h);
                }
            }
            return;
        }

        // First configure for this surface — create the GPU context if needed.
        if self.gpu.is_none() {
            match unsafe { GpuContext::new(self.display_ptr) } {
                Ok(gpu) => self.gpu = Some(Arc::new(gpu)),
                Err(e) => {
                    log::error!("failed to create GPU context: {e}");
                    self.exit = true;
                    return;
                }
            }
        }

        let gpu = Arc::clone(self.gpu.as_ref().unwrap());
        let surface_ptr = self.outputs[idx].surface_ptr;
        let display_ptr = self.display_ptr;
        let shader = &self.shader;
        let fps = self.fps;

        match unsafe { SurfaceRenderer::new(&gpu, display_ptr, surface_ptr, w, h, shader, fps) } {
            Ok(renderer) => {
                self.outputs[idx].renderer = Some(renderer);
                log::info!("output {} configured: {}x{}", idx, w, h);
            }
            Err(e) => {
                log::error!("failed to create surface renderer for output {idx}: {e}");
                self.exit = true;
            }
        }
    }
}

delegate_compositor!(AppState);
delegate_output!(AppState);
delegate_layer!(AppState);
delegate_registry!(AppState);

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