wl-webrtc/src/state.rs

use std::collections::HashMap;
use std::mem;
use std::os::fd::{AsFd, OwnedFd};
use std::os::unix::io::FromRawFd;
use std::path::{Path, PathBuf};
use std::time::Instant;

use anyhow::Result;
use wayland_client::backend::ObjectId;
use wayland_client::globals::{GlobalList, GlobalListContents};
use wayland_client::protocol::wl_buffer::WlBuffer;
use wayland_client::protocol::wl_output::WlOutput;
use wayland_client::protocol::wl_registry::WlRegistry;
use wayland_client::{event_created_child, Dispatch, Proxy, QueueHandle};
use wayland_protocols::wp::linux_dmabuf::zv1::client::zwp_linux_buffer_params_v1::{
    Event as BufferParamsEvent, Flags as BufferParamsFlags, ZwpLinuxBufferParamsV1,
};
use wayland_protocols::wp::linux_dmabuf::zv1::client::zwp_linux_dmabuf_feedback_v1::{
    Event as DmabufFeedbackEvent, ZwpLinuxDmabufFeedbackV1,
};
use wayland_protocols::wp::linux_dmabuf::zv1::client::zwp_linux_dmabuf_v1::{
    Event as DmabufEvent, ZwpLinuxDmabufV1,
};
use wayland_protocols::xdg::xdg_output::zv1::client::zxdg_output_manager_v1::ZxdgOutputManagerV1;
use wayland_protocols::xdg::xdg_output::zv1::client::zxdg_output_v1::{
    Event as XdgOutputEvent, ZxdgOutputV1,
};
use wayland_protocols_wlr::output_management::v1::client::zwlr_output_head_v1::{
    self, Event as WlrHeadEvent, ZwlrOutputHeadV1,
};
use wayland_protocols_wlr::output_management::v1::client::zwlr_output_manager_v1::{
    self, Event as WlrOutputManagerEvent, ZwlrOutputManagerV1,
};
use wayland_protocols_wlr::output_management::v1::client::zwlr_output_mode_v1::ZwlrOutputModeV1;
use wayland_protocols_wlr::screencopy::v1::client::zwlr_screencopy_frame_v1::{
    Event as ScreencopyFrameEvent, ZwlrScreencopyFrameV1,
};
use wayland_protocols_wlr::screencopy::v1::client::zwlr_screencopy_manager_v1::ZwlrScreencopyManagerV1;

use ffmpeg_next as ff;
use ffmpeg_next::ffi;

use crate::args::Args;
use crate::avhw::{AvHwDevCtx, EncState};
use crate::cap_wlr_screencopy::CapWlrScreencopy;
use crate::fps_limit::FpsLimit;
use crate::transform::{transpose_if_transform_transposed, Transform};

// ---------------------------------------------------------------------------
// CaptureSource trait
// ---------------------------------------------------------------------------

/// Screen capture backend trait.
pub trait CaptureSource: Sized + 'static {
    type Frame: Send;

    fn new(
        gm: &GlobalList,
        output: &WlOutput,
        output_info: &OutputInfo,
        qh: &QueueHandle<State<Self>>,
    ) -> Result<Self>;

    fn alloc_frame(&mut self) -> Option<Self::Frame>;

    fn queue_copy(&mut self, buffer: &WlBuffer, qh: &QueueHandle<State<Self>>);

    fn on_done_with_frame(&mut self, frame: Self::Frame);
}

// ---------------------------------------------------------------------------
// Output info types
// ---------------------------------------------------------------------------

pub struct OutputInfo {
    pub name: String,
    pub transform: Transform,
    pub physical_size: (i32, i32),
    pub logical_position: (i32, i32),
}

pub struct PartialOutputInfo {
    pub name: Option<String>,
    /// Name from wl_output::Name (v4) — used to match wlr-output-management heads
    pub wl_name: Option<String>,
    pub transform: Option<Transform>,
    pub physical_size: Option<(i32, i32)>,
    pub logical_position: Option<(i32, i32)>,
    // Pixel dimensions from Mode event — preparatory for Phase 2 resolution logic
    pub mode_size: Option<(i32, i32)>,
    pub done_count: u32,
}

impl Default for PartialOutputInfo {
    fn default() -> Self {
        Self {
            name: None,
            wl_name: None,
            transform: None,
            physical_size: None,
            logical_position: None,
            mode_size: None,
            done_count: 0,
        }
    }
}

/// Stores head info from wlr-output-management for name-based matching with wl_output.
struct WlrHeadInfo {
    position: Option<(i32, i32)>,
}

/// User data for XdgOutput dispatch to identify which WlOutput it belongs to.
pub struct OutputId(pub u32);

// ---------------------------------------------------------------------------
// EncConstructionStage
// ---------------------------------------------------------------------------

pub enum EncConstructionStage<S: CaptureSource> {
    ProbingOutputs {
        outputs: Vec<PartialOutputInfo>,
        bound_outputs: Vec<WlOutput>,
        output_names: Vec<u32>,
        screencopy_manager: Option<ZwlrScreencopyManagerV1>,
        dmabuf: Option<ZwpLinuxDmabufV1>,
        dmabuf_feedback: Option<ZwpLinuxDmabufFeedbackV1>,
        xdg_output_manager: Option<ZxdgOutputManagerV1>,
        wlr_output_manager: Option<ZwlrOutputManagerV1>,
        wlr_manager_done: bool,
        wlr_heads: HashMap<String, WlrHeadInfo>,
        wlr_head_proxy_to_name: HashMap<ObjectId, String>,
    },
    EverythingButFmt {
        output_info: OutputInfo,
        output: WlOutput,
        hw_device_ctx: AvHwDevCtx,
        cap: S,
        screencopy_manager: ZwlrScreencopyManagerV1,
        dmabuf: ZwpLinuxDmabufV1,
    },
    Streaming {
        output_info: OutputInfo,
        output: WlOutput,
        enc: EncState,
        cap: S,
        screencopy_manager: ZwlrScreencopyManagerV1,
        dmabuf: ZwpLinuxDmabufV1,
    },
    Intermediate,
}

