cosmic-applets/cosmic-applet-bluetooth/src/bluetooth.rs

use std::{collections::HashMap, fmt::Debug, hash::Hash, sync::Arc, time::Duration};

use bluer::{
    agent::{Agent, AgentHandle},
    Adapter, Address, DeviceProperty, Session, Uuid,
};
use cosmic::iced::{self, subscription};

use futures::StreamExt;
use rand::Rng;
use tokio::{
    spawn,
    sync::{
        mpsc::{channel, Receiver, Sender},
        Mutex,
    },
    task::JoinHandle,
    time::timeout,
};

pub fn bluetooth_subscription<I: 'static + Hash + Copy + Send + Sync + Debug>(
    id: I,
) -> iced::Subscription<(I, BluerEvent)> {
    subscription::unfold(id, State::Ready, move |state| start_listening(id, state))
}

pub enum State {
    Ready,
    Waiting { session_state: BluerSessionState },
    Finished,
}

async fn start_listening<I: Copy + Debug>(id: I, state: State) -> (Option<(I, BluerEvent)>, State) {
    match state {
        State::Ready => {
            let session = match Session::new().await {
                Ok(s) => s,
                Err(_) => return (Some((id, BluerEvent::Finished)), State::Finished),
            };
            let (tx, rx) = channel(100);

            let session_state = match BluerSessionState::new(session, rx).await {
                Ok(s) => s,
                Err(_) => return (Some((id, BluerEvent::Finished)), State::Finished),
            };

            let state = session_state.bluer_state().await;
            return (
                Some((
                    id,
                    BluerEvent::Init {
                        sender: tx,
                        state: state.clone(),
                    },
                )),
                State::Waiting { session_state },
            );
        }
        State::Waiting { mut session_state } => {
            let mut session_rx = match session_state.rx.take() {
                Some(rx) => rx,
                None => {
                    return (Some((id, BluerEvent::Finished)), State::Finished); // fail if we can't get the rx
                }
            };

            let event = if let Some(event) = session_rx.recv().await {
                match event {
                    BluerSessionEvent::ChangesProcessed(state) => {
                        Some((id, BluerEvent::DevicesChanged { state }))
                    }
                    BluerSessionEvent::RequestResponse {
                        req,
                        state,
                        err_msg,
                    } => Some((
                        id,
                        BluerEvent::RequestResponse {
                            req,
                            state,
                            err_msg,
                        },
                    )),
                    BluerSessionEvent::AgentEvent(e) => Some((id, BluerEvent::AgentEvent(e))),
                    _ => None,
                }
            } else {
                return (Some((id, BluerEvent::Finished)), State::Finished);
            };
            session_state.rx = Some(session_rx);
            (event, State::Waiting { session_state })
        }
        State::Finished => iced::futures::future::pending().await,
    }
}

#[derive(Debug, Clone, Hash, Eq, PartialEq)]
pub enum BluerRequest {
    SetBluetoothEnabled(bool),
    SetPairable(bool),
    SetDiscoverable(bool),
    PairDevice(Address),
    ConnectDevice(Address),
    DisconnectDevice(Address),
    CancelConnect(Address),
    StateUpdate,
}

#[derive(Debug, Clone)]
pub enum BluerEvent {
    RequestResponse {
        req: BluerRequest,
        state: BluerState,
        err_msg: Option<String>,
    },
    Init {
        sender: Sender<BluerRequest>,
        state: BluerState,
    },
    DevicesChanged {
        state: BluerState,
    },
    AgentEvent(BluerAgentEvent),
    Finished,
}

#[derive(Debug, Clone, Default)]
pub struct BluerState {
    pub devices: Vec<BluerDevice>,
    pub bluetooth_enabled: bool,
    pub discoverable: bool,
    pub pairable: bool,
}

#[derive(Debug, Clone, PartialEq, Eq, PartialOrd, Ord)]
pub enum BluerDeviceStatus {
    Connected,
    Connecting,
    Paired,
    /// Pairing is in progress, maybe with a passkey or pincode
    /// passkey or pincode will be 000000 - 999999
    Pairing,
    Disconnected,
    Disconnecting,
}

