haply 1.2.1

//! Device module provides the core functionality for interacting with Haply haptic devices.
//! It handles device communication, state management, and command processing through both
//! HTTP and WebSocket protocols.

use crate::device_model::{
    AngularPositionInput,
    AngularTorqueInput,
    Command,
    Config,
    ExtensionDataInput,
    ForceInput,
    Inverse3CommandClear,
    Inverse3Configure,
    PositionInput,
    SdfHfxObject,
    ServiceData,
    ServiceMsg,
    SessionConfigure,
    TimestampedServiceData,
    VerseGripCommandClear,
    VerseGripConfigure,
    VerseGripDuplicateMode,
    VersionResponse,
};
use crate::state::{
    clear_inverse3_commands_in_msg,
    clear_oneshot_fields,
    clear_verse_grip_commands_in_msg,
    update_entire_msg,
    update_state_on_message,
    VerseGripLane,
};
use crate::http::InverseHttpClient;
use crate::websocket::InverseWebSocketClient;
use log::{ error, trace };
use std::sync::Arc;
use tokio::sync::Mutex;

#[cfg(test)]
#[path = "unit_tests.rs"]
mod unit_tests;

/// Main struct representing a Haply haptic device.
/// Handles all communication and state management for a connected device.
pub struct HaplyDevice {
    http_client: InverseHttpClient,
    pub ws_client: InverseWebSocketClient,
    /// Cached device state
    pub state: Arc<Mutex<TimestampedServiceData>>,
    /// Cached service cmd msg that will be sent to the device on `send_command` call
    pub device_cmd_msg: Arc<Mutex<ServiceMsg>>,
    verse_grip_mode: VerseGripDuplicateMode,
}

impl HaplyDevice {
    /// Creates a new HaplyDevice with default options (`PreferCustom` dedup mode).
    pub async fn new(
        http_base: &str,
        ws_url: &str
    ) -> Result<Self, Box<dyn std::error::Error + Send + Sync>> {
        Self::with_options(http_base, ws_url, VerseGripDuplicateMode::default()).await
    }

    /// Creates a new HaplyDevice with a specific verse grip duplicate mode.
    pub async fn with_options(
        http_base: &str,
        ws_url: &str,
        verse_grip_mode: VerseGripDuplicateMode
    ) -> Result<Self, Box<dyn std::error::Error + Send + Sync>> {
        let http_client = InverseHttpClient::with_options(http_base, verse_grip_mode);
        let mut ws_client = InverseWebSocketClient::new(ws_url);
        ws_client.connect().await?;

        let initial_state = TimestampedServiceData {
            data: ServiceData::default(),
            timestamp: std::time::Instant::now(),
        };
        let device = Self {
            http_client,
            ws_client: ws_client.clone(),
            state: Arc::new(Mutex::new(initial_state)),
            device_cmd_msg: Arc::new(Mutex::new(ServiceMsg::default())),
            verse_grip_mode,
        };

        device.start_listening();
        Ok(device)
    }

    /// Shutdown the device and its connections.
    pub async fn shutdown(&mut self) -> Result<(), Box<dyn std::error::Error + Send + Sync>> {
        self.ws_client.disconnect().await
    }

    /// Starts the WebSocket listening loop for device updates.
    fn start_listening(&self) {
        #[cfg(feature = "tracy")]
        let _span = tracy_client::span!("device::start_listening");

        let state_clone = Arc::clone(&self.state);
        let mut ws_client = self.ws_client.clone();
        let mode = self.verse_grip_mode;
        tokio::spawn(async move {
            let mut first_msg = true;
            ws_client.listen(move |message| {
                let state_clone = Arc::clone(&state_clone);
                async move {
                    update_state_on_message(state_clone, message, &mut first_msg, mode).await;
                }
            }).await;
        });
    }

    /// Retrieves the current state of the device.
    pub async fn read_state(
        &self
    ) -> Result<ServiceData, Box<dyn std::error::Error + Send + Sync>> {
        #[cfg(feature = "tracy")]
        let _span = tracy_client::span!("device::read_state");

        trace!("Reading device state");
        let data = {
            #[cfg(feature = "tracy")]
            let _lock_span = tracy_client::span!("device::read_state_lock_and_clone");
            self.state.lock().await.data.clone()
        };
        Ok(data)
    }