// ---------------------------------------------------------------------------
// InFlightSurface
// ---------------------------------------------------------------------------

pub enum InFlightSurface<S: CaptureSource> {
    None,
    AllocQueued,
    Allocd(S::Frame),
    CopyQueued {
        surface: ff::frame::Video,
        drm_map: ff::ffi::AVDRMFrameDescriptor,
        frame: S::Frame,
        buffer: WlBuffer,
    },
}

// ---------------------------------------------------------------------------
// State
// ---------------------------------------------------------------------------

pub struct State<S: CaptureSource> {
    pub stage: EncConstructionStage<S>,
    pub in_flight_surface: InFlightSurface<S>,
    pub starting_timestamp: Option<i64>,
    pub first_frame: bool,
    pub args: Args,
    pub errored: bool,
    pub gm: GlobalList,
    pub fps_limit: FpsLimit<S::Frame>,
    pub qhandle: QueueHandle<State<S>>,
    pub drm_device: Option<PathBuf>,
    pub drm_device_from_compositor: Option<PathBuf>,
}

// ---------------------------------------------------------------------------
// Helpers
// ---------------------------------------------------------------------------

/// Scan /dev/dri for the first available DRM render node (renderD*).
fn find_drm_render_node() -> Option<PathBuf> {
    std::fs::read_dir("/dev/dri")
        .ok()?
        .filter_map(|e| e.ok())
        .filter(|e| {
            e.file_name()
                .to_str()
                .map(|s| s.starts_with("renderD"))
                .unwrap_or(false)
        })
        .filter_map(|e| {
            let path = e.path();
            std::fs::metadata(&path).ok()?;
            Some(path)
        })
        .min_by_key(|e| e.to_path_buf())
}

impl<S: CaptureSource> State<S> {
    fn resolve_drm_path(&self) -> PathBuf {
        self.drm_device
            .clone()
            .or_else(|| self.drm_device_from_compositor.clone())
            .or_else(find_drm_render_node)
            .unwrap_or_else(|| PathBuf::from("/dev/dri/renderD128"))
    }
}

// ---------------------------------------------------------------------------
// State<S> methods
// ---------------------------------------------------------------------------

impl<S: CaptureSource> State<S> {
    pub fn new(gm: GlobalList, args: Args, qhandle: QueueHandle<State<S>>) -> Self {
        let fps = args.fps;
        let drm_device = args.drm_device.as_ref().map(PathBuf::from);
        let mut state = Self {
            stage: EncConstructionStage::ProbingOutputs {
                outputs: Vec::new(),
                bound_outputs: Vec::new(),
                output_names: Vec::new(),
                screencopy_manager: None,
                dmabuf: None,
                dmabuf_feedback: None,
                xdg_output_manager: None,
                wlr_output_manager: None,
                wlr_manager_done: false,
                wlr_heads: HashMap::new(),
                wlr_head_proxy_to_name: HashMap::new(),
            },
            in_flight_surface: InFlightSurface::None,
            starting_timestamp: None,
            first_frame: true,
            fps_limit: FpsLimit::new(fps),
            args,
            errored: false,
            gm,
            qhandle,
            drm_device,
            drm_device_from_compositor: None,
        };

        // registry_queue_init consumes registry events internally during its
        // initial roundtrip and does NOT forward them to our Dispatch impl.
        // We must manually bind the initial globals here.
        state.bind_initial_globals();

        state
    }

    /// Iterate over the GlobalList from registry_queue_init and bind all
    /// globals we care about. This is necessary because registry_queue_init
    /// consumes registry events during its internal roundtrip without forwarding
    /// them to our Dispatch<WlRegistry> handler.
    fn bind_initial_globals(&mut self) {
        use wayland_client::globals::Global;

        let globals: Vec<Global> = self.gm.contents().clone_list();
        let registry = self.gm.registry();
        let qhandle = &self.qhandle;

        // Sort globals so that managers are bound BEFORE wl_output.
        // This ensures xdg_output_manager and zwlr_output_manager are available
        // when we bind wl_output, so we can immediately get xdg_output / wlr head.
        let globals = {
            fn priority(interface: &str) -> u8 {
                match interface {
                    "zwlr_screencopy_manager_v1" => 0,
                    "zwp_linux_dmabuf_v1" => 0,
                    "zxdg_output_manager_v1" => 1,
                    "zwlr_output_manager_v1" => 1,
                    "wl_output" => 2,
                    _ => 3,
                }
            }
            let mut g = globals;
            g.sort_by_key(|g| priority(&g.interface));
            g
        };

        for Global {
            name,
            interface,
            version,
        } in globals
        {
            match interface.as_str() {
                "zwlr_screencopy_manager_v1" => {
                    let v = version.min(3);
                    tracing::debug!("Init: binding zwlr_screencopy_manager_v1 v{v} (name={name})");
                    let mgr: ZwlrScreencopyManagerV1 = registry.bind(name, v, qhandle, ());
                    if let EncConstructionStage::ProbingOutputs {
                        screencopy_manager, ..
                    } = &mut self.stage
                    {
                        *screencopy_manager = Some(mgr);
                    }
                }
                "zwp_linux_dmabuf_v1" => {
                    let v = version.min(4);
                    tracing::debug!("Init: binding zwp_linux_dmabuf_v1 v{v} (name={name})");
                    let proxy: ZwpLinuxDmabufV1 = registry.bind(name, v, qhandle, ());
                    if let EncConstructionStage::ProbingOutputs {
                        dmabuf,
                        dmabuf_feedback,
                        ..
                    } = &mut self.stage
                    {
                        *dmabuf = Some(proxy.clone());
                        if v >= 4 {
                            let feedback = proxy.get_default_feedback(qhandle, ());
                            *dmabuf_feedback = Some(feedback);
                        }
                    }
                }
                "zxdg_output_manager_v1" => {
                    let v = version.min(3);
                    tracing::debug!("Init: binding zxdg_output_manager_v1 v{v} (name={name})");
                    let xdg_mgr: ZxdgOutputManagerV1 = registry.bind(name, v, qhandle, ());
                    if let EncConstructionStage::ProbingOutputs {
                        bound_outputs,
                        xdg_output_manager,
                        output_names,
                        ..
                    } = &mut self.stage
                    {
                        for (i, output) in bound_outputs.iter().enumerate() {
                            let oname = output_names.get(i).copied().unwrap_or(0);
                            let output_id = OutputId(oname);
                            xdg_mgr.get_xdg_output(output, qhandle, output_id);
                        }
                        *xdg_output_manager = Some(xdg_mgr);
                    }
                }
                "zwlr_output_manager_v1" => {
                    let v = version.min(4);
                    tracing::debug!("Init: binding zwlr_output_manager_v1 v{v} (name={name})");
                    let mgr: ZwlrOutputManagerV1 = registry.bind(name, v, qhandle, ());
                    if let EncConstructionStage::ProbingOutputs {
                        wlr_output_manager, ..
                    } = &mut self.stage
                    {
                        *wlr_output_manager = Some(mgr);
                    }
                }
                "wl_output" => {
                    let v = version.min(4);
                    tracing::debug!("Init: binding wl_output v{v} (name={name})");
                    let output: WlOutput = registry.bind(name, v, qhandle, OutputId(name));
                    if let EncConstructionStage::ProbingOutputs {
                        outputs,
                        bound_outputs,
                        output_names,
                        xdg_output_manager,
                        ..
                    } = &mut self.stage
                    {
                        outputs.push(PartialOutputInfo::default());
                        bound_outputs.push(output.clone());
                        output_names.push(name);
                        if let Some(xdg_mgr) = xdg_output_manager {
                            let output_id = OutputId(name);
                            xdg_mgr.get_xdg_output(&output, qhandle, output_id);
                        }
                    }
                }
                _ => {}
            }
        }
    }