#[derive(Debug, Clone)]
pub struct BluerDevice {
    pub name: String,
    pub address: Address,
    pub status: BluerDeviceStatus,
    pub properties: Vec<DeviceProperty>,
}

impl Eq for BluerDevice {}

impl Ord for BluerDevice {
    fn cmp(&self, other: &Self) -> std::cmp::Ordering {
        match self.status.cmp(&other.status) {
            std::cmp::Ordering::Equal => self.name.to_lowercase().cmp(&other.name.to_lowercase()),
            o => o,
        }
    }
}

impl PartialOrd for BluerDevice {
    fn partial_cmp(&self, other: &Self) -> Option<std::cmp::Ordering> {
        match self.status.cmp(&other.status) {
            std::cmp::Ordering::Equal => {
                Some(self.name.to_lowercase().cmp(&other.name.to_lowercase()))
            }
            o => Some(o),
        }
    }
}

impl PartialEq for BluerDevice {
    fn eq(&self, other: &Self) -> bool {
        self.name == other.name && self.address == other.address
    }
}

impl BluerDevice {
    pub async fn from_device(device: &bluer::Device) -> Self {
        let name = device
            .name()
            .await
            .unwrap_or_default()
            .unwrap_or_else(|| "Unknown".to_string());
        let is_paired = device.is_paired().await.unwrap_or_default();
        let is_connected = device.is_connected().await.unwrap_or_default();
        let properties = device.all_properties().await.unwrap_or_default();
        let status = if is_connected {
            BluerDeviceStatus::Connected
        } else if is_paired {
            BluerDeviceStatus::Paired
        } else {
            BluerDeviceStatus::Disconnected
        };
        Self {
            name,
            address: device.address(),
            status,
            properties,
        }
    }
}

#[derive(Debug, Clone)]
pub enum BluerSessionEvent {
    RequestResponse {
        req: BluerRequest,
        state: BluerState,
        err_msg: Option<String>,
    },
    ChangesProcessed(BluerState),
    ChangeStreamEnded, // TODO can we just restart the stream in a new task?
    AgentEvent(BluerAgentEvent),
}

#[derive(Debug, Clone)]
pub enum BluerAgentEvent {
    DisplayPinCode(BluerDevice, String),
    DisplayPasskey(BluerDevice, String),
    RequestPinCode(BluerDevice),
    RequestPasskey(BluerDevice),
    RequestConfirmation(BluerDevice, String, Sender<bool>), // Note mpsc channel is used bc the sender must be cloned in the iced Message machinery
    RequestDeviceAuthorization(BluerDevice, Sender<bool>),
    RequestServiceAuthorization(BluerDevice, Uuid, Sender<bool>),
}

pub struct BluerSessionState {
    _session: Session,
    _agent_handle: AgentHandle,
    pub adapter: Adapter,
    pub devices: Arc<Mutex<Vec<BluerDevice>>>,
    pub rx: Option<Receiver<BluerSessionEvent>>,
    tx: Sender<BluerSessionEvent>,
    active_requests: Arc<Mutex<HashMap<BluerRequest, JoinHandle<anyhow::Result<()>>>>>,
}

