mendi 0.0.2

//! BLE client for the Mendi neurofeedback headband.
//!
//! Handles scanning, connecting, GATT subscription, and notification dispatch.

use std::collections::BTreeSet;
use std::sync::Arc;
use std::sync::atomic::{AtomicBool, Ordering};
use std::time::Duration;

use anyhow::{anyhow, Result};
use btleplug::api::{
    Central, CentralEvent, Characteristic, Manager as _, Peripheral as _, ScanFilter, WriteType,
};
use btleplug::platform::{Adapter, Manager, Peripheral};
use futures::StreamExt;
use log::{debug, info, warn};
use prost::Message;
use tokio::sync::mpsc;
use uuid::Uuid;

use crate::parse::{parse_adc, parse_calibration, parse_diagnostics, parse_frame, parse_sensor};
use crate::protocol::*;
use crate::types::{DeviceInfo, MendiEvent, MENDI_FCC_ID};
use crate::wire;

// ── MendiDevice ──────────────────────────────────────────────────────────────

/// A Mendi headband discovered during a BLE scan.
///
/// Returned by [`MendiClient::scan`]; pass to [`MendiClient::connect_to`]
/// to establish a connection.
#[derive(Clone, Debug)]
pub struct MendiDevice {
    /// Advertised device name (e.g. "Mendi").
    pub name: String,
    /// Platform BLE identifier (MAC address on Linux, UUID on macOS).
    pub id: String,
    pub(crate) peripheral: Peripheral,
    pub(crate) adapter: Adapter,
}

// ── MendiClientConfig ────────────────────────────────────────────────────────

/// Configuration for [`MendiClient`].
#[derive(Debug, Clone)]
pub struct MendiClientConfig {
    /// BLE scan duration in seconds. Default: `10`.
    pub scan_timeout_secs: u64,
    /// Device name prefix to match during scanning. Default: `"Mendi"`.
    pub name_prefix: String,
    /// Use the Mendi service UUID (`fc3eabb0-…`) as a BLE scan filter.
    ///
    /// When `true`, only devices advertising the Mendi service are returned.
    /// When `false`, scanning relies solely on the name prefix (useful if
    /// the headband does not include the service UUID in its advertisements).
    /// Default: `true`.
    pub filter_by_service_uuid: bool,
}

impl Default for MendiClientConfig {
    fn default() -> Self {
        Self {
            scan_timeout_secs: 10,
            name_prefix: MENDI_NAME_PREFIX.into(),
            filter_by_service_uuid: true,
        }
    }
}

// ── MendiClient ──────────────────────────────────────────────────────────────

/// BLE client for the Mendi neurofeedback headband.
///
/// # Quick start
///
/// ```no_run
/// use mendi::prelude::*;
///
/// #[tokio::main]
/// async fn main() -> anyhow::Result<()> {
///     let client = MendiClient::new(MendiClientConfig::default());
///     let (mut rx, handle) = client.connect().await?;
///
///     while let Some(event) = rx.recv().await {
///         match event {
///             MendiEvent::Frame(f) => {
///                 println!("IR left={} right={} pulse={}",
///                     f.ir_left, f.ir_right, f.ir_pulse);
///             }
///             MendiEvent::Disconnected => break,
///             _ => {}
///         }
///     }
///     Ok(())
/// }
/// ```
pub struct MendiClient {
    config: MendiClientConfig,
}

impl MendiClient {
    /// Create a new client with the given configuration.
    pub fn new(config: MendiClientConfig) -> Self {
        Self { config }
    }

    // ── Adapter helper ───────────────────────────────────────────────────────

    async fn get_adapter() -> Result<Adapter> {
        let manager = Manager::new().await?;
        let adapters = manager.adapters().await?;
        adapters
            .into_iter()
            .next()
            .ok_or_else(|| anyhow!("No Bluetooth adapter found"))
    }