    pub fn queue_alloc_frame(&mut self)
    where
        State<S>: Dispatch<ZwlrScreencopyFrameV1, ()>,
    {
        let (manager, output) = match &self.stage {
            EncConstructionStage::Streaming {
                screencopy_manager,
                output,
                ..
            } => (screencopy_manager.clone(), output.clone()),
            EncConstructionStage::EverythingButFmt {
                screencopy_manager,
                output,
                ..
            } => (screencopy_manager.clone(), output.clone()),
            _ => return,
        };
        match &self.in_flight_surface {
            InFlightSurface::None => {}
            _ => return,
        }
        let _frame_proxy = manager.capture_output(1, &output, &self.qhandle, ());
        self.in_flight_surface = InFlightSurface::AllocQueued;
    }

    pub fn on_frame_allocd(&mut self, frame: S::Frame, format: u32, width: u32, height: u32) {
        let (frames_rgb_ctx, dmabuf, cap) = match &mut self.stage {
            EncConstructionStage::Streaming {
                output_info: _,
                output: _,
                enc,
                dmabuf,
                cap,
                screencopy_manager: _,
            } => (enc.frames_rgb().as_ptr(), dmabuf, cap),
            _ => {
                tracing::warn!("on_frame_allocd: not in Streaming stage");
                return;
            }
        };

        let mut surface = ff::frame::Video::empty();
        // SAFETY: frames_rgb_ctx is a valid AVHWFramesContext pointer; surface
        // is a freshly allocated empty Video frame.
        let ret = unsafe { ffi::av_hwframe_get_buffer(frames_rgb_ctx, surface.as_mut_ptr(), 0) };
        if ret < 0 {
            tracing::error!("av_hwframe_get_buffer failed: error {}", ret);
            self.errored = true;
            return;
        }

        let mut map_frame = ff::frame::Video::empty();
        // SAFETY: Setting format to DRM_PRIME and calling av_hwframe_map creates
        // a mapped view of the GPU surface with DMA-BUF file descriptors.
        unsafe {
            (*map_frame.as_mut_ptr()).format = ffi::AVPixelFormat::AV_PIX_FMT_DRM_PRIME as i32;
        }
        let ret = unsafe { ffi::av_hwframe_map(map_frame.as_mut_ptr(), surface.as_ptr(), 0) };
        if ret < 0 {
            tracing::error!("av_hwframe_map failed: error {}", ret);
            self.errored = true;
            return;
        }

        // SAFETY: After av_hwframe_map with DRM_PRIME format, data[0] points to
        // a valid AVDRMFrameDescriptor.
        let desc: ff::ffi::AVDRMFrameDescriptor = unsafe {
            let desc_ptr = (*map_frame.as_ptr()).data[0] as *const ff::ffi::AVDRMFrameDescriptor;
            std::ptr::read(desc_ptr)
        };

        let params = dmabuf.create_params(&self.qhandle, ());

        for layer_idx in 0..desc.nb_layers as usize {
            let layer = &desc.layers[layer_idx];
            for p in 0..layer.nb_planes as usize {
                let plane = &layer.planes[p];
                let obj = &desc.objects[plane.object_index as usize];
                let mod_hi = (obj.format_modifier >> 32) as u32;
                let mod_lo = (obj.format_modifier & 0xFFFF_FFFF) as u32;
                // SAFETY: obj.fd is a valid DMA-BUF fd. We dup because params.add()
                // takes ownership of the fd, and the original fd is owned by map_frame.
                let fd_dup = unsafe { libc::dup(obj.fd) };
                if fd_dup < 0 {
                    tracing::error!("failed to dup dma-buf fd");
                    // wayland-client does not auto-destroy params on Drop.
                    params.destroy();
                    self.errored = true;
                    return;
                }
                // SAFETY: fd_dup is valid freshly-duped fd.
                let fd_owned = unsafe { OwnedFd::from_raw_fd(fd_dup) };
                params.add(
                    fd_owned.as_fd(),
                    p as u32,
                    plane.offset as u32,
                    plane.pitch as u32,
                    mod_hi,
                    mod_lo,
                );
            }
        }

        let wl_buffer = params.create_immed(
            width as i32,
            height as i32,
            format,
            BufferParamsFlags::empty(),
            &self.qhandle,
            (),
        );
        self.in_flight_surface = InFlightSurface::CopyQueued {
            surface,
            drm_map: desc,
            frame,
            buffer: wl_buffer,
        };
        let buffer_ref = match &self.in_flight_surface {
            InFlightSurface::CopyQueued { buffer, .. } => buffer,
            _ => unreachable!("just set to CopyQueued"),
        };
        cap.queue_copy(buffer_ref, &self.qhandle);
    }