impl BluerSessionState {
    pub(crate) async fn new(
        session: Session,
        request_rx: Receiver<BluerRequest>,
    ) -> anyhow::Result<Self> {
        let adapter = session.default_adapter().await?;
        let devices = build_device_list(&adapter).await;
        let (tx, rx) = tokio::sync::mpsc::channel(100);

        let tx_clone_1 = tx.clone();
        let tx_clone_2 = tx.clone();
        let tx_clone_3 = tx.clone();
        let tx_clone_4 = tx.clone();
        let tx_clone_5 = tx.clone();
        let tx_clone_6 = tx.clone();
        let tx_clone_7 = tx.clone();
        let adapter_clone_1 = adapter.clone();
        let adapter_clone_2 = adapter.clone();
        let adapter_clone_3 = adapter.clone();
        let adapter_clone_4 = adapter.clone();
        let adapter_clone_5 = adapter.clone();
        let adapter_clone_6 = adapter.clone();
        let adapter_clone_7 = adapter.clone();

        let _agent = Agent {
            request_default: false, // TODO which agent should eventually become the default? Maybe the one in the settings app?
            request_pin_code: Some(Box::new(move |req| {
                let agent_clone = adapter_clone_1.clone();
                let tx_clone = tx_clone_1.clone();
                Box::pin(async move {
                    let device = match agent_clone.device(req.device) {
                        Ok(d) => d,
                        Err(_) => return Err(bluer::agent::ReqError::Rejected),
                    };
                    let _ = tx_clone
                        .send(BluerSessionEvent::AgentEvent(
                            BluerAgentEvent::RequestPinCode(
                                BluerDevice::from_device(&device).await,
                            ),
                        ))
                        .await;
                    let mut rng = rand::thread_rng();
                    let pin_code = rng.gen_range(0..999999);
                    Ok(format!("{:06}", pin_code))
                })
            })),
            display_pin_code: Some(Box::new(move |req| {
                let agent_clone = adapter_clone_2.clone();
                let tx_clone = tx_clone_2.clone();
                Box::pin(async move {
                    let device = match agent_clone.device(req.device) {
                        Ok(d) => d,
                        Err(_) => return Err(bluer::agent::ReqError::Rejected),
                    };
                    let _ = tx_clone
                        .send(BluerSessionEvent::AgentEvent(
                            BluerAgentEvent::DisplayPinCode(
                                BluerDevice::from_device(&device).await,
                                req.pincode,
                            ),
                        ))
                        .await;

                    Ok(())
                })
            })),
            request_passkey: Some(Box::new(move |req| {
                let agent_clone = adapter_clone_3.clone();
                let tx_clone = tx_clone_3.clone();
                Box::pin(async move {
                    let device = match agent_clone.device(req.device) {
                        Ok(d) => d,
                        Err(_) => return Err(bluer::agent::ReqError::Rejected),
                    };
                    let _ = tx_clone
                        .send(BluerSessionEvent::AgentEvent(
                            BluerAgentEvent::RequestPasskey(
                                BluerDevice::from_device(&device).await,
                            ),
                        ))
                        .await;
                    let mut rng = rand::thread_rng();
                    let pin_code = rng.gen_range(0..999999);
                    Ok(pin_code)
                })
            })),
            display_passkey: Some(Box::new(move |req| {
                let agent_clone = adapter_clone_4.clone();
                let tx_clone = tx_clone_4.clone();
                Box::pin(async move {
                    let device = match agent_clone.device(req.device) {
                        Ok(d) => d,
                        Err(_) => return Err(bluer::agent::ReqError::Rejected),
                    };
                    let _ = tx_clone
                        .send(BluerSessionEvent::AgentEvent(
                            BluerAgentEvent::DisplayPasskey(
                                BluerDevice::from_device(&device).await,
                                format!("{:06}", req.passkey),
                            ),
                        ))
                        .await;
                    Ok(())
                })
            })),
            request_confirmation: Some(Box::new(move |req| {
                let agent_clone = adapter_clone_5.clone();
                let tx_clone = tx_clone_5.clone();
                Box::pin(async move {
                    let device = match agent_clone.device(req.device) {
                        Ok(d) => d,
                        Err(_) => return Err(bluer::agent::ReqError::Rejected),
                    };
                    let (tx, mut rx) = channel(1);
                    let _ = tx_clone
                        .send(BluerSessionEvent::AgentEvent(
                            BluerAgentEvent::RequestConfirmation(
                                BluerDevice::from_device(&device).await,
                                format!("{:06}", req.passkey),
                                tx,
                            ),
                        ))
                        .await;
                    let res = rx.recv().await;
                    match res {
                        Some(res) if res => Ok(()),
                        _ => Err(bluer::agent::ReqError::Rejected),
                    }
                })
            })),
            request_authorization: Some(Box::new(move |req| {
                let agent_clone = adapter_clone_6.clone();
                let tx_clone = tx_clone_6.clone();
                Box::pin(async move {
                    let device = match agent_clone.device(req.device) {
                        Ok(d) => d,
                        Err(_) => return Err(bluer::agent::ReqError::Rejected),
                    };
                    let (tx, mut rx) = channel(1);
                    let _ = tx_clone
                        .send(BluerSessionEvent::AgentEvent(
                            BluerAgentEvent::RequestDeviceAuthorization(
                                BluerDevice::from_device(&device).await,
                                tx,
                            ),
                        ))
                        .await;
                    let res = rx.recv().await;
                    match res {
                        Some(res) if res => Ok(()),
                        _ => Err(bluer::agent::ReqError::Rejected),
                    }
                })
            })),
            authorize_service: Some(Box::new(move |req| {
                let agent_clone = adapter_clone_7.clone();
                let tx_clone = tx_clone_7.clone();
                Box::pin(async move {
                    let device = match agent_clone.device(req.device) {
                        Ok(d) => d,
                        Err(_) => return Err(bluer::agent::ReqError::Rejected),
                    };
                    let (tx, mut rx) = channel(1);
                    // TODO better describe the service to the user
                    let _ = tx_clone
                        .send(BluerSessionEvent::AgentEvent(
                            BluerAgentEvent::RequestServiceAuthorization(
                                BluerDevice::from_device(&device).await,
                                req.service,
                                tx,
                            ),
                        ))
                        .await;
                    let res = rx.recv().await;
                    match res {
                        Some(res) if res => Ok(()),
                        _ => Err(bluer::agent::ReqError::Rejected),
                    }
                })
            })),
            _non_exhaustive: (),
        };
        let _agent_handle = session.register_agent(_agent).await?;

        let self_ = Self {
            _agent_handle,
            _session: session,
            adapter,
            devices: Arc::new(Mutex::new(devices)),
            rx: Some(rx),
            tx,
            active_requests: Arc::new(Mutex::new(HashMap::new())),
        };
        self_.process_requests(request_rx);
        self_.process_changes();

        Ok(self_)
    }