    #[cfg(target_os = "macos")]
    async fn wait_for_adapter(adapter: &Adapter) {
        use btleplug::api::CentralState;

        let deadline = tokio::time::Instant::now() + Duration::from_secs(3);
        loop {
            match adapter.adapter_state().await {
                Ok(CentralState::PoweredOn) => {
                    info!("macOS: adapter is PoweredOn");
                    break;
                }
                Ok(state) => {
                    if tokio::time::Instant::now() >= deadline {
                        warn!("macOS: adapter still {state:?} after 3s — proceeding");
                        break;
                    }
                    debug!("macOS: adapter state = {state:?}, waiting…");
                }
                Err(e) => {
                    warn!("macOS: adapter_state() error: {e}");
                    break;
                }
            }
            tokio::time::sleep(Duration::from_millis(200)).await;
        }
        tokio::time::sleep(Duration::from_millis(300)).await;
    }

    // ── Scan ─────────────────────────────────────────────────────────────────

    /// Scan for all nearby Mendi headbands.
    ///
    /// Runs for `config.scan_timeout_secs` seconds and returns all matching devices.
    pub async fn scan(&self) -> Result<Vec<MendiDevice>> {
        let adapter = Self::get_adapter().await?;

        #[cfg(target_os = "macos")]
        Self::wait_for_adapter(&adapter).await;

        info!("Scanning for Mendi devices ({} s) …", self.config.scan_timeout_secs);
        let scan_filter = if self.config.filter_by_service_uuid {
            ScanFilter {
                services: vec![MENDI_SERVICE_UUID],
            }
        } else {
            ScanFilter::default()
        };
        adapter.start_scan(scan_filter).await?;
        tokio::time::sleep(Duration::from_secs(self.config.scan_timeout_secs)).await;
        adapter.stop_scan().await.ok();

        let mut found = vec![];
        for p in adapter.peripherals().await? {
            if let Ok(Some(props)) = p.properties().await {
                if let Some(name) = props.local_name {
                    if name.starts_with(&self.config.name_prefix) {
                        let id = p.id().to_string();
                        info!("Found: {name} (id={id})");
                        found.push(MendiDevice {
                            name,
                            id,
                            peripheral: p,
                            adapter: adapter.clone(),
                        });
                    }
                }
            }
        }
        info!("{} device(s) found", found.len());
        Ok(found)
    }

    // ── Connect ──────────────────────────────────────────────────────────────

    /// Connect to a specific device from [`MendiClient::scan`].
    pub async fn connect_to(
        &self,
        device: MendiDevice,
    ) -> Result<(mpsc::Receiver<MendiEvent>, MendiHandle)> {
        self.setup(device.peripheral, device.name, device.adapter)
            .await
    }

    /// Scan for the first Mendi device, connect, and start streaming.
    ///
    /// Convenience method when only one headband is expected.
    pub async fn connect(&self) -> Result<(mpsc::Receiver<MendiEvent>, MendiHandle)> {
        let adapter = Self::get_adapter().await?;

        #[cfg(target_os = "macos")]
        Self::wait_for_adapter(&adapter).await;

        info!("Scanning for Mendi device ({} s) …", self.config.scan_timeout_secs);
        let scan_filter = if self.config.filter_by_service_uuid {
            ScanFilter {
                services: vec![MENDI_SERVICE_UUID],
            }
        } else {
            ScanFilter::default()
        };
        adapter.start_scan(scan_filter).await?;

        let peripheral = self
            .find_first(&adapter, &self.config.name_prefix, self.config.scan_timeout_secs)
            .await?;
        adapter.stop_scan().await.ok();

        let props = peripheral.properties().await?.unwrap_or_default();
        let name = props.local_name.unwrap_or_else(|| "Mendi".into());
        info!("Found: {name}");

        self.setup(peripheral, name, adapter).await
    }

    // ── Setup ────────────────────────────────────────────────────────────────