    pub fn on_copy_complete(&mut self, tv_sec: u64, tv_usec: u32)
    where
        S::Frame: Default,
    {
        let (mut surface, _drm_map, frame, buffer) =
            match mem::replace(&mut self.in_flight_surface, InFlightSurface::None) {
                InFlightSurface::CopyQueued {
                    surface,
                    drm_map,
                    frame,
                    buffer,
                } => (surface, drm_map, frame, buffer),
                other => {
                    tracing::warn!("on_copy_complete: unexpected state");
                    self.in_flight_surface = other;
                    return;
                }
            };
        let fps = self.args.fps as i64;
        // PTS in frame-number units (encoder time_base = 1/fps)
        let pts = (tv_sec as i64) * fps + (tv_usec as i64) * fps / 1_000_000;
        surface.set_pts(Some(pts));
        drop(buffer);
        let cap = match &mut self.stage {
            EncConstructionStage::Streaming { cap, .. } => cap,
            _ => {
                tracing::warn!("on_copy_complete: not in Streaming stage");
                return;
            }
        };
        cap.on_done_with_frame(frame);
        let enc = match &mut self.stage {
            EncConstructionStage::Streaming { enc, .. } => enc,
            _ => unreachable!("already checked Streaming above"),
        };
        let should_encode = if self.first_frame {
            self.first_frame = false;
            true
        } else {
            self.fps_limit
                .on_new_frame(S::Frame::default(), Instant::now())
                .is_some()
        };
        if should_encode {
            if let Err(e) = enc.encode_frame(&surface) {
                tracing::error!("encode_frame failed: {}", e);
                self.errored = true;
            }
        }
    }

    pub fn on_copy_fail(&mut self)
    where
        S::Frame: Default,
    {
        tracing::error!("compositor copy failed");
        let taken = mem::replace(&mut self.in_flight_surface, InFlightSurface::None);
        match taken {
            InFlightSurface::CopyQueued {
                buffer,
                frame,
                ..
            } => {
                drop(buffer);
                if let EncConstructionStage::Streaming { cap, .. } = &mut self.stage {
                    cap.on_done_with_frame(frame);
                }
            }
            other => {
                self.in_flight_surface = other;
            }
        }
        self.errored = true;
    }

    pub fn negotiate_format(&mut self, format: u32, width: u32, height: u32) {
        let stage_data = match mem::replace(&mut self.stage, EncConstructionStage::Intermediate) {
            EncConstructionStage::EverythingButFmt {
                output_info,
                output,
                hw_device_ctx: _hw_device_ctx,
                cap,
                screencopy_manager,
                dmabuf,
            } => (output_info, output, cap, screencopy_manager, dmabuf),
            other => {
                tracing::warn!("negotiate_format: not in EverythingButFmt stage");
                self.stage = other;
                return;
            }
        };
        let (output_info, output, cap, screencopy_manager, dmabuf) = stage_data;
        let drm_path = self.resolve_drm_path();
        let fps = self.args.fps;
        let bitrate = self.args.bitrate.unwrap_or_else(|| {
            2 * (width as u64) * (height as u64) * (fps as u64) / 100
        });
        let enc = match crate::avhw::create_encoder(
            &drm_path,
            Path::new(&self.args.output),
            width,
            height,
            fps,
            output_info.transform,
            self.args.bitrate,
            self.args.gop_size,
        ) {
            Ok(enc) => enc,
            Err(e) => {
                tracing::error!("EncState::new failed: {}", e);
                self.errored = true;
                return;
            }
        };
        tracing::info!(
            "Encoder initialized: {}x{} format={} bitrate={}",
            width,
            height,
            format,
            bitrate
        );
        self.stage = EncConstructionStage::Streaming {
            output_info,
            output,
            enc,
            cap,
            screencopy_manager,
            dmabuf,
        };
    }

