wl-webrtc/src/state/mod.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::sync::atomic::{AtomicBool, Ordering};
use std::sync::Arc;
use std::time::Instant;

use anyhow::Result;
use wayland_client::backend::ObjectId;
use wayland_client::globals::GlobalList;
use wayland_client::protocol::wl_buffer::WlBuffer;
use wayland_client::protocol::wl_output::WlOutput;
use wayland_client::{Dispatch, QueueHandle};
use wayland_protocols::wp::linux_dmabuf::zv1::client::zwp_linux_buffer_params_v1::Flags as BufferParamsFlags;
use wayland_protocols::wp::linux_dmabuf::zv1::client::zwp_linux_dmabuf_feedback_v1::ZwpLinuxDmabufFeedbackV1;
use wayland_protocols::wp::linux_dmabuf::zv1::client::zwp_linux_dmabuf_v1::ZwpLinuxDmabufV1;
use wayland_protocols::xdg::xdg_output::zv1::client::zxdg_output_manager_v1::ZxdgOutputManagerV1;
use wayland_protocols_wlr::output_management::v1::client::zwlr_output_manager_v1::ZwlrOutputManagerV1;
use wayland_protocols_wlr::screencopy::v1::client::zwlr_screencopy_frame_v1::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, EncodedH264Frame, SwEncState};
use crate::fps_limit::FpsLimit;
use crate::stats::{FrameTimings, PipelineStats};
use crate::transform::{transpose_if_transform_transposed, Transform};
use crate::webrtc::WebRtcState;

mod dispatch;

// ---------------------------------------------------------------------------
// 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 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,
}

#[derive(Default)]
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>,
    // Pixel dimensions from Mode event — preparatory for Phase 2 resolution logic
    pub mode_size: Option<(i32, i32)>,
    pub done_count: u32,
}

/// Marker for wlr-output-management heads seen during probing; tracked by name
/// in `EncConstructionStage::ProbingOutputs.wlr_heads`.
// `pub(crate)` (not module-private): exposed via `EncConstructionStage::ProbingOutputs.wlr_heads`
// which is reached from main.rs during the wlr-screencopy probing loop.
pub(crate) struct WlrHeadInfo {}

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

// ---------------------------------------------------------------------------
// StreamingEncoder
// ---------------------------------------------------------------------------

/// Wraps the two possible encoder backends for the streaming stage.
///
/// - `Mp4(EncState)`  — hardware VAAPI encoder writing to an MP4 file
/// - `WebRtc(SwEncState)` — software encoder feeding H.264 NALUs into a WebRTC channel
pub enum StreamingEncoder {
    Mp4(EncState),
    WebRtc(SwEncState),
}

impl StreamingEncoder {
    fn frames_rgb(&self) -> &crate::avhw::AvHwFrameCtx {
        match self {
            StreamingEncoder::Mp4(enc) => enc.frames_rgb(),
            StreamingEncoder::WebRtc(enc) => enc.frames_rgb(),
        }
    }

    fn encode_frame(
        &mut self,
        hw_frame: &ffmpeg_next::frame::Video,
    ) -> anyhow::Result<crate::avhw::EncodeStages> {
        match self {
            StreamingEncoder::Mp4(enc) => enc.encode_frame(hw_frame),
            StreamingEncoder::WebRtc(enc) => enc.encode_frame(hw_frame),
        }
    }

    pub fn flush(&mut self) -> anyhow::Result<()> {
        match self {
            StreamingEncoder::Mp4(enc) => enc.flush(),
            StreamingEncoder::WebRtc(enc) => enc.flush(),
        }
    }
}

// ---------------------------------------------------------------------------
// EncConstructionStage
// ---------------------------------------------------------------------------
//
// `pub(crate)` (not `pub`): this enum leaks the private `WlrHeadInfo` type via
// its `wlr_heads` field, and the construction-stage state machine is an
// internal implementation detail. Crate-internal consumers (main.rs) get there
// via `crate::state::`; there is no need to expose this across the crate
// boundary. See Oracle audit 2026-06-28.

pub(crate) 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: WlOutput,
        enc: StreamingEncoder,
        cap: S,
        screencopy_manager: ZwlrScreencopyManagerV1,
        dmabuf: ZwpLinuxDmabufV1,
    },
    Intermediate,
}

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

pub enum InFlightSurface<S: CaptureSource> {
    None,
    AllocQueued,
    CopyQueued {
        surface: ff::frame::Video,
        // Boxed: AVDRMFrameDescriptor is ~592 bytes (4 objects + 4 layers),
        // which would balloon every InFlightSurface variant via enum alignment.
        // The box shrinks the enum to ~32 bytes regardless of variant.
        drm_map: Box<ff::ffi::AVDRMFrameDescriptor>,
        frame: S::Frame,
        buffer: WlBuffer,
    },
}

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

