df_displmgr 0.1.0

//! Optimized Wayland object management for wlroots-based compositors.
//! Integrated with color-management-v1 for HDR and color space control.

use wayland_client::{
    Connection, Dispatch, QueueHandle, Proxy,
    protocol::{wl_registry, wl_callback, wl_output::Transform}
};
use wayland_protocols_wlr::output_management::v1::client::{
    zwlr_output_manager_v1, zwlr_output_head_v1, zwlr_output_mode_v1, zwlr_output_configuration_v1
};
use async_trait::async_trait;
use tokio::sync::oneshot;

use crate::error::{DisplayError, DisplayResult};
use crate::types::{OutputState, DisplayRotation, HdrState, HdrMode, VideoMode, Extent2D, DisplayId};
use crate::traits::{UniversalTopology, OutputEditable};

/// Internal state for tracking Wayland objects and async synchronization.
pub struct WlrInternalState {
    pub manager: Option<zwlr_output_manager_v1::ZwlrOutputManagerV1>,
    pub heads: Vec<HeadInfo>,
    pub serial: u32,
    pub sync_tx: Option<oneshot::Sender<()>>,
    pub persist: bool,
}

/// Represents a physical output device and its capabilities.
pub struct HeadInfo {
    pub handle: zwlr_output_head_v1::ZwlrOutputHeadV1,
    pub current_state: OutputState,
    pub modes: Vec<VideoMode>,
}

pub struct WlrTopology {
    connection: Connection,
    state: WlrInternalState,
    event_queue: wayland_client::EventQueue<WlrInternalState>,
}

impl WlrTopology {
    /// Wake any currently inactive outputs and return a list of their IDs
    /// which can later be passed to `restore_inactive_outputs` to revert.
    pub fn snapshot_wake_inactive_outputs(&mut self, keep_active: Option<&DisplayId>) -> Vec<DisplayId> {
        let mut activated = Vec::new();
        for head in &mut self.state.heads {
            let id = head.current_state.identity.id.clone();
            if !head.current_state.enabled {
                if let Some(k) = keep_active {
                    if &id == k {
                        continue;
                    }
                }
                head.current_state.enabled = true;
                activated.push(id);
            }
        }
        activated
    }

    /// Restore previously deactivated outputs by ID. Does not disable an
    /// output that was explicitly kept active during snapshot.
    pub fn restore_inactive_outputs(&mut self, inactive_ids: &[DisplayId]) {
        for id in inactive_ids {
            if let Some(head) = self.state.heads.iter_mut().find(|h| &h.current_state.identity.id == id) {
                head.current_state.enabled = false;
            }
        }
    }
}

pub struct WlrOutputEditor<'a> {
    head: &'a mut HeadInfo,
    all_heads: &'a [HeadInfo],
}

fn wl_transform_to_rotation(transform: Transform) -> DisplayRotation {
    match transform {
        Transform::Normal => DisplayRotation::Rotate0,
        Transform::Rotate90 => DisplayRotation::Rotate90,
        Transform::Rotate180 => DisplayRotation::Rotate180,
        Transform::Rotate270 => DisplayRotation::Rotate270,
        _ => DisplayRotation::Rotate0,
    }
}

fn rotation_to_wl_transform(rotation: DisplayRotation) -> Transform {
    match rotation {
        DisplayRotation::Rotate0 => Transform::Normal,
        DisplayRotation::Rotate90 => Transform::Rotate90,
        DisplayRotation::Rotate180 => Transform::Rotate180,
        DisplayRotation::Rotate270 => Transform::Rotate270,
    }
}

// --- Dispatch Implementations (Registry, Manager, Head, Mode, Callback) ---

impl Dispatch<wl_registry::WlRegistry, ()> for WlrInternalState {
    fn event(state: &mut Self, proxy: &wl_registry::WlRegistry, event: wl_registry::Event, _: &(), _: &Connection, qh: &QueueHandle<Self>) {
        if let wl_registry::Event::Global { name, interface, version } = event {
            if interface == zwlr_output_manager_v1::ZwlrOutputManagerV1::interface().name {
                state.manager = Some(proxy.bind::<zwlr_output_manager_v1::ZwlrOutputManagerV1, _, _>(name, version, qh, ()));
            }
        }
    }
}

impl Dispatch<zwlr_output_manager_v1::ZwlrOutputManagerV1, ()> for WlrInternalState {
    fn event(state: &mut Self, _: &zwlr_output_manager_v1::ZwlrOutputManagerV1, event: zwlr_output_manager_v1::Event, _: &(), _: &Connection, _: &QueueHandle<Self>) {
        if let zwlr_output_manager_v1::Event::Done { serial } = event {
            state.serial = serial;
        }
    }
}