    fn try_finalize_output(&mut self, _idx: usize) -> bool {
        // Merge wlr head position info into outputs (needed for niri path)
        if let EncConstructionStage::ProbingOutputs {
            outputs, wlr_heads, ..
        } = &mut self.stage
        {
            for info in outputs.iter_mut() {
                if info.logical_position.is_none() {
                    if let Some(ref wl_name) = info.wl_name {
                        if let Some(head_info) = wlr_heads.get(wl_name) {
                            info.logical_position = head_info.position;
                        }
                    }
                }
            }
        }

        let (target_idx, output_count) = match &self.stage {
            EncConstructionStage::ProbingOutputs {
                outputs,
                xdg_output_manager,
                wlr_manager_done,
                ..
            } => {
                let has_xdg = xdg_output_manager.is_some();
                let output_count = outputs.len();
                let idx = if let Some(ref name) = self.args.output_name {
                    let pos = outputs
                        .iter()
                        .position(|o| o.name.as_deref() == Some(name.as_str()));
                    match pos {
                        Some(i) => Some(i),
                        None => {
                            let all_probed = outputs.iter().all(|o| o.done_count >= 1);
                            if all_probed {
                                let available: Vec<&str> =
                                    outputs.iter().filter_map(|o| o.name.as_deref()).collect();
                                tracing::error!(
                                    "Output '{}' not found. Available outputs: {:?}",
                                    name,
                                    available
                                );
                                self.errored = true;
                            }
                            None
                        }
                    }
                } else if outputs.iter().all(|o| o.done_count >= 1) {
                    if outputs.is_empty() {
                        return false;
                    }
                    Some(0)
                } else {
                    None
                };
                match idx {
                    Some(i) => {
                        let info = &outputs[i];
                        if has_xdg {
                            // xdg-output path (Sway/Hyprland) — strict checks
                            if info.done_count < 2
                                || info.name.is_none()
                                || info.transform.is_none()
                                || info.physical_size.is_none()
                                || info.logical_position.is_none()
                            {
                                return false;
                            }
                        } else {
                            // wlr-output-management path (niri) — relaxed checks
                            if info.done_count < 1 || !wlr_manager_done {
                                return false;
                            }
                            if info.transform.is_none() || info.physical_size.is_none() {
                                return false;
                            }
                            // name and logical_position can use defaults
                        }
                        (i, output_count)
                    }
                    None => return false,
                }
            }
            _ => return false,
        };

        let probing = match mem::replace(&mut self.stage, EncConstructionStage::Intermediate) {
            s @ EncConstructionStage::ProbingOutputs { .. } => s,
            other => {
                self.stage = other;
                return false;
            }
        };

        let (
            outputs,
            bound_outputs,
            output_names,
            screencopy_manager,
            dmabuf,
            dmabuf_feedback,
            _xdg_output_manager,
            _wlr_output_manager,
            _wlr_manager_done,
            _wlr_heads,
            _wlr_head_proxy_to_name,
        ) = match probing {
            EncConstructionStage::ProbingOutputs {
                outputs,
                bound_outputs,
                output_names,
                screencopy_manager,
                dmabuf,
                dmabuf_feedback,
                xdg_output_manager,
                wlr_output_manager,
                wlr_manager_done,
                wlr_heads,
                wlr_head_proxy_to_name,
            } => (
                outputs,
                bound_outputs,
                output_names,
                screencopy_manager,
                dmabuf,
                dmabuf_feedback,
                xdg_output_manager,
                wlr_output_manager,
                wlr_manager_done,
                wlr_heads,
                wlr_head_proxy_to_name,
            ),
            _ => unreachable!(),
        };
        // Destroy feedback object — prevents server-side resource leak
        if let Some(feedback) = dmabuf_feedback {
            feedback.destroy();
        }

        let info = &outputs[target_idx];
        let output_info = OutputInfo {
            name: info
                .name
                .clone()
                .or(info.wl_name.clone())
                .unwrap_or_else(|| format!("output-{}", output_names[target_idx])),
            transform: info.transform.unwrap(),
            physical_size: info.physical_size.unwrap(),
            logical_position: info.logical_position.unwrap_or((0, 0)),
        };
        let output = bound_outputs[target_idx].clone();

        let screencopy_manager = match screencopy_manager {
            Some(m) => m,
            None => {
                tracing::error!("No screencopy manager bound");
                self.errored = true;
                return false;
            }
        };
        let dmabuf = match dmabuf {
            Some(d) => d,
            None => {
                tracing::error!("No dmabuf manager bound");
                self.errored = true;
                return false;
            }
        };

        let drm_path = self.resolve_drm_path();

        let hw_device_ctx = match AvHwDevCtx::new_vaapi(&drm_path) {
            Ok(ctx) => ctx,
            Err(e) => {
                tracing::error!("Failed to create VAAPI device: {}", e);
                self.errored = true;
                return false;
            }
        };

        let cap = match S::new(&self.gm, &output, &output_info, &self.qhandle) {
            Ok(c) => c,
            Err(e) => {
                tracing::error!("Failed to create capture source: {}", e);
                self.errored = true;
                return false;
            }
        };

        tracing::info!("Selected output: {}", output_info.name);
        if self.args.output_name.is_none() && output_count > 1 {
            tracing::warn!(
                "Multiple outputs found, using '{}'. Use --output-name to select.",
                output_info.name
            );
        }
        self.stage = EncConstructionStage::EverythingButFmt {
            output_info,
            output,
            hw_device_ctx,
            cap,
            screencopy_manager,
            dmabuf,
        };

        true
    }
}

// ---------------------------------------------------------------------------
// Dispatch<WlRegistry, GlobalListContents>
// ---------------------------------------------------------------------------

impl<S: CaptureSource> Dispatch<WlRegistry, GlobalListContents> for State<S> {
    fn event(
        state: &mut Self,
        registry: &WlRegistry,
        event: wayland_client::protocol::wl_registry::Event,
        _data: &GlobalListContents,
        _conn: &wayland_client::Connection,
        qhandle: &QueueHandle<State<S>>,
    ) {
        use wayland_client::protocol::wl_registry::Event as RegistryEvent;
        tracing::debug!("Dispatch<WlRegistry>::event fired: {:?}", event);

        match event {
            RegistryEvent::Global {
                name,
                interface,
                version,
            } => match interface.as_str() {
                "zwlr_screencopy_manager_v1" => {
                    let v = version.min(3);
                    tracing::debug!("Binding zwlr_screencopy_manager_v1 v{v} (name={name})");
                    let mgr: ZwlrScreencopyManagerV1 = registry.bind(name, v, qhandle, ());
                    if let EncConstructionStage::ProbingOutputs {
                        screencopy_manager, ..
                    } = &mut state.stage
                    {
                        *screencopy_manager = Some(mgr);
                    }
                }
                "zwp_linux_dmabuf_v1" => {
                    let v = version.min(4);
                    tracing::debug!("Binding zwp_linux_dmabuf_v1 v{v} (name={name})");
                    let proxy: ZwpLinuxDmabufV1 = registry.bind(name, v, qhandle, ());
                    if let EncConstructionStage::ProbingOutputs {
                        dmabuf,
                        dmabuf_feedback,
                        ..
                    } = &mut state.stage
                    {
                        *dmabuf = Some(proxy.clone());
                        if v >= 4 {
                            let feedback = proxy.get_default_feedback(qhandle, ());
                            *dmabuf_feedback = Some(feedback);
                        }
                    }
                }
                "wl_output" => {
                    let v = version.min(4);
                    tracing::debug!("Binding wl_output v{v} (name={name})");
                    let output: WlOutput = registry.bind(name, v, qhandle, OutputId(name));
                    if let EncConstructionStage::ProbingOutputs {
                        outputs,
                        bound_outputs,
                        output_names,
                        xdg_output_manager,
                        ..
                    } = &mut state.stage
                    {
                        outputs.push(PartialOutputInfo::default());
                        bound_outputs.push(output.clone());
                        output_names.push(name);
                        if let Some(xdg_mgr) = xdg_output_manager {
                            let output_id = OutputId(name);
                            xdg_mgr.get_xdg_output(&output, qhandle, output_id);
                        }
                    }
                }
                "zxdg_output_manager_v1" => {
                    let v = version.min(3);
                    tracing::debug!("Binding zxdg_output_manager_v1 v{v} (name={name})");
                    let xdg_mgr: ZxdgOutputManagerV1 = registry.bind(name, v, qhandle, ());
                    if let EncConstructionStage::ProbingOutputs {
                        bound_outputs,
                        xdg_output_manager,
                        output_names,
                        ..
                    } = &mut state.stage
                    {
                        for (i, output) in bound_outputs.iter().enumerate() {
                            let oname = output_names.get(i).copied().unwrap_or(0);
                            let output_id = OutputId(oname);
                            xdg_mgr.get_xdg_output(output, qhandle, output_id);
                        }
                        *xdg_output_manager = Some(xdg_mgr);
                    }
                }
                "zwlr_output_manager_v1" => {
                    let v = version.min(4);
                    tracing::debug!("Binding zwlr_output_manager_v1 v{v} (name={name})");
                    let mgr: ZwlrOutputManagerV1 = registry.bind(name, v, qhandle, ());
                    if let EncConstructionStage::ProbingOutputs {
                        wlr_output_manager, ..
                    } = &mut state.stage
                    {
                        *wlr_output_manager = Some(mgr);
                    }
                }
                _ => {}
            },
            RegistryEvent::GlobalRemove { name } => {
                tracing::debug!("Global removed: name={name}");
            }
            _ => {}
        }
    }
}