    /// Forces a full state update from the service then returns the current state.
    pub async fn force_read_full_state(
        &mut self,
        timeout: Option<std::time::Duration>
    ) -> Result<ServiceData, Box<dyn std::error::Error + Send + Sync>> {
        let start = std::time::Instant::now();
        self.send_force_full_render().await?;
        if let Some(dur) = timeout {
            loop {
                let elapsed = start.elapsed();
                if elapsed >= dur {
                    error!("Timeout waiting for full state update");
                    return Err("Timeout waiting for full state update".into());
                }
                let last_update = self.state.lock().await;
                if last_update.timestamp >= start {
                    trace!("State updated after force render");
                    break;
                }
                tokio::time::sleep(std::time::Duration::from_millis(10)).await;
            }
        }
        self.read_state().await
    }

    // -------------------------------------------------------
    // Per-tick command convenience methods
    // -------------------------------------------------------

    /// Updates force values for one or more Inverse3 devices.
    pub async fn update_force(
        &mut self,
        forces_input: Vec<ForceInput>,
        execute_check: Option<bool>,
        send: Option<bool>
    ) -> Result<(), Box<dyn std::error::Error + Send + Sync>> {
        let commands: Vec<Command> = forces_input
            .into_iter()
            .map(|fi| Command::SetCursorForce {
                device_id: fi.device_id,
                vector: fi.forces,
                execute: execute_check.unwrap_or(true),
            })
            .collect();

        self.update_command_msg(commands).await?;
        if send.unwrap_or(false) {
            self.send_command().await?;
        }
        Ok(())
    }

    /// Updates position values for one or more Inverse3 devices.
    pub async fn update_position(
        &mut self,
        positions_input: Vec<PositionInput>,
        send: Option<bool>
    ) -> Result<(), Box<dyn std::error::Error + Send + Sync>> {
        let commands: Vec<Command> = positions_input
            .into_iter()
            .map(|pi| Command::SetCursorPosition {
                device_id: pi.device_id,
                position: pi.positions,
                execute: true,
            })
            .collect();

        self.update_command_msg(commands).await?;
        if send.unwrap_or(false) {
            self.send_command().await?;
        }
        Ok(())
    }

    /// Updates angular torques for one or more Inverse3 devices.
    pub async fn update_angular_torques(
        &mut self,
        inputs: Vec<AngularTorqueInput>,
        send: Option<bool>
    ) -> Result<(), Box<dyn std::error::Error + Send + Sync>> {
        let commands: Vec<Command> = inputs
            .into_iter()
            .map(|i| Command::SetAngularTorques {
                device_id: i.device_id,
                torques: i.torques,
                execute: true,
            })
            .collect();

        self.update_command_msg(commands).await?;
        if send.unwrap_or(false) {
            self.send_command().await?;
        }
        Ok(())
    }

    /// Updates angular position for one or more Inverse3 devices.
    pub async fn update_angular_position(
        &mut self,
        inputs: Vec<AngularPositionInput>,
        send: Option<bool>
    ) -> Result<(), Box<dyn std::error::Error + Send + Sync>> {
        let commands: Vec<Command> = inputs
            .into_iter()
            .map(|i| Command::SetAngularPosition {
                device_id: i.device_id,
                angles: i.angles,
                execute: true,
            })
            .collect();

        self.update_command_msg(commands).await?;
        if send.unwrap_or(false) {
            self.send_command().await?;
        }
        Ok(())
    }

    /// Updates extension data for wireless verse grip devices.
    pub async fn update_extension_data(
        &mut self,
        extension_data_inputs: Vec<ExtensionDataInput>,
        send: Option<bool>
    ) -> Result<(), Box<dyn std::error::Error + Send + Sync>> {
        let commands: Vec<Command> = extension_data_inputs
            .into_iter()
            .map(|edi| Command::SetExtensionData {
                device_id: edi.device_id,
                extension_data: edi.extension_data,
            })
            .collect();

        self.update_command_msg(commands).await?;
        if send.unwrap_or(false) {
            self.send_command().await?;
        }
        Ok(())
    }

    /// Requests position probe from specified Inverse3 devices.
    pub async fn probe_cursor_position(
        &mut self,
        device_ids: Vec<String>,
        send: Option<bool>
    ) -> Result<(), Box<dyn std::error::Error + Send + Sync>> {
        let commands: Vec<Command> = device_ids
            .into_iter()
            .map(|id| Command::ProbePosition { device_id: id })
            .collect();

        self.update_command_msg(commands).await?;
        if send.unwrap_or(false) {
            self.send_command().await?;
        }
        Ok(())
    }