pub struct State<S: CaptureSource> {
    pub(crate) stage: EncConstructionStage<S>,
    pub in_flight_surface: InFlightSurface<S>,
    pub stats_start_time: Option<Instant>,
    pub stats_last_time: Option<Instant>,
    pub stats_frames: u64,
    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>,
    pub webrtc: Option<WebRtcState>,
    pub webrtc_tx: Option<crossbeam_channel::Sender<EncodedH264Frame>>,
    webrtc_rx: Option<crossbeam_channel::Receiver<EncodedH264Frame>>,
    webrtc_frames_sent: u64,
    webrtc_paused: Option<Arc<AtomicBool>>,
    stats: PipelineStats,
}

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

/// Scan /dev/dri for all available DRM render nodes (renderD*), sorted by node number.
pub(crate) fn find_drm_render_nodes() -> Vec<PathBuf> {
    let Ok(entries) = std::fs::read_dir("/dev/dri") else {
        return Vec::new();
    };

    let mut nodes: Vec<(u32, PathBuf)> = entries
        .filter_map(Result::ok)
        .filter_map(|entry| {
            let path = entry.path();
            let name = path.file_name()?.to_str()?;
            let number = name.strip_prefix("renderD")?.parse::<u32>().ok()?;
            std::fs::metadata(&path).ok()?;
            Some((number, path))
        })
        .collect();
    nodes.sort_by_key(|(number, _)| *number);
    nodes.into_iter().map(|(_, path)| path).collect()
}

