lightbar/src/modules/tray.rs

use std::{
    collections::{BTreeMap, BTreeSet},
    sync::{
        Arc, Mutex,
        atomic::{AtomicBool, Ordering},
    },
    thread::{self, JoinHandle},
    time::Duration,
};

use anyhow::{Context, Result, anyhow};
use calloop::channel::Sender;
use futures_lite::StreamExt;
use tokio::sync::mpsc;
use zbus::{
    Connection, MatchRule, MessageStream, Proxy,
    fdo::{DBusProxy, NameOwnerChanged},
    message::{Header, Type},
    names::BusName,
    object_server::SignalEmitter,
};

use crate::{
    config::TrayModule,
    model::{IconPixmap, ModuleEvent, ModuleSnapshot, PopupContent, PopupModel, PopupRow, Segment},
};

const WATCHER_SERVICE: &str = "org.kde.StatusNotifierWatcher";
const WATCHER_PATH: &str = "/StatusNotifierWatcher";
const WATCHER_INTERFACE: &str = "org.kde.StatusNotifierWatcher";
const ITEM_INTERFACE: &str = "org.kde.StatusNotifierItem";

#[derive(Debug, Clone, Eq, Ord, PartialEq, PartialOrd)]
struct TrayAddress {
    service: String,
    path: String,
}

impl TrayAddress {
    fn parse(registered: &str) -> Option<Self> {
        if registered.starts_with('/') {
            return None;
        }
        let (service, path) = registered.find('/').map_or_else(
            || (registered, "/StatusNotifierItem"),
            |index| (&registered[..index], &registered[index..]),
        );
        (!service.is_empty() && path.starts_with('/')).then(|| Self {
            service: service.to_owned(),
            path: path.to_owned(),
        })
    }

    fn id(&self) -> String {
        format!("{}|{}", self.service, self.path)
    }

    fn from_id(id: &str) -> Option<Self> {
        let (service, path) = id.split_once('|')?;
        (!service.is_empty() && path.starts_with('/')).then(|| Self {
            service: service.to_owned(),
            path: path.to_owned(),
        })
    }

    fn registered_name(&self) -> String {
        if self.path == "/StatusNotifierItem" {
            self.service.clone()
        } else {
            format!("{}{}", self.service, self.path)
        }
    }
}

#[derive(Debug, Copy, Clone)]
enum RegistryEvent {
    Changed,
}

#[derive(Debug, Clone, Eq, PartialEq)]
struct ItemRegistration {
    owner: String,
    service: String,
}

type ItemRegistry = Arc<Mutex<BTreeMap<String, ItemRegistration>>>;

#[derive(Debug)]
struct Watcher {
    items: ItemRegistry,
    hosts: Arc<Mutex<BTreeSet<String>>>,
    changed: mpsc::Sender<RegistryEvent>,
}

#[zbus::interface(name = "org.kde.StatusNotifierWatcher")]
impl Watcher {
    async fn register_status_notifier_item(
        &self,
        service_or_path: String,
        #[zbus(header)] header: Header<'_>,
        #[zbus(signal_emitter)] emitter: SignalEmitter<'_>,
    ) -> zbus::fdo::Result<()> {
        let sender = header
            .sender()
            .map(ToString::to_string)
            .ok_or_else(|| zbus::fdo::Error::Failed("registration has no sender".to_owned()))?;
        let owner = sender.clone();
        let address = if service_or_path.starts_with('/') {
            TrayAddress {
                service: sender,
                path: service_or_path,
            }
        } else {
            TrayAddress {
                service: service_or_path,
                path: "/StatusNotifierItem".to_owned(),
            }
        };
        let registered = address.registered_name();
        let registration = ItemRegistration {
            owner,
            service: address.service,
        };
        let previous = self
            .items
            .lock()
            .map_err(|_| zbus::fdo::Error::Failed("watcher lock poisoned".to_owned()))?
            .insert(registered.clone(), registration.clone());
        let inserted = previous.is_none();
        let changed = previous.as_ref() != Some(&registration);
        if inserted {
            Self::status_notifier_item_registered(&emitter, &registered).await?;
        }
        if changed {
            let _ = self.changed.try_send(RegistryEvent::Changed);
        }
        Ok(())
    }