    async fn setup(
        &self,
        peripheral: Peripheral,
        device_name: String,
        adapter: Adapter,
    ) -> Result<(mpsc::Receiver<MendiEvent>, MendiHandle)> {
        // Connect with timeout
        tokio::time::timeout(Duration::from_secs(10), peripheral.connect())
            .await
            .map_err(|_| anyhow!("BLE connect() timed out after 10 s"))??;

        // Linux: wait for GATT cache to populate
        #[cfg(target_os = "linux")]
        tokio::time::sleep(Duration::from_millis(600)).await;

        tokio::time::timeout(Duration::from_secs(15), peripheral.discover_services())
            .await
            .map_err(|_| anyhow!("discover_services() timed out after 15 s"))??;
        info!("Connected, services discovered: {device_name}");

        let chars: BTreeSet<Characteristic> = peripheral.characteristics();

        let find_char = |uuid: Uuid| -> Option<Characteristic> {
            chars.iter().find(|c| c.uuid == uuid).cloned()
        };

        // Read device info from standard BLE characteristics
        let firmware_version = if let Some(c) = find_char(FIRMWARE_REVISION) {
            peripheral
                .read(&c)
                .await
                .ok()
                .and_then(|v| String::from_utf8(v).ok())
        } else {
            None
        };

        let hardware_version = if let Some(c) = find_char(HARDWARE_REVISION) {
            peripheral
                .read(&c)
                .await
                .ok()
                .and_then(|v| String::from_utf8(v).ok())
        } else {
            None
        };

        let device_info = DeviceInfo {
            firmware_version,
            hardware_version,
            id: peripheral.id().to_string(),
            name: device_name.clone(),
            fcc_id: MENDI_FCC_ID.into(),
        };
        info!("Device: {device_info}");

        // ── Subscribe to characteristics ─────────────────────────────────────

        // Frame (0xABB1) — main sensor data stream
        if let Some(c) = find_char(FRAME_CHARACTERISTIC) {
            peripheral.subscribe(&c).await?;
            info!("Subscribed to Frame (0xABB1)");
        } else {
            warn!("Frame characteristic (0xABB1) not found");
        }

        // ADC / Battery (0xABB4)
        if let Some(c) = find_char(ADC_CHARACTERISTIC) {
            peripheral.subscribe(&c).await?;
            info!("Subscribed to ADC (0xABB4)");
        }

        // Calibration (0xABB6)
        if let Some(c) = find_char(CALIBRATION_CHARACTERISTIC) {
            peripheral.subscribe(&c).await?;
            info!("Subscribed to Calibration (0xABB6)");
        }

        // Sensor (0xABB2) — subscribe for register read responses
        if let Some(c) = find_char(SENSOR_CHARACTERISTIC) {
            peripheral.subscribe(&c).await?;
            info!("Subscribed to Sensor (0xABB2)");
        }

        // ── Event channel ────────────────────────────────────────────────────

        let (tx, rx) = mpsc::channel::<MendiEvent>(256);
        let _ = tx.send(MendiEvent::Connected(device_info)).await;

        // Diagnostics (0xABB5) — read once and send as event
        if let Some(c) = find_char(DIAGNOSTICS_CHARACTERISTIC) {
            match peripheral.read(&c).await {
                Ok(data) => {
                    if let Some(diag) = parse_diagnostics(&data) {
                        info!(
                            "Diagnostics: imu_ok={}, sensor_ok={}, adc={:?}",
                            diag.imu_ok,
                            diag.sensor_ok,
                            diag.adc.as_ref().map(|a| format!("{}mV", a.voltage_mv))
                        );
                        let _ = tx.send(MendiEvent::Diagnostics(diag)).await;
                    }
                }
                Err(e) => debug!("Could not read Diagnostics: {e}"),
            }
        }

        // ── Disconnect guard ──────────────────────────────────────────────
        // Both the adapter event watcher and the notification stream can
        // detect disconnection. This flag ensures only one Disconnected
        // event reaches the consumer.
        let disconnected = Arc::new(AtomicBool::new(false));

        // ── Disconnect watcher ───────────────────────────────────────────────
        let disconnect_tx = tx.clone();
        let peripheral_id = peripheral.id();
        let disconnected_w = Arc::clone(&disconnected);
        tokio::spawn(async move {
            match adapter.events().await {
                Ok(mut events) => {
                    while let Some(event) = events.next().await {
                        if let CentralEvent::DeviceDisconnected(id) = event {
                            if id == peripheral_id {
                                info!("Disconnect watcher: device disconnected");
                                if !disconnected_w.swap(true, Ordering::SeqCst) {
                                    let _ = disconnect_tx.send(MendiEvent::Disconnected).await;
                                }
                                break;
                            }
                        }
                    }
                }
                Err(e) => warn!("Could not subscribe to adapter events: {e}"),
            }
        });

        // ── Notification dispatch ────────────────────────────────────────────

        let peripheral_clone = peripheral.clone();
        let disconnected_n = Arc::clone(&disconnected);
        tokio::spawn(async move {
            let mut notifications = match peripheral_clone.notifications().await {
                Ok(n) => n,
                Err(e) => {
                    warn!("Could not get notifications stream: {e}");
                    return;
                }
            };

            let mut count: u64 = 0;
            while let Some(notif) = notifications.next().await {
                let data = &notif.value;
                let uuid = notif.uuid;
                count += 1;

                if count <= 5 || count.is_multiple_of(1000) {
                    debug!("Notification #{count}: uuid={uuid} len={}", data.len());
                }

                if uuid == FRAME_CHARACTERISTIC {
                    match parse_frame(data) {
                        Some(frame) => {
                            if tx.try_send(MendiEvent::Frame(frame)).is_err() {
                                debug!("Frame channel full, dropping frame");
                            }
                        }
                        None => {
                            debug!("Invalid frame (len={}), skipping", data.len());
                        }
                    }
                } else if uuid == ADC_CHARACTERISTIC {
                    if let Some(battery) = parse_adc(data) {
                        let _ = tx.send(MendiEvent::Battery(battery)).await;
                    }
                } else if uuid == CALIBRATION_CHARACTERISTIC {
                    if let Some(cal) = parse_calibration(data) {
                        let _ = tx.send(MendiEvent::Calibration(cal)).await;
                    }
                } else if uuid == SENSOR_CHARACTERISTIC {
                    if let Some(sensor) = parse_sensor(data) {
                        let _ = tx.send(MendiEvent::SensorRead(sensor)).await;
                    }
                } else if uuid == IMU_CHARACTERISTIC {
                    // IMU register read responses (if subscribed)
                    debug!("IMU notification (len={})", data.len());
                    // IMU read responses are raw protobuf Imu messages;
                    // no dedicated event type yet — log for now.
                } else {
                    debug!("Unknown notification from {uuid}");
                }
            }

            info!("Notification stream ended — device disconnected.");
            if !disconnected_n.swap(true, Ordering::SeqCst) {
                let _ = tx.send(MendiEvent::Disconnected).await;
            }
        });

        let handle = MendiHandle {
            peripheral,
            chars: chars.into_iter().collect(),
        };

        Ok((rx, handle))
    }

