scrim 0.1.0

Wayland full-screen color overlay with file interface.
use std::sync::mpsc;
use std::time::{Duration, Instant};

use smithay_client_toolkit::{
    compositor::{CompositorHandler, CompositorState, Region},
    delegate_compositor, delegate_layer, delegate_output, delegate_registry,
    delegate_seat, delegate_shm,
    output::{OutputHandler, OutputState},
    registry::{ProvidesRegistryState, RegistryState},
    registry_handlers,
    seat::{SeatHandler, SeatState, Capability},
    shell::{
        wlr_layer::{
            Anchor, KeyboardInteractivity, Layer, LayerShell, LayerShellHandler,
            LayerSurface, LayerSurfaceConfigure,
        },
        WaylandSurface,
    },
    shm::{slot::SlotPool, Shm, ShmHandler},
};
use wayland_client::{
    globals::GlobalList,
    protocol::{wl_output, wl_seat, wl_shm, wl_surface},
    Connection, QueueHandle,
};

use crate::watcher::Msg;

pub struct Scrim {
    registry: RegistryState,
    output_state: OutputState,
    seat_state: SeatState,
    compositor: CompositorState,
    shm: Shm,

    layer: Option<LayerSurface>,
    pool: Option<SlotPool>,
    current_buffer: Option<smithay_client_toolkit::shm::slot::Buffer>,
    width: u32,
    height: u32,

    r: u8,
    g: u8,
    b: u8,
    alpha: u8,
    fade_from: u8,
    fade_to: u8,
    fade_dur: Duration,
    fade_start: Instant,
    fading: bool,
    pub dirty: bool,
    wait_frame: bool,

    rx: mpsc::Receiver<Msg>,
}

delegate_compositor!(Scrim);
delegate_shm!(Scrim);
delegate_layer!(Scrim);
delegate_output!(Scrim);
delegate_seat!(Scrim);
delegate_registry!(Scrim);

impl Scrim {
    pub fn new(
        _conn: &Connection,
        globals: &GlobalList,
        qh: &QueueHandle<Self>,
        rx: mpsc::Receiver<Msg>,
    ) -> Self {
        let registry = RegistryState::new(globals);
        let output_state = OutputState::new(globals, qh);
        let seat_state = SeatState::new(globals, qh);
        let compositor = CompositorState::bind(globals, qh).expect("compositor");
        let shm = Shm::bind(globals, qh).expect("shm");
        let layer_shell = LayerShell::bind(globals, qh).expect("layer-shell");

        let wl_surface = compositor.create_surface(qh);
        let layer = layer_shell.create_layer_surface(
            qh, wl_surface, Layer::Overlay, Some("scrim"), None,
        );

        layer.set_anchor(Anchor::TOP | Anchor::BOTTOM | Anchor::LEFT | Anchor::RIGHT);
        layer.set_size(0, 0);
        layer.set_exclusive_zone(-1);
        layer.set_keyboard_interactivity(KeyboardInteractivity::None);

        // 点击穿透
        if let Ok(region) = Region::new(&compositor) {
            layer.wl_surface().set_input_region(Some(region.wl_region()));
        }

        layer.commit();

        let r = read_file_u8("/tmp/scrim/r", 0);
        let g = read_file_u8("/tmp/scrim/g", 0);
        let b = read_file_u8("/tmp/scrim/b", 0);
        let alpha = read_file_u8("/tmp/scrim/alpha", 0);

        Self {
            registry, output_state, seat_state, compositor, shm,
            layer: Some(layer),
            pool: None,
            current_buffer: None,
            width: 0, height: 0,
            r, g, b, alpha,
            fade_from: 0, fade_to: 0,
            fade_dur: Duration::ZERO,
            fade_start: Instant::now(),
            fading: false,
            dirty: true,
            rx,
            wait_frame: false,
        }
    }

    pub fn poll_messages(&mut self) {
        while let Ok(msg) = self.rx.try_recv() {
            self.handle_msg(msg);
        }
    }

    pub fn wait_message(&mut self, timeout: Duration) -> Option<Msg> {
        self.rx.recv_timeout(timeout).ok()
    }