    async fn register_status_notifier_host(
        &self,
        service: String,
        #[zbus(signal_emitter)] emitter: SignalEmitter<'_>,
    ) -> zbus::fdo::Result<()> {
        let inserted = self
            .hosts
            .lock()
            .map_err(|_| zbus::fdo::Error::Failed("watcher lock poisoned".to_owned()))?
            .insert(service.clone());
        if inserted {
            Self::status_notifier_host_registered(&emitter).await?;
            let _ = self.changed.try_send(RegistryEvent::Changed);
        }
        Ok(())
    }

    #[zbus(property, name = "RegisteredStatusNotifierItems")]
    fn registered_status_notifier_items(&self) -> Vec<String> {
        self.items
            .lock()
            .map_or_else(|_| Vec::new(), |items| items.keys().cloned().collect())
    }

    #[zbus(property, name = "IsStatusNotifierHostRegistered")]
    fn is_status_notifier_host_registered(&self) -> bool {
        self.hosts.lock().is_ok_and(|hosts| !hosts.is_empty())
    }

    #[zbus(property, name = "ProtocolVersion")]
    fn protocol_version(&self) -> i32 {
        0
    }

    #[zbus(signal, name = "StatusNotifierItemRegistered")]
    async fn status_notifier_item_registered(
        emitter: &SignalEmitter<'_>,
        service: &str,
    ) -> zbus::Result<()>;

    #[zbus(signal, name = "StatusNotifierItemUnregistered")]
    async fn status_notifier_item_unregistered(
        emitter: &SignalEmitter<'_>,
        service: &str,
    ) -> zbus::Result<()>;