    /// Requests position probe (3.5 unified probe — replaces separate angular probe).
    #[deprecated(
        note = "Use probe_cursor_position instead — 3.5 unified probe_position covers both"
    )]
    pub async fn probe_angular_position(
        &mut self,
        device_ids: Vec<String>,
        send: Option<bool>
    ) -> Result<(), Box<dyn std::error::Error + Send + Sync>> {
        self.probe_cursor_position(device_ids, send).await
    }

    /// Requests orientation probe from specified VerseGrip/WVG devices.
    pub async fn probe_orientation(
        &mut self,
        device_ids: Vec<String>,
        send: Option<bool>
    ) -> Result<(), Box<dyn std::error::Error + Send + Sync>> {
        let commands: Vec<Command> = device_ids
            .into_iter()
            .map(|id| Command::ProbeOrientation { device_id: id })
            .collect();

        self.update_command_msg(commands).await?;
        if send.unwrap_or(false) {
            self.send_command().await?;
        }
        Ok(())
    }

    // -------------------------------------------------------
    // Configure convenience methods (one-shot)
    // -------------------------------------------------------

    /// Configures an Inverse3 device (preset, basis, mount, damping, force_gate, navigation).
    pub async fn configure_inverse3(
        &mut self,
        device_id: &str,
        config: Inverse3Configure,
        send: Option<bool>
    ) -> Result<(), Box<dyn std::error::Error + Send + Sync>> {
        let command = Command::ConfigureInverse3 {
            device_id: device_id.to_string(),
            config,
        };
        self.update_command_msg(vec![command]).await?;
        if send.unwrap_or(false) {
            self.send_command().await?;
        }
        Ok(())
    }

    /// Configures a VerseGrip/WVG device (preset, basis, mount).
    pub async fn configure_versegrip(
        &mut self,
        device_id: &str,
        config: VerseGripConfigure,
        send: Option<bool>
    ) -> Result<(), Box<dyn std::error::Error + Send + Sync>> {
        let command = Command::ConfigureVerseGrip {
            device_id: device_id.to_string(),
            config,
        };
        self.update_command_msg(vec![command]).await?;
        if send.unwrap_or(false) {
            self.send_command().await?;
        }
        Ok(())
    }

    /// Configures the session (profile, basis, serialization, sdf output).
    pub async fn configure_session(
        &mut self,
        config: SessionConfigure,
        send: Option<bool>
    ) -> Result<(), Box<dyn std::error::Error + Send + Sync>> {
        let command = Command::ConfigureSession(config);
        self.update_command_msg(vec![command]).await?;
        if send.unwrap_or(false) {
            self.send_command().await?;
        }
        Ok(())
    }

    // -------------------------------------------------------
    // SDF HFX convenience methods
    // -------------------------------------------------------

    /// Atomic replace — clears all SDF effects for this device, then inserts `objects`.
    pub async fn sdf_set(
        &mut self,
        device_id: &str,
        objects: Vec<SdfHfxObject>,
        from_space: Option<String>,
        send: Option<bool>
    ) -> Result<(), Box<dyn std::error::Error + Send + Sync>> {
        let command = Command::SdfSet {
            device_id: device_id.to_string(),
            objects,
            from_space,
        };
        self.update_command_msg(vec![command]).await?;
        if send.unwrap_or(false) {
            self.send_command().await?;
        }
        Ok(())
    }

    /// Partial update — only provided fields are applied; `id` required on each object.
    pub async fn sdf_update(
        &mut self,
        device_id: &str,
        objects: Vec<SdfHfxObject>,
        send: Option<bool>
    ) -> Result<(), Box<dyn std::error::Error + Send + Sync>> {
        let command = Command::SdfUpdate {
            device_id: device_id.to_string(),
            objects,
        };
        self.update_command_msg(vec![command]).await?;
        if send.unwrap_or(false) {
            self.send_command().await?;
        }
        Ok(())
    }

    /// Remove SDF effects by ID.
    pub async fn sdf_remove(
        &mut self,
        device_id: &str,
        ids: Vec<String>,
        send: Option<bool>
    ) -> Result<(), Box<dyn std::error::Error + Send + Sync>> {
        let command = Command::SdfRemove {
            device_id: device_id.to_string(),
            ids,
        };
        self.update_command_msg(vec![command]).await?;
        if send.unwrap_or(false) {
            self.send_command().await?;
        }
        Ok(())
    }

    /// Enable or disable SDF state output streaming for the session.
    /// The `device_id` argument is retained for backward compatibility and is ignored.
    pub async fn configure_sdf_output(
        &mut self,
        device_id: &str,
        state_output: bool,
        send: Option<bool>
    ) -> Result<(), Box<dyn std::error::Error + Send + Sync>> {
        let command = Command::ConfigureSdfOutput {
            device_id: device_id.to_string(),
            state_output,
        };
        self.update_command_msg(vec![command]).await?;
        if send.unwrap_or(false) {
            self.send_command().await?;
        }
        Ok(())
    }

    // -------------------------------------------------------
    // Session-level commands
    // -------------------------------------------------------

    /// Sends a force full render command to the service.
    pub async fn send_force_full_render(
        &mut self
    ) -> Result<(), Box<dyn std::error::Error + Send + Sync>> {
        self.update_command_msg(vec![Command::ForceRenderFullState]).await?;
        self.send_command().await?;
        Ok(())
    }

    /// Sends a lightweight ping to keep the WebSocket connection alive.
    pub async fn send_ping(&mut self) -> Result<(), Box<dyn std::error::Error + Send + Sync>> {
        self.ws_client.send_message("{}").await?;
        Ok(())
    }

    // -------------------------------------------------------
    // Command accumulation and transmission
    // -------------------------------------------------------

    /// Accumulates commands into the cached service message.
    pub async fn update_command_msg(
        &mut self,
        commands: Vec<Command>
    ) -> Result<(), Box<dyn std::error::Error + Send + Sync>> {
        #[cfg(feature = "tracy")]
        let _span = tracy_client::span!("device::update_command_msg");

        let mut service_msg = {
            #[cfg(feature = "tracy")]
            let _lock_span = tracy_client::span!("device::update_command_msg_lock");
            self.device_cmd_msg.lock().await
        };
        {
            #[cfg(feature = "tracy")]
            let _merge_span = tracy_client::span!("device::update_command_msg_merge");
            update_entire_msg(&commands, &mut service_msg)?;
        }
        Ok(())
    }

    /// Sends the cached command message through the WebSocket connection.
    /// After sending, one-shot fields (configure, force_render_full_state) are cleared.
    /// Per-tick commands are preserved for the next frame.
    pub async fn send_command(&mut self) -> Result<(), Box<dyn std::error::Error + Send + Sync>> {
        #[cfg(feature = "tracy")]
        let _span = tracy_client::span!("device::send_command");

        let mut msg = {
            #[cfg(feature = "tracy")]
            let _lock_span = tracy_client::span!("device::send_command_lock");
            self.device_cmd_msg.lock().await
        };
        let cmd_text = {
            #[cfg(feature = "tracy")]
            let _serialize_span = tracy_client::span!("device::send_command_serialize");
            serde_json::to_string(&*msg)?
        };
        {
            #[cfg(feature = "tracy")]
            let _send_span = tracy_client::span!("device::send_command_ws_send");
            self.ws_client.send_message(&cmd_text).await?;
        }

        // Clear one-shot fields after send
        {
            #[cfg(feature = "tracy")]
            let _clear_span = tracy_client::span!("device::send_command_clear_oneshot");
            clear_oneshot_fields(&mut msg);
        }

        Ok(())
    }

    // -------------------------------------------------------
    // Per-tick command clearing
    // -------------------------------------------------------
    //
    // Per-tick command fields (`set_cursor_force`, `set_cursor_position`,
    // `set_extension_data`, etc.) are deliberately preserved in the cached
    // `ServiceMsg` across `send_command` flushes so callers can rely on
    // zero-order hold behavior in slow control loops. When the controlling
    // layer is intentionally switching modes — for example moving from
    // force-rendered control to absolute cursor position — the previous
    // mode's field would otherwise stay cached and continue to appear in
    // every outgoing WS frame, causing the firmware to honor both modes
    // each frame. The methods below let the caller explicitly remove the
    // stale fields so subsequent frames serialize without them.
    //
    // The clear is synchronous and one-shot: it operates on the cached
    // `ServiceMsg` immediately and persists across frames just like any
    // other state mutation. Calling with a default-constructed selector
    // (no fields selected) is a cheap no-op.

    /// Clears the selected per-tick command fields cached for one Inverse3
    /// device, so subsequent WS frames emit those fields as absent until the
    /// caller sets them again. Fields not selected in `what` are left intact.
    ///
    /// If clearing leaves the device's command bag entirely empty, the bag
    /// itself is dropped so the WS frame omits the `commands` object for
    /// that device altogether.
    pub async fn clear_inverse3_commands(
        &mut self,
        device_id: &str,
        what: Inverse3CommandClear,
    ) -> Result<(), Box<dyn std::error::Error + Send + Sync>> {
        if !what.any() {
            return Ok(());
        }

        #[cfg(feature = "tracy")]
        let _span = tracy_client::span!("device::clear_inverse3_commands");

        let mut msg = {
            #[cfg(feature = "tracy")]
            let _lock_span = tracy_client::span!("device::clear_inverse3_commands_lock");
            self.device_cmd_msg.lock().await
        };
        clear_inverse3_commands_in_msg(&mut msg, device_id, what);
        Ok(())
    }

    /// Clears the selected per-tick command fields cached for one wired
    /// VerseGrip device (the `verse_grip` lane in `ServiceMsg`).
    ///
    /// Most callers writing to VerseGrip should use
    /// [`HaplyDevice::clear_wireless_verse_grip_commands`] instead, since the
    /// existing high-level helpers (`update_extension_data`, `probe_orientation`,
    /// `configure_versegrip`) all route to the wireless lane.
    pub async fn clear_verse_grip_commands(
        &mut self,
        device_id: &str,
        what: VerseGripCommandClear,
    ) -> Result<(), Box<dyn std::error::Error + Send + Sync>> {
        if !what.any() {
            return Ok(());
        }

        #[cfg(feature = "tracy")]
        let _span = tracy_client::span!("device::clear_verse_grip_commands");

        let mut msg = self.device_cmd_msg.lock().await;
        clear_verse_grip_commands_in_msg(&mut msg, VerseGripLane::Wired, device_id, what);
        Ok(())
    }

    /// Clears the selected per-tick command fields cached for one
    /// wireless or custom VerseGrip device (the `wireless_verse_grip` lane
    /// in `ServiceMsg`). This is the lane that the existing
    /// `update_extension_data`, `probe_orientation`, and `configure_versegrip`
    /// helpers populate.
    pub async fn clear_wireless_verse_grip_commands(
        &mut self,
        device_id: &str,
        what: VerseGripCommandClear,
    ) -> Result<(), Box<dyn std::error::Error + Send + Sync>> {
        if !what.any() {
            return Ok(());
        }

        #[cfg(feature = "tracy")]
        let _span = tracy_client::span!("device::clear_wireless_verse_grip_commands");

        let mut msg = self.device_cmd_msg.lock().await;
        clear_verse_grip_commands_in_msg(&mut msg, VerseGripLane::Wireless, device_id, what);
        Ok(())
    }

    // -------------------------------------------------------
    // Device querying
    // -------------------------------------------------------

    /// Lists all available Haply devices via HTTP.
    pub async fn list_devices(
        &self
    ) -> Result<Vec<Config>, Box<dyn std::error::Error + Send + Sync>> {
        let configs = self.http_client.get_devices().await?;
        Ok(configs)
    }

    /// Returns all device configs from the cached service state.
    pub async fn get_devices(
        &self
    ) -> Result<Vec<Config>, Box<dyn std::error::Error + Send + Sync>> {
        let state = self.state.lock().await;
        let mut configs = Vec::new();
        for d in &state.data.inverse3 {
            let config = d.config.clone().ok_or("Missing config for Inverse3 device")?;
            configs.push(Config::DeviceConfig(config));
        }
        for d in &state.data.verse_grip {
            let config = d.config.clone().ok_or("Missing config for Verse Grip device")?;
            configs.push(Config::VGConfig(config));
        }
        for d in &state.data.wireless_verse_grip {
            let config = d.config.clone().ok_or("Missing config for Wireless Verse Grip device")?;
            configs.push(Config::WVGConfig(config));
        }
        for d in &state.data.custom_verse_grip {
            let config = d.config.clone().ok_or("Missing config for Custom Verse Grip device")?;
            configs.push(Config::WVGConfig(config));
        }
        Ok(configs)
    }

    /// Finds a specific device by ID.
    pub async fn get_device_info(
        &self,
        device_id: &str
    ) -> Result<Config, Box<dyn std::error::Error + Send + Sync>> {
        let devices = self.list_devices().await?;
        let device = devices.into_iter().find(|c| {
            match c {
                Config::DeviceConfig(dc) => dc.device_info.id == device_id,
                Config::VGConfig(vgc) => vgc.id == device_id,
                Config::WVGConfig(wvgc) => wvgc.id == device_id,
            }
        });
        match device {
            Some(cfg) => Ok(cfg),
            None => Err(format!("Device {} not found", device_id).into()),
        }
    }

    /// Gets the service version from the HTTP client.
    pub async fn get_service_version(
        &self
    ) -> Result<VersionResponse, Box<dyn std::error::Error + Send + Sync>> {
        let version = self.http_client.get_version().await?;
        Ok(version)
    }
}