    // ── Find first ───────────────────────────────────────────────────────────

    async fn find_first(
        &self,
        adapter: &Adapter,
        prefix: &str,
        timeout_secs: u64,
    ) -> Result<Peripheral> {
        tokio::time::timeout(Duration::from_secs(timeout_secs), async {
            loop {
                for p in adapter.peripherals().await.unwrap_or_default() {
                    if let Ok(Some(props)) = p.properties().await {
                        if let Some(name) = &props.local_name {
                            if name.starts_with(prefix) {
                                return p;
                            }
                        }
                    }
                }
                tokio::time::sleep(Duration::from_millis(250)).await;
            }
        })
        .await
        .map_err(|_| anyhow!("Timed out scanning for Mendi after {timeout_secs} s"))
    }
}

// ── MendiHandle ──────────────────────────────────────────────────────────────

/// Handle to an active Mendi connection.
///
/// Provides methods to write calibration settings, read sensor registers, and
/// disconnect.
pub struct MendiHandle {
    peripheral: Peripheral,
    chars: Vec<Characteristic>,
}

impl MendiHandle {
    fn find_char(&self, uuid: Uuid) -> Option<&Characteristic> {
        self.chars.iter().find(|c| c.uuid == uuid)
    }

    /// Write a calibration message to the headband.
    ///
    /// Sets LED current offsets and auto-calibration mode.
    pub async fn write_calibration(
        &self,
        offset_left: f32,
        offset_right: f32,
        offset_pulse: f32,
        enable_auto_cal: bool,
        low_power: bool,
    ) -> Result<()> {
        let char = self
            .find_char(CALIBRATION_CHARACTERISTIC)
            .ok_or_else(|| anyhow!("Calibration characteristic not found"))?;
        let msg = wire::Calibration {
            offset_l: offset_left,
            offset_r: offset_right,
            offset_p: offset_pulse,
            enable: enable_auto_cal,
            low_power_mode: low_power,
        };
        let payload = msg.encode_to_vec();
        self.peripheral
            .write(char, &payload, WriteType::WithResponse)
            .await?;
        Ok(())
    }