    #[zbus(signal, name = "StatusNotifierHostRegistered")]
    async fn status_notifier_host_registered(emitter: &SignalEmitter<'_>) -> zbus::Result<()>;
}

pub fn spawn(
    module: String,
    settings: TrayModule,
    sender: Sender<ModuleEvent>,
    stop: Arc<AtomicBool>,
) -> JoinHandle<()> {
    thread::Builder::new()
        .name(format!("lightbar-tray-{module}"))
        .spawn(move || {
            let runtime = tokio::runtime::Builder::new_current_thread()
                .enable_all()
                .build();
            let result = runtime
                .context("could not create tray runtime")
                .and_then(|runtime| runtime.block_on(run(module.clone(), settings, sender, stop)));
            if let Err(error) = result {
                tracing::warn!(module, error = %format_args!("{error:#}"), "tray stopped");
            }
        })
        .expect("tray worker thread")
}

async fn run(
    module: String,
    settings: TrayModule,
    sender: Sender<ModuleEvent>,
    stop: Arc<AtomicBool>,
) -> Result<()> {
    let connection = Connection::session()
        .await
        .context("could not connect to the session D-Bus")?;
    let (registry_sender, mut registry_receiver) = mpsc::channel(32);
    let owned_items = ensure_watcher(&connection, registry_sender).await?;

    let host_name = format!("org.kde.StatusNotifierHost-{}-1", std::process::id());
    connection
        .request_name(host_name.as_str())
        .await
        .context("could not own the StatusNotifierHost bus name")?;
    let watcher = watcher_proxy(&connection).await?;
    watcher
        .call::<_, _, ()>("RegisterStatusNotifierHost", &(host_name.clone()))
        .await
        .context("watcher rejected StatusNotifierHost registration")?;

    let watcher_rule = MatchRule::builder()
        .msg_type(Type::Signal)
        .interface(WATCHER_INTERFACE)?
        .build();
    let item_rule = MatchRule::builder()
        .msg_type(Type::Signal)
        .interface(ITEM_INTERFACE)?
        .build();
    let properties_rule = MatchRule::builder()
        .msg_type(Type::Signal)
        .interface("org.freedesktop.DBus.Properties")?
        .member("PropertiesChanged")?
        .add_arg(ITEM_INTERFACE)?
        .build();
    let owner_rule = MatchRule::builder()
        .msg_type(Type::Signal)
        .sender("org.freedesktop.DBus")?
        .interface("org.freedesktop.DBus")?
        .member("NameOwnerChanged")?
        .build();
    let mut watcher_signals = MessageStream::for_match_rule(watcher_rule, &connection, Some(64))
        .await
        .context("could not subscribe to tray watcher changes")?;
    let mut item_signals = MessageStream::for_match_rule(item_rule, &connection, Some(128))
        .await
        .context("could not subscribe to tray item changes")?;
    let mut property_signals =
        MessageStream::for_match_rule(properties_rule, &connection, Some(128))
            .await
            .context("could not subscribe to tray property changes")?;
    let mut owner_signals = MessageStream::for_match_rule(owner_rule, &connection, Some(128))
        .await
        .context("could not subscribe to D-Bus owner changes")?;
    let mut reconcile = tokio::time::interval(Duration::from_secs(30));
    let mut stop_check = tokio::time::interval(Duration::from_millis(250));

    publish(&connection, &module, &settings, &sender).await;
    loop {
        tokio::select! {
            _ = stop_check.tick() => {
                if stop.load(Ordering::Acquire) {
                    break;
                }
                continue;
            }
            _ = reconcile.tick() => {}
            event = registry_receiver.recv() => {
                if event.is_none() {
                    break;
                }
            }
            event = watcher_signals.next() => {
                if event.is_none() {
                    break;
                }
            }
            event = item_signals.next() => {
                if event.is_none() {
                    break;
                }
            }
            event = property_signals.next() => {
                if event.is_none() {
                    break;
                }
            }
            event = owner_signals.next() => {
                let Some(event) = event else {
                    break;
                };
                let Some(vanished) = event.ok().and_then(vanished_bus_name) else {
                    continue;
                };
                let Some(items) = &owned_items else {
                    continue;
                };
                match unregister_vanished_items(&connection, items, &vanished).await {
                    Ok(0) => continue,
                    Ok(count) => {
                        tracing::debug!(%vanished, count, "removed vanished tray items");
                    }
                    Err(error) => {
                        tracing::warn!(%vanished, %error, "could not unregister vanished tray items");
                        continue;
                    }
                }
            }
        }
        // Collapse icon/property signal bursts into a single bounded refresh.
        tokio::time::sleep(Duration::from_millis(30)).await;
        publish(&connection, &module, &settings, &sender).await;
    }
    Ok(())
}

async fn ensure_watcher(
    connection: &Connection,
    changed: mpsc::Sender<RegistryEvent>,
) -> Result<Option<ItemRegistry>> {
    if watcher_proxy(connection)
        .await?
        .get_property::<i32>("ProtocolVersion")
        .await
        .is_ok()
    {
        tracing::debug!("using existing StatusNotifierWatcher");
        return Ok(None);
    }
    let items = Arc::new(Mutex::new(BTreeMap::new()));
    let watcher = Watcher {
        items: items.clone(),
        hosts: Arc::new(Mutex::new(BTreeSet::new())),
        changed,
    };
    connection
        .object_server()
        .at(WATCHER_PATH, watcher)
        .await
        .context("could not export StatusNotifierWatcher")?;
    connection
        .request_name(WATCHER_SERVICE)
        .await
        .context("could not own StatusNotifierWatcher")?;
    tracing::info!("providing StatusNotifierWatcher");
    Ok(Some(items))
}

fn vanished_bus_name(message: zbus::Message) -> Option<String> {
    let signal = NameOwnerChanged::from_message(message)?;
    let args = signal.args().ok()?;
    args.new_owner()
        .as_ref()
        .is_none()
        .then(|| args.name().as_str().to_owned())
}

async fn unregister_vanished_items(
    connection: &Connection,
    items: &ItemRegistry,
    vanished: &str,
) -> Result<usize> {
    let removed = {
        let mut items = items
            .lock()
            .map_err(|_| anyhow!("watcher item registry lock poisoned"))?;
        remove_registrations_for_name(&mut items, vanished)
    };
    if removed.is_empty() {
        return Ok(0);
    }

    let emitter = SignalEmitter::new(connection, WATCHER_PATH)?;
    for registered in &removed {
        Watcher::status_notifier_item_unregistered(&emitter, registered).await?;
    }
    Ok(removed.len())
}

fn remove_registrations_for_name(
    items: &mut BTreeMap<String, ItemRegistration>,
    vanished: &str,
) -> Vec<String> {
    let mut removed = Vec::new();
    items.retain(|registered, registration| {
        let keep = registration.owner != vanished && registration.service != vanished;
        if !keep {
            removed.push(registered.clone());
        }
        keep
    });
    removed
}

async fn watcher_proxy(connection: &Connection) -> Result<Proxy<'_>> {
    Proxy::new(connection, WATCHER_SERVICE, WATCHER_PATH, WATCHER_INTERFACE)
        .await
        .map_err(Into::into)
}