    // Note: For some reason, this doesn't actually seem to work so well. it seems unreliable...
    pub(crate) fn process_changes(&self) {
        let tx = self.tx.clone();
        let adapter_clone = self.adapter.clone();
        let _monitor_devices: tokio::task::JoinHandle<Result<(), anyhow::Error>> =
            spawn(async move {
                let mut change_stream = adapter_clone.discover_devices_with_changes().await?;

                let mut devices_changed = false;
                let mut milli_timeout = 10;
                'outer: loop {
                    while let Ok(event) =
                        timeout(Duration::from_millis(milli_timeout), change_stream.next()).await
                    {
                        if event.is_none() {
                            break 'outer;
                        }
                        devices_changed = true;
                    }
                    if devices_changed {
                        devices_changed = false;
                        let _ = tx
                            .send(BluerSessionEvent::ChangesProcessed(BluerState {
                                devices: build_device_list(&adapter_clone).await,
                                bluetooth_enabled: adapter_clone
                                    .is_powered()
                                    .await
                                    .unwrap_or_default(),
                                discoverable: adapter_clone
                                    .is_discoverable()
                                    .await
                                    .unwrap_or_default(),
                                pairable: adapter_clone.is_pairable().await.unwrap_or_default(),
                            }))
                            .await;
                        // reset timeout
                        milli_timeout = 10;
                    } else {
                        // slow down if no changes occur
                        milli_timeout = (milli_timeout * 2).max(5120);
                    }
                }
                let _ = tx.send(BluerSessionEvent::ChangeStreamEnded).await;
                Ok(())
            });
    }

    pub(crate) fn process_requests(&self, request_rx: Receiver<BluerRequest>) {
        let active_requests = self.active_requests.clone();
        let adapter = self.adapter.clone();
        let tx = self.tx.clone();

        let _handle: JoinHandle<anyhow::Result<()>> = spawn(async move {
            let mut request_rx = request_rx;

            while let Some(req) = request_rx.recv().await {
                let req_clone = req.clone();
                let req_clone_2 = req.clone();
                let active_requests_clone = active_requests.clone();
                let tx_clone = tx.clone();
                let adapter_clone = adapter.clone();
                let handle = spawn(async move {
                    let mut err_msg = None;
                    match &req_clone {
                        BluerRequest::SetBluetoothEnabled(enabled) => {
                            let res = adapter_clone.set_powered(*enabled).await;
                            if let Err(e) = res {
                                err_msg = Some(e.to_string());
                            }
                            if *enabled {
                                let res = adapter_clone.set_discoverable(*enabled).await;
                                if let Err(e) = res {
                                    err_msg = Some(e.to_string());
                                }
                            }
                        }
                        BluerRequest::PairDevice(address) => {
                            let res = adapter_clone.device(address.clone());
                            if let Err(err) = res {
                                err_msg = Some(err.to_string());
                            } else if let Ok(device) = res {
                                let res = device.pair().await;
                                if let Err(err) = res {
                                    err_msg = Some(err.to_string());
                                }
                            }
                        }
                        BluerRequest::ConnectDevice(address) => {
                            let res = adapter_clone.device(address.clone());
                            if let Err(err) = res {
                                err_msg = Some(err.to_string());
                            } else if let Ok(device) = res {
                                let res = device.connect().await;
                                if let Err(err) = res {
                                    err_msg = Some(err.to_string());
                                }
                            }
                        }
                        BluerRequest::DisconnectDevice(address) => {
                            let res = adapter_clone.device(address.clone());
                            if let Err(err) = res {
                                err_msg = Some(err.to_string());
                            } else if let Ok(device) = res {
                                let res = device.disconnect().await;
                                if let Err(err) = res {
                                    err_msg = Some(err.to_string());
                                }
                            }
                        }
                        BluerRequest::CancelConnect(_) => {
                            if let Some(handle) = active_requests_clone.lock().await.get(&req_clone)
                            {
                                handle.abort();
                            } else {
                                err_msg = Some("No active connection request found".to_string());
                            }
                        }
                        BluerRequest::StateUpdate => {}
                        BluerRequest::SetPairable(enabled) => {
                            let res = adapter_clone.set_pairable(*enabled).await;
                            if let Err(e) = res {
                                err_msg = Some(e.to_string());
                            }
                        }
                        BluerRequest::SetDiscoverable(enabled) => {
                            let res = adapter_clone.set_discoverable(*enabled).await;
                            if let Err(e) = res {
                                err_msg = Some(e.to_string());
                            }
                        }
                    };

                    let state = BluerState {
                        devices: build_device_list(&adapter_clone).await,
                        bluetooth_enabled: adapter_clone.is_powered().await.unwrap_or_default(),
                        discoverable: adapter_clone.is_discoverable().await.unwrap_or_default(),
                        pairable: adapter_clone.is_pairable().await.unwrap_or_default(),
                    };

                    let _ = tx_clone
                        .send(BluerSessionEvent::RequestResponse {
                            req: req_clone,
                            state,
                            err_msg,
                        })
                        .await;

                    let mut active_requests_clone = active_requests_clone.lock().await;
                    let _ = active_requests_clone.remove(&req_clone_2);

                    Ok(())
                });

                active_requests.lock().await.insert(req, handle);
            }
            Ok(())
        });
    }

    pub(crate) async fn bluer_state(&self) -> BluerState {
        BluerState {
            devices: build_device_list(&self.adapter).await,
            // TODO is this a proper way of checking if bluetooth is enabled?
            bluetooth_enabled: self.adapter.is_powered().await.unwrap_or_default(),
            discoverable: self.adapter.is_discoverable().await.unwrap_or_default(),
            pairable: self.adapter.is_pairable().await.unwrap_or_default(),
        }
    }
}

async fn build_device_list(adapter: &Adapter) -> Vec<BluerDevice> {
    let addrs = adapter.device_addresses().await.unwrap_or_default();
    let mut devices = Vec::with_capacity(addrs.len());

    for address in addrs {
        let device = match adapter.device(address) {
            Ok(device) => device,
            Err(_) => continue,
        };
        let name = device.name().await.unwrap_or_default().unwrap_or_default();
        let is_paired = device.is_paired().await.unwrap_or_default();
        let is_connected = device.is_connected().await.unwrap_or_default();
        let properties = device.all_properties().await.unwrap_or_default();
        let status = if is_connected {
            BluerDeviceStatus::Connected
        } else if is_paired {
            BluerDeviceStatus::Paired
        } else {
            BluerDeviceStatus::Disconnected
        };
        devices.push(BluerDevice {
            name,
            address,
            status,
            properties,
        });
    }
    devices.sort();
    devices
}