    /// Write a sensor register.
    ///
    /// Low-level access to optical sensor registers via the Sensor characteristic (0xABB2).
    pub async fn write_sensor_register(&self, address: u32, data: u32) -> Result<()> {
        let char = self
            .find_char(SENSOR_CHARACTERISTIC)
            .ok_or_else(|| anyhow!("Sensor characteristic not found"))?;
        let msg = wire::Sensor {
            read: false,
            address,
            data,
        };
        let payload = msg.encode_to_vec();
        self.peripheral
            .write(char, &payload, WriteType::WithResponse)
            .await?;
        Ok(())
    }

    /// Read a sensor register.
    ///
    /// Sends a read request; the response arrives as a notification on the
    /// Sensor characteristic (0xABB2).
    pub async fn read_sensor_register(&self, address: u32) -> Result<()> {
        let char = self
            .find_char(SENSOR_CHARACTERISTIC)
            .ok_or_else(|| anyhow!("Sensor characteristic not found"))?;
        let msg = wire::Sensor {
            read: true,
            address,
            data: 0,
        };
        let payload = msg.encode_to_vec();
        self.peripheral
            .write(char, &payload, WriteType::WithResponse)
            .await?;
        Ok(())
    }

    /// Write to an IMU register.
    ///
    /// Low-level access to IMU registers via the IMU characteristic (0xABB3).
    pub async fn write_imu_register(&self, address: u32, data: &[u8]) -> Result<()> {
        let char = self
            .find_char(IMU_CHARACTERISTIC)
            .ok_or_else(|| anyhow!("IMU characteristic not found"))?;
        let msg = wire::Imu {
            read: 0,
            address,
            data: data.to_vec(),
        };
        let payload = msg.encode_to_vec();
        self.peripheral
            .write(char, &payload, WriteType::WithResponse)
            .await?;
        Ok(())
    }

    /// Read from an IMU register.
    ///
    /// Sends a read request; the response arrives as a notification on the
    /// IMU characteristic (0xABB3).
    pub async fn read_imu_register(&self, address: u32, num_bytes: u32) -> Result<()> {
        let char = self
            .find_char(IMU_CHARACTERISTIC)
            .ok_or_else(|| anyhow!("IMU characteristic not found"))?;
        let msg = wire::Imu {
            read: num_bytes,
            address,
            data: vec![],
        };
        let payload = msg.encode_to_vec();
        self.peripheral
            .write(char, &payload, WriteType::WithResponse)
            .await?;
        Ok(())
    }

    /// Enable the optical sensor by writing a Sensor message with `read=true`.
    ///
    /// This mirrors the Mendi app's `makeSensorPayload(true)` /
    /// `listenForSensorUpdates` flow. Must be called to start receiving
    /// Frame notifications on some firmware versions.
    pub async fn enable_sensor(&self) -> Result<()> {
        let char = self
            .find_char(SENSOR_CHARACTERISTIC)
            .ok_or_else(|| anyhow!("Sensor characteristic not found"))?;
        let msg = wire::Sensor {
            read: true,
            address: 0,
            data: 0,
        };
        let payload = msg.encode_to_vec();
        self.peripheral
            .write(char, &payload, WriteType::WithResponse)
            .await?;
        info!("Sensor enabled");
        Ok(())
    }

    /// Disable the optical sensor by writing a Sensor message with `read=false`.
    ///
    /// This mirrors the Mendi app's `makeSensorPayload(false)` /
    /// `cancelSensorUpdates` flow.
    pub async fn disable_sensor(&self) -> Result<()> {
        let char = self
            .find_char(SENSOR_CHARACTERISTIC)
            .ok_or_else(|| anyhow!("Sensor characteristic not found"))?;
        let msg = wire::Sensor {
            read: false,
            address: 0,
            data: 0,
        };
        let payload = msg.encode_to_vec();
        self.peripheral
            .write(char, &payload, WriteType::WithResponse)
            .await?;
        info!("Sensor disabled");
        Ok(())
    }

    /// Check if the headband is still connected.
    pub async fn is_connected(&self) -> bool {
        self.peripheral.is_connected().await.unwrap_or(false)
    }

    /// Disconnect from the headband.
    pub async fn disconnect(&self) -> Result<()> {
        self.peripheral.disconnect().await?;
        Ok(())
    }
}