async fn publish(
    connection: &Connection,
    module: &str,
    settings: &TrayModule,
    sender: &Sender<ModuleEvent>,
) {
    match tokio::time::timeout(
        Duration::from_secs(2),
        collect_snapshot(connection, settings),
    )
    .await
    {
        Ok(Ok(snapshot)) => {
            let _ = sender.send(ModuleEvent {
                module: module.to_owned(),
                snapshot,
            });
        }
        Ok(Err(error)) => tracing::warn!(%error, "could not refresh tray"),
        Err(_) => tracing::warn!("tray refresh exceeded two seconds"),
    }
}

async fn collect_snapshot(
    connection: &Connection,
    settings: &TrayModule,
) -> Result<ModuleSnapshot> {
    let registered = watcher_proxy(connection)
        .await?
        .get_property::<Vec<String>>("RegisteredStatusNotifierItems")
        .await
        .context("could not read registered tray items")?;
    let mut addresses: BTreeSet<_> = registered
        .iter()
        .filter_map(|item| TrayAddress::parse(item))
        .collect();
    let dbus = DBusProxy::new(connection).await?;
    let mut segments = Vec::new();
    let mut rows = Vec::new();
    for address in &addresses {
        let bus_name = match BusName::try_from(address.service.as_str()) {
            Ok(name) => name,
            Err(error) => {
                tracing::debug!(service = %address.service, %error, "ignoring invalid tray service name");
                continue;
            }
        };
        match dbus.name_has_owner(bus_name).await {
            Ok(false) => {
                tracing::debug!(service = %address.service, "ignoring vanished tray item");
                continue;
            }
            Ok(true) => {}
            Err(error) => {
                tracing::debug!(service = %address.service, %error, "could not verify tray item owner");
            }
        }
        match read_item(connection, address, settings).await {
            Ok(Some((segment, row))) => {
                segments.push(segment);
                rows.push(row);
            }
            Ok(None) => {}
            Err(error) => {
                tracing::debug!(service = %address.service, %error, "ignoring broken tray item");
            }
        }
    }
    // Deduplicate watcher bugs and ensure the set cannot grow while iterating.
    addresses.clear();
    Ok(ModuleSnapshot {
        visible: !segments.is_empty(),
        segments,
        tooltip: (!rows.is_empty()).then(|| {
            rows.iter()
                .map(|row| row.label.as_str())
                .collect::<Vec<_>>()
                .join("\n")
        }),
        popup: (!rows.is_empty()).then(|| PopupModel {
            title: "System tray".to_owned(),
            content: PopupContent::Rows(rows),
        }),
        stale: false,
    })
}