impl Dispatch<zwlr_output_head_v1::ZwlrOutputHeadV1, ()> for WlrInternalState {
    fn event(state: &mut Self, proxy: &zwlr_output_head_v1::ZwlrOutputHeadV1, event: zwlr_output_head_v1::Event, _: &(), _: &Connection, qh: &QueueHandle<Self>) {
        let head_idx = if let Some(idx) = state.heads.iter().position(|h| h.handle == *proxy) {
            idx
        } else {
            state.heads.push(HeadInfo {
                handle: proxy.clone(),
                current_state: OutputState::default(),
                modes: Vec::new(),
            });
            state.heads.len() - 1
        };

        let head = &mut state.heads[head_idx];
        match event {
            zwlr_output_head_v1::Event::Name { name } => {
                head.current_state.identity.id = DisplayId(name.clone());
                head.current_state.identity.monitor_name = name;
            }
            zwlr_output_head_v1::Event::Position { x, y } => {
                head.current_state.geometry.origin.x = x;
                head.current_state.geometry.origin.y = y;
            }
            zwlr_output_head_v1::Event::Scale { scale } => {
                head.current_state.scale = scale as f64;
            }
            zwlr_output_head_v1::Event::Transform { transform } => {
                head.current_state.rotation = wl_transform_to_rotation(transform);
            }
            zwlr_output_head_v1::Event::Mode { mode } => {
                qh.make_handle().insert_instance::<zwlr_output_mode_v1::ZwlrOutputModeV1, usize>(&mode, head_idx);
            }
            zwlr_output_head_v1::Event::Enabled { enabled } => {
                head.current_state.enabled = enabled != 0;
            }
            _ => {}
        }
    }
}

impl Dispatch<zwlr_output_mode_v1::ZwlrOutputModeV1, usize> for WlrInternalState {
    fn event(state: &mut Self, proxy: &zwlr_output_mode_v1::ZwlrOutputModeV1, event: zwlr_output_mode_v1::Event, head_idx: &usize, _: &Connection, _: &QueueHandle<Self>) {
        if let Some(head) = state.heads.get_mut(*head_idx) {
            match event {
                zwlr_output_mode_v1::Event::Size { width, height } => {
                    let mode = VideoMode {
                        resolution: Extent2D { width, height },
                        refresh_rate: 0,
                    };
                    if !head.modes.iter().any(|m| m.resolution == mode.resolution) {
                        head.modes.push(mode);
                    }
                }
                zwlr_output_mode_v1::Event::Refresh { refresh } => {
                    if let Some(mode) = head.modes.last_mut() {
                        mode.refresh_rate = refresh;
                    }
                }
                zwlr_output_mode_v1::Event::Preferred => {
                    // Preferred mode handling is not yet required for this simplified model.
                }
                _ => {}
            }
        }
    }
}

impl Dispatch<wl_callback::WlCallback, ()> for WlrInternalState {
    fn event(state: &mut Self, _: &wl_callback::WlCallback, _: wl_callback::Event, _: &(), _: &Connection, _: &QueueHandle<Self>) {
        if let Some(tx) = state.sync_tx.take() {
            let _ = tx.send(());
        }
    }
}

impl Dispatch<zwlr_output_configuration_v1::ZwlrOutputConfigurationV1, ()> for WlrInternalState {}

// --- Trait Implementations ---

#[async_trait]
impl UniversalTopology for WlrTopology {
    fn acquire() -> DisplayResult<Self> {
        let conn = Connection::connect_to_env().map_err(|_| DisplayError::ConnectionFailed)?;
        let event_queue = conn.new_event_queue();
        let qh = event_queue.handle();

        let state = WlrInternalState {
            manager: None,
            heads: Vec::new(),
            serial: 0,
            sync_tx: None,
            persist: false,
        };

        let _registry = conn.display().get_registry(&qh, ());
        let mut wrapper = WlrTopology { connection: conn, state, event_queue };
        
        wrapper.event_queue.roundtrip(&mut wrapper.state)
            .map_err(|e| DisplayError::BackendError(e.to_string()))?;

        Ok(wrapper)
    }

    fn get_outputs(&self) -> Vec<OutputState> {
        self.state.heads.iter().map(|h| h.current_state.clone()).collect()
    }