    pub fn handle_msg(&mut self, msg: Msg) {
        match msg {
            Msg::SetAlpha(a) => {
                self.alpha = a;
                self.fading = false;
                self.dirty = true;
            }
            Msg::SetR(v) => {
                self.r = v;
                self.dirty = true;
            }
            Msg::SetG(v) => {
                self.g = v;
                self.dirty = true;
            }
            Msg::SetB(v) => {
                self.b = v;
                self.dirty = true;
            }
            Msg::Fade { target, duration_ms } => {
                self.fade_from = self.alpha;
                self.fade_to = target;
                self.fade_dur = Duration::from_millis(duration_ms);
                self.fade_start = Instant::now();
                self.fading = true;
                self.dirty = true;
            }
        }
    }

    pub fn tick(&mut self) {
        if !self.fading { return; }
        let t = self.fade_start.elapsed().as_secs_f32() / self.fade_dur.as_secs_f32();
        if t >= 1.0 {
            self.alpha = self.fade_to;
            self.fading = false;
        } else {
            let from = self.fade_from as f32;
            let to = self.fade_to as f32;
            self.alpha = (from + (to - from) * t) as u8;
        }
        self.dirty = true;
    }

    pub fn animating(&self) -> bool { self.fading }
    pub fn waiting_frame(&self) -> bool { self.wait_frame }

    pub fn draw(&mut self, qh: &QueueHandle<Self>) {
        if self.wait_frame { return; }  // 等合成器回调,不重叠绘制

        let layer = match &self.layer { Some(l) => l, None => return };
        let pool = match &mut self.pool { Some(p) => p, None => return };
        if self.width == 0 || self.height == 0 { return; }

        self.current_buffer = None;

        let pixel: [u8; 4] = [self.b, self.g, self.r, self.alpha];

        let w = self.width as usize;
        let h = self.height as usize;
        let stride = w * 4;

        let (buf, canvas) = pool
            .create_buffer(w as i32, h as i32, stride as i32, wl_shm::Format::Argb8888)
            .expect("创建 buffer 失败");
        // 填第一行
        for x in 0..w {
            canvas[x * 4..x * 4 + 4].copy_from_slice(&pixel);
        }
        // 复制到剩余行
        for y in 1..h {
            canvas.copy_within(0..stride, y * stride);
        }

        let s = layer.wl_surface();
        s.damage_buffer(0, 0, self.width as i32, self.height as i32);
        s.attach(Some(buf.wl_buffer()), 0, 0);
        s.frame(qh, s.clone());
        s.commit();

        self.current_buffer = Some(buf);
        self.wait_frame = true;        // 锁住,等回调
    }
}

// --- CompositorHandler ---

impl CompositorHandler for Scrim {
    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>,
        _: &wl_surface::WlSurface, _: u32)
    {
        self.wait_frame = false;       // 合成器说:可以画下一帧了
        if self.alpha > 0 || self.fading {
            self.dirty = true;         // 还需要继续画
        }
    }
    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) {}
}

// --- LayerShellHandler ---

impl LayerShellHandler for Scrim {
    fn configure(&mut self, _: &Connection, _: &QueueHandle<Self>,
        _: &LayerSurface, configure: LayerSurfaceConfigure, _: u32)
    {
        let (w, h) = configure.new_size;
        if w == 0 || h == 0 { return; }

        self.width = w;
        self.height = h;
        eprintln!("scrim: 收到配置尺寸 {}x{}", w, h);

        let size = (w * h * 4) as usize;
        match SlotPool::new(size, &self.shm) {
            Ok(p) => self.pool = Some(p),
            Err(e) => eprintln!("scrim: SHM pool 错误: {}", e),
        }

        self.dirty = true;
    }

    fn closed(&mut self, _: &Connection, _: &QueueHandle<Self>, _: &LayerSurface) {
        std::process::exit(0);
    }
}

// --- OutputHandler ---

impl OutputHandler for Scrim {
    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) {}
}

// --- SeatHandler ---

// --- SeatHandler ---

impl SeatHandler for Scrim {
    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, _: &Connection, _: &QueueHandle<Self>,
        _: wl_seat::WlSeat, _: Capability) {}
    fn remove_capability(&mut self, _: &Connection, _: &QueueHandle<Self>,
        _: wl_seat::WlSeat, _: Capability) {}
    fn remove_seat(&mut self, _: &Connection, _: &QueueHandle<Self>, _: wl_seat::WlSeat) {}
}

// --- ShmHandler ---

impl ShmHandler for Scrim {
    fn shm_state(&mut self) -> &mut Shm { &mut self.shm }
}

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

// --- 工具函数 ---

fn read_file_u8(path: &str, default: u8) -> u8 {
    std::fs::read_to_string(path)
        .ok()
        .and_then(|s| s.trim().parse().ok())
        .unwrap_or(default)
}