async fn read_item(
    connection: &Connection,
    address: &TrayAddress,
    settings: &TrayModule,
) -> Result<Option<(Segment, PopupRow)>> {
    let proxy = Proxy::new(
        connection,
        address.service.as_str(),
        address.path.as_str(),
        ITEM_INTERFACE,
    )
    .await?;
    let status = proxy
        .get_property::<String>("Status")
        .await
        .context("tray item has no readable Status property")?;
    if status == "Passive" {
        return Ok(None);
    }
    let title = proxy
        .get_property::<String>("Title")
        .await
        .unwrap_or_else(|_| address.service.clone());
    let icon_name = proxy
        .get_property::<String>("IconName")
        .await
        .unwrap_or_default();
    let pixmaps = proxy
        .get_property::<Vec<(i32, i32, Vec<u8>)>>("IconPixmap")
        .await
        .unwrap_or_default();
    let icon = choose_pixmap(pixmaps, settings)?;
    let mut segment = Segment::new(if icon.is_some() {
        String::new()
    } else if icon_name.is_empty() {
        "●".to_owned()
    } else {
        // Keep a deterministic, font-renderable fallback when an item only
        // exports a theme icon name. Pixmap-exporting items use the real icon.
        "●".to_owned()
    });
    segment.id = Some(address.id());
    segment.state = if status == "NeedsAttention" {
        "urgent".to_owned()
    } else {
        "normal".to_owned()
    };
    segment.icon = icon;
    let row = PopupRow {
        id: Some(address.id()),
        label: title,
        detail: (!icon_name.is_empty()).then_some(icon_name),
        state: segment.state.clone(),
        enabled: true,
    };
    Ok(Some((segment, row)))
}

fn choose_pixmap(
    pixmaps: Vec<(i32, i32, Vec<u8>)>,
    settings: &TrayModule,
) -> Result<Option<IconPixmap>> {
    let target = i64::from(settings.icon_size);
    let selected = pixmaps
        .into_iter()
        .filter(|(width, height, bytes)| {
            if *width <= 0 || *height <= 0 || *width > 512 || *height > 512 {
                return false;
            }
            let expected = i64::from(*width)
                .checked_mul(i64::from(*height))
                .and_then(|pixels| pixels.checked_mul(4));
            expected.is_some_and(|expected| {
                usize::try_from(expected).ok() == Some(bytes.len())
                    && bytes.len() <= settings.pixmap_max_bytes
            })
        })
        .min_by_key(|(width, height, _)| {
            (i64::from(*width) - target).abs() + (i64::from(*height) - target).abs()
        });
    let Some((width, height, argb)) = selected else {
        return Ok(None);
    };
    let mut pixels = Vec::with_capacity(argb.len());
    for pixel in argb.chunks_exact(4) {
        let alpha = pixel[0];
        let premultiply = |channel: u8| {
            u8::try_from((u16::from(channel) * u16::from(alpha) + 127) / 255).unwrap_or(u8::MAX)
        };
        pixels.extend_from_slice(&[
            premultiply(pixel[3]),
            premultiply(pixel[2]),
            premultiply(pixel[1]),
            alpha,
        ]);
    }
    Ok(Some(IconPixmap {
        width: u32::try_from(width).map_err(|_| anyhow!("negative tray pixmap width"))?,
        height: u32::try_from(height).map_err(|_| anyhow!("negative tray pixmap height"))?,
        pixels,
    }))
}

pub fn activate_tray_item(id: &str, method: &'static str, x: i32, y: i32) -> Result<()> {
    let address = TrayAddress::from_id(id).context("invalid tray item identity")?;
    thread::Builder::new()
        .name("lightbar-tray-action".to_owned())
        .spawn(move || {
            let runtime = tokio::runtime::Builder::new_current_thread()
                .enable_all()
                .build();
            let result: Result<()> = match runtime {
                Ok(runtime) => runtime
                    .block_on(async move {
                        let connection = Connection::session().await?;
                        let proxy =
                            Proxy::new(&connection, address.service, address.path, ITEM_INTERFACE)
                                .await?;
                        tokio::time::timeout(
                            Duration::from_secs(2),
                            proxy.call::<_, _, ()>(method, &(x, y)),
                        )
                        .await
                        .map_err(|_| zbus::Error::Failure("tray action timed out".to_owned()))??;
                        Ok::<_, zbus::Error>(())
                    })
                    .map_err(Into::into),
                Err(error) => Err(error.into()),
            };
            if let Err(error) = result {
                tracing::warn!(%error, method, "tray action failed");
            }
        })
        .context("could not start tray action")?;
    Ok(())
}

