mqttmarionette/src/session.rs

use std::collections::HashMap;
use std::sync::Arc;
use std::time::Duration;

use tokio::sync::Notify;
use tracing::{debug, error, info, trace, warn};

use crate::browser::{Browser, BrowserError};
use crate::config::Configuration;
use crate::marionette::error::{CommandError, ConnectionError};
use crate::marionette::message::WindowType;
use crate::marionette::{Marionette, MarionetteConnection};
use crate::monitor::Monitor;
use crate::mqtt::{Message, MqttClient, MqttPublisher};
#[cfg(unix)]
use crate::x11::{xrandr, Display};

#[derive(Debug)]
pub enum SessionError {
    Browser(BrowserError),
    Io(std::io::Error),
    Connection(ConnectionError),
    InvalidState(String),
    Command(CommandError),
}

impl From<BrowserError> for SessionError {
    fn from(e: BrowserError) -> Self {
        Self::Browser(e)
    }
}

impl From<std::io::Error> for SessionError {
    fn from(e: std::io::Error) -> Self {
        Self::Io(e)
    }
}

impl From<ConnectionError> for SessionError {
    fn from(e: ConnectionError) -> Self {
        Self::Connection(e)
    }
}

impl From<CommandError> for SessionError {
    fn from(e: CommandError) -> Self {
        Self::Command(e)
    }
}

impl std::fmt::Display for SessionError {
    fn fmt(&self, f: &mut std::fmt::Formatter) -> std::fmt::Result {
        match self {
            Self::Browser(e) => write!(f, "Error launching browser: {}", e),
            Self::Io(e) => write!(f, "I/O error: {}", e),
            Self::Connection(e) => write!(f, "Connection error: {}", e),
            Self::InvalidState(e) => write!(f, "Invalid state: {}", e),
            Self::Command(e) => write!(f, "Marionette command failed: {}", e),
        }
    }
}

impl std::error::Error for SessionError {
    fn source(&self) -> Option<&(dyn std::error::Error + 'static)> {
        match self {
            Self::Browser(e) => Some(e),
            Self::Io(e) => Some(e),
            Self::Connection(e) => Some(e),
            Self::InvalidState(_) => None,
            Self::Command(e) => Some(e),
        }
    }
}

pub struct Session {
    config: Configuration,
    browser: Browser,
    marionette: Marionette,
}

impl Session {
    pub async fn begin(config: Configuration) -> Result<Self, SessionError> {
        debug!("Launching Firefox");
        let browser = Browser::launch()?;
        browser.wait_ready().await?;
        debug!("Firefox Marionette is now ready");
        let Some(port) = browser.marionette_port() else {
            return Err(SessionError::InvalidState("No active Marionette port".into()));
        };
        debug!("Connecting to Firefox Marionette on port {}", port);
        let conn = MarionetteConnection::connect(("127.0.0.1", port)).await?;
        let mut marionette = Marionette::new(conn);
        info!("Successfully connected to Firefox Marionette");
        let ses = marionette.new_session().await?;
        debug!("Started Marionette session {}", ses.session_id);
        Ok(Self {
            config,
            browser,
            marionette,
        })
    }

    pub async fn run(mut self, stop: Arc<Notify>) {
        let windows = match self.init_windows().await {
            Ok(w) => w,
            Err(e) => {
                error!("Failed to initialize browser windows: {}", e);
                return;
            }
        };

        let mut client = MqttClient::new(&self.config).unwrap();
        loop {
            if let Err(e) = client.connect().await {
                warn!("Failed to connect to MQTT server: {}", e);
                tokio::time::sleep(Duration::from_secs(1)).await;
                continue;
            }
            break;
        }

        let publisher = client.publisher();
        for w in windows.keys() {
            if let Err(e) = publisher.publish_config(w).await {
                error!(
                    "Failed to publish Home Assistant config for screen {}: {}",
                    w, e
                );
            }
        }

        let handler = MessageHandler {
            marionette: &mut self.marionette,
            windows,
        };
        client.run(handler, stop).await;
    }

    async fn init_windows(
        &mut self,
    ) -> Result<HashMap<String, String>, SessionError> {
        debug!("Getting monitor configuration");
        let monitors = get_monitors();
        let mut handles = self.marionette.get_window_handles().await?;
        for handle in handles.iter().take(handles.len() - 1) {
            self.marionette.close_window(handle).await?;
        }
        let mut windowmap = HashMap::new();
        let mut window = Some(handles.remove(handles.len() - 1));
        for monitor in monitors {
            debug!(
                "Creating window for monitor {}: {}x{} ({}, {})",
                monitor.name,
                monitor.width,
                monitor.height,
                monitor.x,
                monitor.y
            );
            if window.is_none() {
                window = Some(
                    self.marionette
                        .new_window(WindowType::Window, false)
                        .await?,
                );
            }
            let w = window.take().unwrap();
            self.marionette.switch_to_window(w.clone(), false).await?;
            self.marionette
                .set_window_rect(Some(monitor.x), Some(monitor.y), None, None)
                .await?;
            self.marionette.fullscreen().await?;
            windowmap.insert(monitor.name, w);
        }
        trace!("Built window map: {:?}", windowmap);
        Ok(windowmap)
    }
}

pub struct MessageHandler<'a> {
    marionette: &'a mut Marionette,
    windows: HashMap<String, String>,
}

#[async_trait::async_trait]
impl<'a> crate::mqtt::MessageHandler for MessageHandler<'a> {
    async fn navigate(&mut self, publisher: &MqttPublisher, msg: &Message) {
        let url = msg.payload_str();
        let parts: Vec<&str> = msg.topic().split('/').rev().collect();
        let screen = match parts.get(1) {
            Some(&"") | None => {
                warn!("Invalid navigate request: no screen");
                return;
            }
            Some(s) => s,
        };
        if let Some(window) = self.windows.get(*screen) {
            debug!("Switching to window {}", window);
            if let Err(e) =
                self.marionette.switch_to_window(window.into(), false).await
            {
                error!(
                    "Failed to switch to window on screen {}: {}",
                    screen, e
                );
                return;
            }
        } else {
            error!("Invalid navigate request: unknown screen {}", screen);
            return;
        }
        debug!("Handling navigate request: {}", url);
        info!("Navigate screen {} to {}", screen, url);
        if let Err(e) = self.marionette.navigate(url.to_string()).await {
            error!("Failed to navigate: {}", e);
        }
        match self.marionette.get_current_url().await {
            Ok(u) => {
                if let Err(e) = publisher.publish_url(screen, &u).await {
                    error!("Failed to publish URL: {}", e);
                }
            }
            Err(e) => error!("Failed to get current browser URL: {}", e),
        }
        match self.marionette.get_title().await {
            Ok(t) => {
                if let Err(e) = publisher.publish_title(screen, &t).await {
                    error!("Failed to publish title: {}", e);
                }
            }
            Err(e) => error!("Error getting title: {}", e),
        }
    }
}

#[cfg(unix)]
fn get_monitors() -> Vec<Monitor> {
    let display = Display::open().unwrap();
    xrandr::get_monitors(&display)
        .iter()
        .map(|m| Monitor {
            name: m.name().unwrap().to_string(),
            width: m.width(),
            height: m.height(),
            x: m.x(),
            y: m.y(),
        })
        .collect()
}