    fn edit_output(&mut self, id: &DisplayId) -> DisplayResult<Box<dyn OutputEditable + '_>> {
        let head_idx = self.state.heads.iter()
            .position(|h| h.current_state.identity.id == *id)
            .ok_or_else(|| DisplayError::NotFound(id.clone()))?;
        
        Ok(Box::new(WlrOutputEditor { 
            head: &mut self.state.heads[head_idx],
            all_heads: &self.state.heads 
        }))
    }

    fn set_persistence(&mut self, enabled: bool) -> &mut Self {
        self.state.persist = enabled;
        self
    }

    async fn validate(&self) -> DisplayResult<()> {
        // On Wayland, validation happens during the test phase of the configuration application.
        // For the trait, we assume staged changes are valid until commit.
        Ok(())
    }

    async fn commit(&mut self) -> DisplayResult<()> {
        let manager = self.state.manager.as_ref()
            .ok_or_else(|| DisplayError::BackendError("Output manager missing".into()))?;

        let qh = self.event_queue.handle();
        let config = manager.create_configuration(self.state.serial, &qh, ());

        for head in &self.state.heads {
            if head.current_state.enabled {
                let head_config = config.enable_head(&head.handle, &qh, ());
                head_config.set_position(head.current_state.geometry.origin.x, head.current_state.geometry.origin.y);
                head_config.set_scale(head.current_state.scale as f32);
                head_config.set_transform(rotation_to_wl_transform(head.current_state.rotation));
            } else {
                config.disable_head(&head.handle);
            }

            if head.current_state.hdr_state == HdrState::Enabled {
                return Err(DisplayError::BackendError("HDR is not supported in this simplified WLR path".into()));
            }
        }

        config.apply();

        let (tx, rx) = oneshot::channel();
        self.state.sync_tx = Some(tx);
        self.connection.display().sync(&qh, ());

        while self.state.sync_tx.is_some() {
            self.event_queue.blocking_dispatch(&mut self.state)
                .map_err(|e| DisplayError::BackendError(e.to_string()))?;
        }

        rx.await.map_err(|_| DisplayError::BackendError("Sync failed during commit".into()))?;
        Ok(())
    }
}

impl<'a> OutputEditable for WlrOutputEditor<'a> {
    fn set_rotation(&mut self, rotation: DisplayRotation) -> DisplayResult<&mut dyn OutputEditable> {
        self.head.current_state.rotation = rotation;
        Ok(self)
    }

    fn set_resolution(&mut self, extent: Extent2D) -> DisplayResult<&mut dyn OutputEditable> {
        self.head.current_state.geometry.size = extent;
        Ok(self)
    }

    fn set_position(&mut self, position: crate::types::Point2D) -> DisplayResult<&mut dyn OutputEditable> {
        self.head.current_state.geometry.origin = position;
        Ok(self)
    }

    fn set_primary(&mut self) -> DisplayResult<&mut dyn OutputEditable> {
        self.head.current_state.is_primary = true;
        Ok(self)
    }

    fn set_refresh_rate(&mut self, rate: u32) -> DisplayResult<&mut dyn OutputEditable> {
        self.head.current_state.refresh_rate = rate;
        Ok(self)
    }

    fn set_hdr(&mut self, state: HdrState, mode: HdrMode) -> DisplayResult<&mut dyn OutputEditable> {
        self.head.current_state.hdr_state = state;
        self.head.current_state.hdr_mode = mode;
        Ok(self)
    }

    fn set_scale(&mut self, scale: f64) -> DisplayResult<&mut dyn OutputEditable> {
        self.head.current_state.scale = scale;
        Ok(self)
    }

    fn set_enabled(&mut self, enabled: bool) -> DisplayResult<&mut dyn OutputEditable> {
        self.head.current_state.enabled = enabled;
        Ok(self)
    }

    fn clone_from(&mut self, source_id: &DisplayId) -> DisplayResult<&mut dyn OutputEditable> {
        let source_state = self.all_heads.iter()
            .find(|h| h.current_state.identity.id == *source_id)
            .map(|h| &h.current_state)
            .ok_or_else(|| DisplayError::NotFound(source_id.clone()))?;

        self.head.current_state.geometry = source_state.geometry;
        self.head.current_state.refresh_rate = source_state.refresh_rate;
        self.head.current_state.scale = source_state.scale;
        self.head.current_state.rotation = source_state.rotation;
        Ok(self)
    }

    fn get_state(&self) -> OutputState {
        self.head.current_state.clone()
    }
}