/// Scan /dev/dri for the first available DRM render node (renderD*).
fn find_drm_render_node() -> Option<PathBuf> {
    find_drm_render_nodes().into_iter().next()
}

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>>) -> Result<Self> {
        let fps = args.fps;
        let drm_device = args.drm_device.as_ref().map(PathBuf::from);

        let (webrtc, webrtc_tx, webrtc_rx, webrtc_paused) = if args.port > 0 {
            let (tx, rx) = crossbeam_channel::bounded(32);
            let wrtc = WebRtcState::new(args.port, args.fps)?;
            // paused=true until first WebRTC client connects
            let paused = Arc::new(AtomicBool::new(true));
            (Some(wrtc), Some(tx), Some(rx), Some(paused))
        } else {
            (None, None, None, None)
        };

        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,
            stats_start_time: None,
            stats_last_time: None,
            stats_frames: 0,
            first_frame: true,
            fps_limit: FpsLimit::new(fps),
            args,
            errored: false,
            gm,
            qhandle,
            drm_device,
            drm_device_from_compositor: None,
            webrtc,
            webrtc_tx,
            webrtc_rx,
            webrtc_frames_sent: 0,
            webrtc_paused,
            stats: PipelineStats::new(),
        };

        // 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();

        Ok(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: _,
                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: {}", crate::avhw::ff_err(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;
        }
        // SAFETY: map_frame and surface are valid, owned AVFrame pointers from
        // av_hwframe_get/surface.alloc above. AV_HWFRAME_MAP_READ flag (0 here)
        // requests a read-only mapping. The DRM_PRIME format set above instructs
        // FFmpeg to populate data[0] with an AVDRMFrameDescriptor on success.
        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: {}", crate::avhw::ff_err(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: {}",
                        std::io::Error::last_os_error()
                    );
                    // 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: Box::new(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,
    {
        self.stats.record_capture();

        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;
                }
            };
        // PTS in 90kHz media-clock ticks (WebRTC encoder time_base = 1/90000).
        // Must match Portal path's compute_capture_pts unit. See issue #25.
        let pts = (tv_sec as i64) * 90_000 + (tv_usec as i64) * 90_000 / 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 {
            let encode_start = Instant::now();
            match enc.encode_frame(&surface) {
                Ok(stages) => {
                    let encode_elapsed = encode_start.elapsed().as_micros() as u64;
                    self.stats.record_encode(&FrameTimings {
                        scale_us: stages.scale_us,
                        transfer_us: stages.transfer_us,
                        encode_us: stages.encode_us,
                        total_us: encode_elapsed,
                        ..Default::default()
                    });
                }
                Err(e) => {
                    tracing::error!("encode_frame failed: {}", e);
                    self.errored = true;
                }
            }
        }
        self.stats_frames += 1;
        if let Some(last) = self.stats_last_time {
            if last.elapsed() >= std::time::Duration::from_secs(10) {
                let delta = self.stats_frames;
                let fps = delta as f64 / last.elapsed().as_secs_f64();
                tracing::info!(
                    frames = self.stats_frames,
                    fps = format!("{fps:.1}"),
                    "encoding stats"
                );
                self.stats_last_time = Some(std::time::Instant::now());
                self.stats_frames = 0;
            }
        } else {
            self.stats_start_time = Some(std::time::Instant::now());
            self.stats_last_time = Some(std::time::Instant::now());
        }
    }

    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 poll_webrtc(&mut self) -> Result<()> {
        let Some(ref mut wrtc) = self.webrtc else {
            return Ok(());
        };

        wrtc.handle_signaling()?;
        wrtc.poll_and_feed()?;

        let connected = wrtc.is_connected();

        if let Some(ref paused) = self.webrtc_paused {
            let was_paused = paused.load(Ordering::Relaxed);
            let now_paused = !connected;
            if was_paused && !now_paused {
                tracing::info!("WebRTC client connected, resuming encoding");
            } else if !was_paused && now_paused {
                tracing::warn!("WebRTC client disconnected, pausing encoding");
            }
            paused.store(now_paused, Ordering::Relaxed);
        }

        if let Some(ref rx) = self.webrtc_rx {
            let mut count = 0u32;
            while let Ok(enc_frame) = rx.try_recv() {
                if !connected {
                    continue;
                }
                count += 1;
                if let Err(e) = wrtc.write_h264_frame(&enc_frame.data, enc_frame.pts_ticks) {
                    tracing::debug!("WebRTC write frame error: {e}");
                }
                self.stats.record_send(0.0, None);
                self.webrtc_frames_sent = self.webrtc_frames_sent.saturating_add(1);
            }
            if count > 0 {
                tracing::debug!("WebRTC forwarded {count} frames from channel");
            }
        }

        if self.args.stats && self.stats.should_snapshot() {
            self.stats
                .set_queue_depths(0, self.webrtc_rx.as_ref().map(|r| r.len()).unwrap_or(0));
            let snap = self.stats.snapshot_and_reset();
            tracing::info!("stats: {snap}");
        }

        Ok(())
    }

    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,
                cap,
                screencopy_manager,
                dmabuf,
            } => (
                output_info,
                output,
                hw_device_ctx,
                cap,
                screencopy_manager,
                dmabuf,
            ),
            other => {
                tracing::warn!("negotiate_format: not in EverythingButFmt stage");
                self.stage = other;
                return;
            }
        };
        let (output_info, output, hw_device_ctx, 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 = if let Some(ref tx) = self.webrtc_tx {
            let (enc_w, enc_h) = transpose_if_transform_transposed(
                output_info.transform,
                width as i32,
                height as i32,
            );
            let actual_gop_size = self.args.gop_size.unwrap_or((fps * 2).max(20));
            match SwEncState::new_webrtc(
                &drm_path,
                width,
                height,
                enc_w as u32,
                enc_h as u32,
                fps,
                bitrate,
                actual_gop_size,
                tx.clone(),
                self.webrtc_paused
                    .as_ref()
                    .expect("webrtc_paused must exist when webrtc_tx exists")
                    .clone(),
            ) {
                Ok(enc) => StreamingEncoder::WebRtc(enc),
                Err(e) => {
                    tracing::error!("SwEncState::new_webrtc failed: {}", e);
                    self.errored = true;
                    return;
                }
            }
        } else {
            let output_path = self
                .args
                .output
                .as_deref()
                .expect("output required for MP4 mode");
            match crate::avhw::create_encoder(
                &drm_path,
                Path::new(output_path),
                width,
                height,
                fps,
                output_info.transform,
                self.args.bitrate,
                self.args.gop_size,
                Some(hw_device_ctx),
            ) {
                Ok(enc) => StreamingEncoder::Mp4(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,
            enc,
            cap,
            screencopy_manager,
            dmabuf,
        };
    }

    fn try_finalize_output(&mut self, _idx: usize) -> bool {
        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 {
                            // done_count >= 2 implies physical_size and logical_position
                            // already arrived (Wayland: Geometry/Mode/Position fire before Done).
                            if info.done_count < 2
                                || info.name.is_none()
                                || info.transform.is_none()
                            {
                                return false;
                            }
                        } else {
                            // done_count >= 1 implies transform arrived (Geometry precedes Done).
                            if info.done_count < 1 || !wlr_manager_done || info.transform.is_none()
                            {
                                return false;
                            }
                        }
                        (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(),
        };
        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
    }
}