// ---------------------------------------------------------------------------
// Dispatch<WlOutput, ()>
// ---------------------------------------------------------------------------

impl<S: CaptureSource> Dispatch<WlOutput, OutputId> for State<S> {
    fn event(
        state: &mut Self,
        _proxy: &WlOutput,
        event: wayland_client::protocol::wl_output::Event,
        data: &OutputId,
        _conn: &wayland_client::Connection,
        _qhandle: &QueueHandle<State<S>>,
    ) {
        use wayland_client::protocol::wl_output::Event as OutputEvent;
        use wayland_client::protocol::wl_output::Mode as WlMode;
        use wayland_client::protocol::wl_output::Transform as WlTransform;

        let OutputId(target_name) = data;
        let idx = match &state.stage {
            EncConstructionStage::ProbingOutputs { output_names, .. } => {
                output_names.iter().position(|&n| n == *target_name)
            }
            _ => None,
        };
        let idx = match idx {
            Some(i) => i,
            None => return,
        };

        match event {
            OutputEvent::Geometry {
                transform,
                physical_width,
                physical_height,
                ..
            } => {
                let t = match transform {
                    wayland_client::WEnum::Value(WlTransform::Normal) => Transform::Normal,
                    wayland_client::WEnum::Value(WlTransform::_90) => Transform::Normal90,
                    wayland_client::WEnum::Value(WlTransform::_180) => Transform::Normal180,
                    wayland_client::WEnum::Value(WlTransform::_270) => Transform::Normal270,
                    wayland_client::WEnum::Value(WlTransform::Flipped) => Transform::Flipped,
                    wayland_client::WEnum::Value(WlTransform::Flipped90) => Transform::Flipped90,
                    wayland_client::WEnum::Value(WlTransform::Flipped180) => Transform::Flipped180,
                    wayland_client::WEnum::Value(WlTransform::Flipped270) => Transform::Flipped270,
                    _ => Transform::Normal,
                };
                if let EncConstructionStage::ProbingOutputs { outputs, .. } = &mut state.stage {
                    if let Some(info) = outputs.get_mut(idx) {
                        info.transform = Some(t);
                        info.physical_size = Some((physical_width, physical_height));
                    }
                }
            }
            OutputEvent::Mode {
                width,
                height,
                flags,
                ..
            } => {
                let is_current = matches!(flags, wayland_client::WEnum::Value(WlMode::Current));
                if is_current {
                    if let EncConstructionStage::ProbingOutputs { outputs, .. } = &mut state.stage {
                        if let Some(info) = outputs.get_mut(idx) {
                            info.mode_size = Some((width, height));
                        }
                    }
                }
            }
            OutputEvent::Done => {
                if let EncConstructionStage::ProbingOutputs { outputs, .. } = &mut state.stage {
                    if let Some(info) = outputs.get_mut(idx) {
                        info.done_count += 1;
                        if info.done_count >= 1 {
                            state.try_finalize_output(idx);
                        }
                    }
                }
            }
            OutputEvent::Name { name } => {
                if let EncConstructionStage::ProbingOutputs { outputs, .. } = &mut state.stage {
                    if let Some(info) = outputs.get_mut(idx) {
                        info.wl_name = Some(name);
                    }
                }
            }
            _ => {}
        }
    }
}

// ---------------------------------------------------------------------------
// Dispatch<ZxdgOutputV1, OutputId>
// ---------------------------------------------------------------------------

impl<S: CaptureSource> Dispatch<ZxdgOutputV1, OutputId> for State<S> {
    fn event(
        state: &mut Self,
        _proxy: &ZxdgOutputV1,
        event: <ZxdgOutputV1 as Proxy>::Event,
        data: &OutputId,
        _conn: &wayland_client::Connection,
        _qhandle: &QueueHandle<State<S>>,
    ) {
        let target_name = data.0;
        let idx = match &state.stage {
            EncConstructionStage::ProbingOutputs { output_names, .. } => {
                output_names.iter().position(|&n| n == target_name)
            }
            _ => None,
        };
        let idx = match idx {
            Some(i) => i,
            None => return,
        };

        match event {
            XdgOutputEvent::Name { name } => {
                if let EncConstructionStage::ProbingOutputs { outputs, .. } = &mut state.stage {
                    if let Some(info) = outputs.get_mut(idx) {
                        info.name = Some(name);
                    }
                }
            }
            XdgOutputEvent::LogicalPosition { x, y } => {
                if let EncConstructionStage::ProbingOutputs { outputs, .. } = &mut state.stage {
                    if let Some(info) = outputs.get_mut(idx) {
                        info.logical_position = Some((x, y));
                    }
                }
            }
            XdgOutputEvent::LogicalSize { .. } => {}
            XdgOutputEvent::Done => {
                if let EncConstructionStage::ProbingOutputs { outputs, .. } = &mut state.stage {
                    if let Some(info) = outputs.get_mut(idx) {
                        info.done_count += 1;
                        if info.done_count >= 1 {
                            state.try_finalize_output(idx);
                        }
                    }
                }
            }
            _ => {}
        }
    }
}