pub fn scroll_tray_item(id: &str, delta: i32, orientation: &'static str) -> Result<()> {
    let address = TrayAddress::from_id(id).context("invalid tray item identity")?;
    thread::Builder::new()
        .name("lightbar-tray-scroll".to_owned())
        .spawn(move || {
            let runtime = tokio::runtime::Builder::new_current_thread()
                .enable_all()
                .build();
            let result: Result<()> = match runtime {
                Ok(runtime) => runtime
                    .block_on(async move {
                        let connection = Connection::session().await?;
                        let proxy =
                            Proxy::new(&connection, address.service, address.path, ITEM_INTERFACE)
                                .await?;
                        tokio::time::timeout(
                            Duration::from_secs(2),
                            proxy.call::<_, _, ()>("Scroll", &(delta, orientation)),
                        )
                        .await
                        .map_err(|_| zbus::Error::Failure("tray scroll timed out".to_owned()))??;
                        Ok::<_, zbus::Error>(())
                    })
                    .map_err(Into::into),
                Err(error) => Err(error.into()),
            };
            if let Err(error) = result {
                tracing::warn!(%error, "tray scroll failed");
            }
        })
        .context("could not start tray scroll action")?;
    Ok(())
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn parses_registered_item_forms() {
        assert_eq!(
            TrayAddress::parse("org.example.Item").unwrap().path,
            "/StatusNotifierItem"
        );
        let item = TrayAddress::parse(":1.42/Custom/Item").unwrap();
        assert_eq!(item.service, ":1.42");
        assert_eq!(TrayAddress::from_id(&item.id()), Some(item));
    }

    #[test]
    fn rejects_oversized_and_converts_argb() {
        let settings = TrayModule::default();
        let icon = choose_pixmap(vec![(1, 1, vec![128, 255, 0, 0])], &settings)
            .unwrap()
            .unwrap();
        assert_eq!(icon.pixels, vec![0, 0, 128, 128]);
        assert!(
            choose_pixmap(vec![(1024, 1024, vec![0; 4])], &settings)
                .unwrap()
                .is_none()
        );
    }

    #[test]
    fn removes_all_registrations_owned_by_a_vanished_bus_connection() {
        let mut items = BTreeMap::from([
            (
                ":1.10".to_owned(),
                ItemRegistration {
                    owner: ":1.10".to_owned(),
                    service: ":1.10".to_owned(),
                },
            ),
            (
                "org.example.First".to_owned(),
                ItemRegistration {
                    owner: ":1.10".to_owned(),
                    service: "org.example.First".to_owned(),
                },
            ),
            (
                "org.example.Other".to_owned(),
                ItemRegistration {
                    owner: ":1.11".to_owned(),
                    service: "org.example.Other".to_owned(),
                },
            ),
        ]);

        let removed = remove_registrations_for_name(&mut items, ":1.10");
        assert_eq!(removed.len(), 2);
        assert!(removed.contains(&":1.10".to_owned()));
        assert!(removed.contains(&"org.example.First".to_owned()));
        assert_eq!(items.len(), 1);
        assert!(items.contains_key("org.example.Other"));
    }

    #[test]
    fn removes_registration_when_its_well_known_name_is_released() {
        let mut items = BTreeMap::from([(
            "org.example.Tray".to_owned(),
            ItemRegistration {
                owner: ":1.20".to_owned(),
                service: "org.example.Tray".to_owned(),
            },
        )]);

        assert_eq!(
            remove_registrations_for_name(&mut items, "org.example.Tray"),
            vec!["org.example.Tray".to_owned()]
        );
        assert!(items.is_empty());
    }
}