// ---------------------------------------------------------------------------
// Dispatch<ZwpLinuxDmabufV1, ()>
// ---------------------------------------------------------------------------

impl<S: CaptureSource> Dispatch<ZwpLinuxDmabufV1, ()> for State<S> {
    fn event(
        _state: &mut Self,
        _proxy: &ZwpLinuxDmabufV1,
        event: <ZwpLinuxDmabufV1 as Proxy>::Event,
        _data: &(),
        _conn: &wayland_client::Connection,
        _qhandle: &QueueHandle<State<S>>,
    ) {
        match event {
            DmabufEvent::Format { .. } => {}
            DmabufEvent::Modifier { .. } => {}
            _ => {}
        }
    }
}

impl<S: CaptureSource> Dispatch<ZwpLinuxDmabufFeedbackV1, ()> for State<S> {
    fn event(
        state: &mut Self,
        _proxy: &ZwpLinuxDmabufFeedbackV1,
        event: <ZwpLinuxDmabufFeedbackV1 as Proxy>::Event,
        _data: &(),
        _conn: &wayland_client::Connection,
        _qhandle: &QueueHandle<State<S>>,
    ) {
        match event {
            DmabufFeedbackEvent::MainDevice { device } => {
                if device.len() >= 8 {
                    let dev_bytes: [u8; 8] = device[..8].try_into().unwrap_or([0u8; 8]);
                    let dev = u64::from_ne_bytes(dev_bytes);
                    let minor = ((dev & 0xFF) | ((dev >> 12) & 0xFFFFFF00)) as u32;
                    let path = PathBuf::from(format!("/dev/dri/renderD{}", minor));
                    if path.exists() {
                        tracing::info!(
                            "Compositor DRM device: {} (dev_t: {})",
                            path.display(),
                            dev
                        );
                        state.drm_device_from_compositor = Some(path);
                    } else {
                        tracing::warn!(
                            "Compositor reported DRM device {} (dev_t: {}) but path does not exist",
                            path.display(),
                            dev
                        );
                    }
                } else {
                    tracing::warn!(
                        "main_device event with unexpected data length: {}",
                        device.len()
                    );
                }
            }
            DmabufFeedbackEvent::FormatTable { .. } => {}
            DmabufFeedbackEvent::Done => {}
            DmabufFeedbackEvent::TrancheDone => {}
            DmabufFeedbackEvent::TrancheTargetDevice { .. } => {}
            DmabufFeedbackEvent::TrancheFormats { .. } => {}
            DmabufFeedbackEvent::TrancheFlags { .. } => {}
            _ => {}
        }
    }
}

// ---------------------------------------------------------------------------
// Dispatch<ZwpLinuxBufferParamsV1, ()>
// ---------------------------------------------------------------------------

impl<S: CaptureSource> Dispatch<ZwpLinuxBufferParamsV1, ()> for State<S> {
    fn event(
        state: &mut Self,
        proxy: &ZwpLinuxBufferParamsV1,
        event: <ZwpLinuxBufferParamsV1 as Proxy>::Event,
        _data: &(),
        _conn: &wayland_client::Connection,
        _qhandle: &QueueHandle<State<S>>,
    ) {
        match event {
            BufferParamsEvent::Created { .. } => {
                tracing::debug!("DMA-BUF buffer created");
            }
            BufferParamsEvent::Failed => {
                tracing::error!("DMA-BUF buffer creation failed");
                let taken = mem::replace(&mut state.in_flight_surface, InFlightSurface::None);
                match taken {
                    InFlightSurface::CopyQueued {
                        buffer,
                        frame,
                        ..
                    } => {
                        drop(buffer);
                        if let EncConstructionStage::Streaming { cap, .. } = &mut state.stage {
                            cap.on_done_with_frame(frame);
                        }
                    }
                    other => {
                        state.in_flight_surface = other;
                    }
                }
                proxy.destroy();
                state.errored = true;
            }
            _ => {}
        }
    }
}

// ---------------------------------------------------------------------------
// Dispatch<ZwlrScreencopyFrameV1, ()> for CapWlrScreencopy
// ---------------------------------------------------------------------------

impl Dispatch<ZwlrScreencopyFrameV1, ()> for State<CapWlrScreencopy> {
    fn event(
        state: &mut Self,
        proxy: &ZwlrScreencopyFrameV1,
        event: <ZwlrScreencopyFrameV1 as Proxy>::Event,
        _data: &(),
        _conn: &wayland_client::Connection,
        _qhandle: &QueueHandle<State<CapWlrScreencopy>>,
    ) {
        match event {
            ScreencopyFrameEvent::Buffer { .. } => {
                tracing::warn!(
                    "Received SHM Buffer event — only DMA-BUF capture is supported. Ignoring."
                );
                return;
            }
            ScreencopyFrameEvent::LinuxDmabuf {
                format,
                width,
                height,
            } => {
                tracing::debug!("Screencopy LinuxDmabuf: format={format}, {width}x{height}");
                if matches!(state.stage, EncConstructionStage::EverythingButFmt { .. }) {
                    state.negotiate_format(format, width, height);
                    if state.errored {
                        return;
                    }
                }
                if let EncConstructionStage::Streaming { cap, .. } = &mut state.stage {
                    cap.current_frame = Some(proxy.clone());
                }
                state.on_frame_allocd((), format, width, height);
            }
            ScreencopyFrameEvent::Ready {
                tv_sec_hi,
                tv_sec_lo,
                tv_nsec,
            } => {
                let tv_sec = (tv_sec_hi as u64) << 32 | tv_sec_lo as u64;
                let tv_usec = tv_nsec / 1000;
                tracing::trace!("Screencopy ready: tv_sec={tv_sec}, tv_usec={tv_usec}");
                state.on_copy_complete(tv_sec, tv_usec);
            }
            ScreencopyFrameEvent::Failed => {
                tracing::error!("Screencopy frame failed");
                state.on_copy_fail();
            }
            ScreencopyFrameEvent::Damage { .. } => {}
            _ => {}
        }
    }
}

// ---------------------------------------------------------------------------
// Dispatch<ZxdgOutputManagerV1, ()>
// ---------------------------------------------------------------------------

impl<S: CaptureSource> Dispatch<ZxdgOutputManagerV1, ()> for State<S> {
    fn event(
        _state: &mut Self,
        _proxy: &ZxdgOutputManagerV1,
        _event: <ZxdgOutputManagerV1 as Proxy>::Event,
        _data: &(),
        _conn: &wayland_client::Connection,
        _qhandle: &QueueHandle<State<S>>,
    ) {
    }
}

// ---------------------------------------------------------------------------
// Dispatch<ZwlrOutputManagerV1, ()>
// ---------------------------------------------------------------------------

impl<S: CaptureSource> Dispatch<ZwlrOutputManagerV1, ()> for State<S> {
    fn event(
        state: &mut Self,
        _proxy: &ZwlrOutputManagerV1,
        event: <ZwlrOutputManagerV1 as Proxy>::Event,
        _data: &(),
        _conn: &wayland_client::Connection,
        qhandle: &QueueHandle<State<S>>,
    ) {
        match event {
            WlrOutputManagerEvent::Head { head } => {
                let _head: ZwlrOutputHeadV1 = head;
                tracing::debug!("wlr output head advertised");
            }
            WlrOutputManagerEvent::Done { .. } => {
                if let EncConstructionStage::ProbingOutputs {
                    wlr_manager_done,
                    outputs,
                    ..
                } = &mut state.stage
                {
                    *wlr_manager_done = true;
                    let count = outputs.len();
                    for idx in 0..count {
                        state.try_finalize_output(idx);
                    }
                }
            }
            WlrOutputManagerEvent::Finished { .. } => {
                tracing::warn!("zwlr_output_manager_v1::Finished received during probing");
            }
            _ => {}
        }
    }

    event_created_child!(State<S>, ZwlrOutputManagerV1, [
        zwlr_output_manager_v1::EVT_HEAD_OPCODE => (ZwlrOutputHeadV1, ()),
    ]);
}

// ---------------------------------------------------------------------------
// Dispatch<ZwlrOutputHeadV1, ()>
// ---------------------------------------------------------------------------

impl<S: CaptureSource> Dispatch<ZwlrOutputHeadV1, ()> for State<S> {
    fn event(
        state: &mut Self,
        proxy: &ZwlrOutputHeadV1,
        event: <ZwlrOutputHeadV1 as Proxy>::Event,
        _data: &(),
        _conn: &wayland_client::Connection,
        _qhandle: &QueueHandle<State<S>>,
    ) {
        match event {
            WlrHeadEvent::Name { name } => {
                if let EncConstructionStage::ProbingOutputs {
                    wlr_heads,
                    wlr_head_proxy_to_name,
                    ..
                } = &mut state.stage
                {
                    wlr_heads
                        .entry(name.clone())
                        .or_insert(WlrHeadInfo { position: None });
                    wlr_head_proxy_to_name.insert(proxy.id(), name);
                }
            }
            WlrHeadEvent::Position { x, y } => {
                if let EncConstructionStage::ProbingOutputs {
                    wlr_heads,
                    wlr_head_proxy_to_name,
                    ..
                } = &mut state.stage
                {
                    if let Some(name) = wlr_head_proxy_to_name.get(&proxy.id()) {
                        if let Some(head) = wlr_heads.get_mut(name) {
                            head.position = Some((x, y));
                        }
                    }
                }
            }
            WlrHeadEvent::Finished { .. } => {
                tracing::debug!("zwlr_output_head_v1::Finished received");
            }
            _ => {}
        }
    }

    event_created_child!(State<S>, ZwlrOutputHeadV1, [
        zwlr_output_head_v1::EVT_MODE_OPCODE => (ZwlrOutputModeV1, ()),
    ]);
}

// ---------------------------------------------------------------------------
// Dispatch<ZwlrOutputModeV1, ()>
// ---------------------------------------------------------------------------

impl<S: CaptureSource> Dispatch<ZwlrOutputModeV1, ()> for State<S> {
    fn event(
        _state: &mut Self,
        _proxy: &ZwlrOutputModeV1,
        _event: <ZwlrOutputModeV1 as Proxy>::Event,
        _data: &(),
        _conn: &wayland_client::Connection,
        _qhandle: &QueueHandle<State<S>>,
    ) {
    }
}

// ---------------------------------------------------------------------------
// Dispatch<ZwlrScreencopyManagerV1, ()>
// ---------------------------------------------------------------------------

impl<S: CaptureSource> Dispatch<ZwlrScreencopyManagerV1, ()> for State<S> {
    fn event(
        _state: &mut Self,
        _proxy: &ZwlrScreencopyManagerV1,
        _event: <ZwlrScreencopyManagerV1 as Proxy>::Event,
        _data: &(),
        _conn: &wayland_client::Connection,
        _qhandle: &QueueHandle<State<S>>,
    ) {
    }
}

// ---------------------------------------------------------------------------
// Dispatch<WlBuffer, ()>
// ---------------------------------------------------------------------------

impl<S: CaptureSource> Dispatch<WlBuffer, ()> for State<S> {
    fn event(
        _state: &mut Self,
        _proxy: &WlBuffer,
        event: <WlBuffer as Proxy>::Event,
        _data: &(),
        _conn: &wayland_client::Connection,
        _qhandle: &QueueHandle<State<S>>,
    ) {
        if let wayland_client::protocol::wl_buffer::Event::Release = event {
            tracing::trace!("WlBuffer released");
        }
    }
}