haply 1.0.0

//! State management and command building — pure functions used by `HaplyDevice`
//! and the WebSocket listener for processing incoming/outgoing messages.

use crate::device_model::{
    Command,
    CustomVerseGripDevice,
    ForceRenderFullStateMsg,
    Inverse3Device,
    Inverse3OutMsg,
    ProbeCmd,
    SdfCommand,
    SdfCommandMode,
    SdfHfxObject,
    SdfOutputConfigure,
    ServiceData,
    ServiceMsg,
    ServiceState,
    SetAngularPositionCmd,
    SetAngularTorquesCmd,
    SetCursorForceCmd,
    SetCursorPositionCmd,
    SetExtensionDataCmd,
    SetTransformCmd,
    TimestampedServiceData,
    VerseGripConfigure,
    VerseGripDevice,
    VerseGripDuplicateMode,
    VerseGripOutMsg,
    WirelessVerseGripDevice,
};
use log::{error, trace};
use serde_json::Value;
use std::sync::Arc;
use tokio::sync::Mutex;

// ============================================================
// Command building helpers
// ============================================================

/// Find-or-create an Inverse3 entry in the service message.
fn ensure_inverse3<'a>(msg: &'a mut ServiceMsg, device_id: &str) -> &'a mut Inverse3OutMsg {
    if let Some(pos) = msg.inverse3.iter().position(|d| d.device_id == device_id) {
        &mut msg.inverse3[pos]
    } else {
        msg.inverse3.push(Inverse3OutMsg {
            device_id: device_id.to_string(),
            commands: None,
            configure: None,
        });
        msg.inverse3.last_mut().unwrap()
    }
}

/// Find-or-create a VerseGrip entry in a given vec.
fn ensure_versegrip<'a>(
    vec: &'a mut Vec<VerseGripOutMsg>,
    device_id: &str,
) -> &'a mut VerseGripOutMsg {
    if let Some(pos) = vec.iter().position(|d| d.device_id == device_id) {
        &mut vec[pos]
    } else {
        vec.push(VerseGripOutMsg {
            device_id: device_id.to_string(),
            commands: None,
            configure: None,
        });
        vec.last_mut().unwrap()
    }
}

pub fn update_entire_msg(
    commands: &[Command],
    msg: &mut ServiceMsg,
) -> Result<(), Box<dyn std::error::Error + Send + Sync>> {
    for command in commands {
        if let Err(e) = update_msg_from_command(command.clone(), msg) {
            error!("Error while creating command message: {}", e);
            return Err(e);
        }
    }
    Ok(())
}

/// Processes a single command and merges it into the service message.
pub fn update_msg_from_command(
    command: Command,
    msg: &mut ServiceMsg,
) -> Result<(), Box<dyn std::error::Error + Send + Sync>> {
    match command {
        // --- Inverse3 per-tick commands ---
        Command::SetCursorForce {
            device_id,
            vector,
            execute,
        } => {
            let entry = ensure_inverse3(msg, &device_id);
            let cmds = entry.commands.get_or_insert_with(Default::default);
            cmds.set_cursor_force = Some(SetCursorForceCmd { vector, execute });
        }
        Command::SetCursorPosition {
            device_id,
            position,
            execute,
        } => {
            let entry = ensure_inverse3(msg, &device_id);
            let cmds = entry.commands.get_or_insert_with(Default::default);
            cmds.set_cursor_position = Some(SetCursorPositionCmd { position, execute });
        }
        Command::SetAngularTorques {
            device_id,
            torques,
            execute,
        } => {
            let entry = ensure_inverse3(msg, &device_id);
            let cmds = entry.commands.get_or_insert_with(Default::default);
            cmds.set_angular_torques = Some(SetAngularTorquesCmd { torques, execute });
        }
        Command::SetAngularPosition {
            device_id,
            angles,
            execute,
        } => {
            let entry = ensure_inverse3(msg, &device_id);
            let cmds = entry.commands.get_or_insert_with(Default::default);
            cmds.set_angular_position = Some(SetAngularPositionCmd { angles, execute });
        }
        Command::SetI3Transform {
            device_id,
            transform,
            execute,
        } => {
            let entry = ensure_inverse3(msg, &device_id);
            let cmds = entry.commands.get_or_insert_with(Default::default);
            cmds.set_transform = Some(SetTransformCmd { transform, execute });
        }
        Command::ProbePosition { device_id } => {
            let entry = ensure_inverse3(msg, &device_id);
            let cmds = entry.commands.get_or_insert_with(Default::default);
            cmds.probe_position = Some(ProbeCmd {});
        }

        // --- VerseGrip per-tick commands ---
        Command::ProbeOrientation { device_id } => {
            let wvg = ensure_versegrip(&mut msg.wireless_verse_grip, &device_id);
            let wvg_cmds = wvg.commands.get_or_insert_with(Default::default);
            wvg_cmds.probe_orientation = Some(ProbeCmd {});

            let cvg = ensure_versegrip(&mut msg.custom_verse_grip, &device_id);
            let cvg_cmds = cvg.commands.get_or_insert_with(Default::default);
            cvg_cmds.probe_orientation = Some(ProbeCmd {});
        }
        Command::SetExtensionData {
            device_id,
            extension_data,
        } => {
            let wvg = ensure_versegrip(&mut msg.wireless_verse_grip, &device_id);
            let wvg_cmds = wvg.commands.get_or_insert_with(Default::default);
            wvg_cmds.set_extension_data = Some(SetExtensionDataCmd {
                extension_data: extension_data.clone(),
            });

            let cvg = ensure_versegrip(&mut msg.custom_verse_grip, &device_id);
            let cvg_cmds = cvg.commands.get_or_insert_with(Default::default);
            cvg_cmds.set_extension_data = Some(SetExtensionDataCmd { extension_data });
        }
        Command::SetVgTransform {
            device_id,
            transform,
            execute,
        } => {
            let wvg = ensure_versegrip(&mut msg.wireless_verse_grip, &device_id);
            let wvg_cmds = wvg.commands.get_or_insert_with(Default::default);
            wvg_cmds.set_transform = Some(SetTransformCmd {
                transform: transform.clone(),
                execute,
            });

            let cvg = ensure_versegrip(&mut msg.custom_verse_grip, &device_id);
            let cvg_cmds = cvg.commands.get_or_insert_with(Default::default);
            cvg_cmds.set_transform = Some(SetTransformCmd { transform, execute });
        }

        // --- Configure commands ---
        Command::ConfigureInverse3 { device_id, config } => {
            let entry = ensure_inverse3(msg, &device_id);
            entry.configure = Some(config);
        }
        Command::ConfigureVerseGrip { device_id, config } => {
            let wvg = ensure_versegrip(&mut msg.wireless_verse_grip, &device_id);
            wvg.configure = Some(VerseGripConfigure { ..config.clone() });

            let cvg = ensure_versegrip(&mut msg.custom_verse_grip, &device_id);
            cvg.configure = Some(config);
        }
        Command::ConfigureSession(config) => {
            msg.session.configure = Some(config);
        }
        Command::ConfigureSdfOutput {
            device_id,
            state_output,
        } => {
            let entry = ensure_inverse3(msg, &device_id);
            let cfg = entry.configure.get_or_insert_with(Default::default);
            cfg.sdf = Some(SdfOutputConfigure { state_output });
        }

        // --- SDF HFX commands ---
        Command::SdfSet {
            device_id,
            objects,
            from_space,
        } => {
            let entry = ensure_inverse3(msg, &device_id);
            let cmds = entry.commands.get_or_insert_with(Default::default);
            cmds.sdf = Some(SdfCommand {
                mode: SdfCommandMode::Set,
                from_space,
                objects,
            });
        }
        Command::SdfUpdate {
            device_id,
            objects,
        } => {
            let entry = ensure_inverse3(msg, &device_id);
            let cmds = entry.commands.get_or_insert_with(Default::default);
            cmds.sdf = Some(SdfCommand {
                mode: SdfCommandMode::Update,
                from_space: None,
                objects,
            });
        }
        Command::SdfRemove { device_id, ids } => {
            let entry = ensure_inverse3(msg, &device_id);
            let cmds = entry.commands.get_or_insert_with(Default::default);
            cmds.sdf = Some(SdfCommand {
                mode: SdfCommandMode::Remove,
                from_space: None,
                objects: ids
                    .into_iter()
                    .map(|id| SdfHfxObject {
                        id,
                        ..Default::default()
                    })
                    .collect(),
            });
        }

        // --- Session-level ---
        Command::ForceRenderFullState => {
            msg.session.force_render_full_state = Some(ForceRenderFullStateMsg {});
        }
        Command::PingService => {
            // No modification — ping bypasses ServiceMsg entirely
        }
    }
    Ok(())
}

/// Clears one-shot fields after sending. Per-tick commands are preserved.
pub(crate) fn clear_oneshot_fields(msg: &mut ServiceMsg) {
    msg.session.configure = None;
    msg.session.force_render_full_state = None;

    for entry in &mut msg.inverse3 {
        entry.configure = None;
    }
    for entry in &mut msg.verse_grip {
        entry.configure = None;
    }
    for entry in &mut msg.wireless_verse_grip {
        entry.configure = None;
    }
    for entry in &mut msg.custom_verse_grip {
        entry.configure = None;
    }
}

// ============================================================
// State update handlers
// ============================================================

/// Applies verse grip dedup based on the configured mode.
pub fn apply_verse_grip_dedup(data: &mut ServiceData, mode: VerseGripDuplicateMode) {
    match mode {
        VerseGripDuplicateMode::PreferCustom => {
            let custom_ids: Vec<String> = data.custom_verse_grip.iter()
                .map(|d| d.device_id.clone()).collect();
            data.wireless_verse_grip.retain(|d| !custom_ids.contains(&d.device_id));
        }
        VerseGripDuplicateMode::PreferWireless => {
            let wireless_ids: Vec<String> = data.wireless_verse_grip.iter()
                .map(|d| d.device_id.clone()).collect();
            data.custom_verse_grip.retain(|d| !wireless_ids.contains(&d.device_id));
        }
        VerseGripDuplicateMode::KeepBoth => {}
    }
}

/// Updates the device state based on incoming WebSocket messages.
pub async fn update_state_on_message(
    state: Arc<Mutex<TimestampedServiceData>>,
    message: String,
    first: &mut bool,
    verse_grip_mode: VerseGripDuplicateMode,
) {
    let v: Value = match serde_json::from_str(&message) {
        Ok(val) => val,
        Err(e) => {
            error!("Invalid JSON from WS: {} | error: {}", message, e);
            return;
        }
    };


    fn is_full_payload(v: &Value) -> bool {
        for key in &[
            "inverse3",
            "verse_grip",
            "wireless_verse_grip",
            "custom_verse_grip",
        ] {
            if let Some(arr) = v.get(key).and_then(Value::as_array) {
                if arr.iter().any(|item| item.get("config").is_some()) {
                    return true;
                }
            }
        }
        false
    }

    if is_full_payload(&v) {
        match serde_json::from_value::<ServiceData>(v) {
            Ok(sd) => {
                let mut lock = state.lock().await;
                let _ = update_service_struct(&mut lock.data, &sd, verse_grip_mode);
                lock.timestamp = std::time::Instant::now();
                if *first {
                    trace!("Initialized from full payload");
                } else {
                    trace!("Merged full payload");
                }
            }
            Err(e) => {
                if *first {
                    error!("Failed to parse first-message full payload: {}", e);
                } else {
                    error!("Failed to parse full payload: {}", e);
                }
            }
        }
    } else {
        match serde_json::from_value::<ServiceState>(v.clone()) {
            Ok(ss) => {
                update_state_from_service_state(state.clone(), ss, verse_grip_mode).await;
                trace!("Applied state-only update");
            }
            Err(e) => {
                error!("Failed to parse state-only payload: {}", e);
            }
        }
    }

    if *first {
        *first = false;
    }
}

/// Updates the device state using data from a ServiceState message (state-only updates)
async fn update_state_from_service_state(
    state: Arc<Mutex<TimestampedServiceData>>,
    state_only: ServiceState,
    verse_grip_mode: VerseGripDuplicateMode,
) {
    trace!("Updating device state with partial data");
    let mut state_lock = state.lock().await;
    state_lock.timestamp = std::time::Instant::now();

    for msg in &state_only.inverse3 {
        match state_lock
            .data
            .inverse3
            .iter_mut()
            .find(|d| d.device_id == msg.device_id)
        {
            Some(dev) => {
                dev.state = msg.state.clone();
                dev.status = msg.status.clone();
            }
            None => {
                state_lock.data.inverse3.push(Inverse3Device {
                    device_id: msg.device_id.clone(),
                    config: None,
                    state: msg.state.clone(),
                    status: msg.status.clone(),
                });
            }
        }
    }

    for msg in &state_only.verse_grip {
        match state_lock
            .data
            .verse_grip
            .iter_mut()
            .find(|d| d.device_id == msg.device_id)
        {
            Some(dev) => {
                dev.state = msg.state.clone();
                dev.status = msg.status.clone();
            }
            None => {
                state_lock.data.verse_grip.push(VerseGripDevice {
                    device_id: msg.device_id.clone(),
                    config: None,
                    state: msg.state.clone(),
                    status: msg.status.clone(),
                });
            }
        }
    }

    for msg in &state_only.wireless_verse_grip {
        match state_lock
            .data
            .wireless_verse_grip
            .iter_mut()
            .find(|d| d.device_id == msg.device_id)
        {
            Some(dev) => {
                dev.state = msg.state.clone();
                dev.status = msg.status.clone();
            }
            None => {
                state_lock
                    .data
                    .wireless_verse_grip
                    .push(WirelessVerseGripDevice {
                        device_id: msg.device_id.clone(),
                        config: None,
                        state: msg.state.clone(),
                        status: msg.status.clone(),
                    });
            }
        }
    }

    for msg in &state_only.custom_verse_grip {
        match state_lock
            .data
            .custom_verse_grip
            .iter_mut()
            .find(|d| d.device_id == msg.device_id)
        {
            Some(dev) => {
                dev.state = msg.state.clone();
                dev.status = msg.status.clone();
            }
            None => {
                state_lock
                    .data
                    .custom_verse_grip
                    .push(CustomVerseGripDevice {
                        device_id: msg.device_id.clone(),
                        config: None,
                        state: msg.state.clone(),
                        status: msg.status.clone(),
                    });
            }
        }
    }

    state_lock
        .data
        .inverse3
        .retain(|d| state_only.inverse3.iter().any(|msg| msg.device_id == d.device_id));
    state_lock
        .data
        .verse_grip
        .retain(|d| state_only.verse_grip.iter().any(|msg| msg.device_id == d.device_id));
    state_lock.data.wireless_verse_grip.retain(|d| {
        state_only
            .wireless_verse_grip
            .iter()
            .any(|msg| msg.device_id == d.device_id)
    });
    state_lock.data.custom_verse_grip.retain(|d| {
        state_only
            .custom_verse_grip
            .iter()
            .any(|msg| msg.device_id == d.device_id)
    });

    apply_verse_grip_dedup(&mut state_lock.data, verse_grip_mode);

    if state_only.session_id != 0 {
        state_lock.data.session_id = state_only.session_id;
    }
    if state_only.session.is_some() {
        state_lock.data.session = state_only.session;
    }
}

/// Merges incoming device state data with existing state.
pub fn update_service_struct(
    service_data: &mut ServiceData,
    input_data: &ServiceData,
    verse_grip_mode: VerseGripDuplicateMode,
) -> Result<(), Box<dyn std::error::Error + Send + Sync>> {
    for device in &input_data.inverse3 {
        match service_data
            .inverse3
            .iter_mut()
            .find(|d| d.device_id == device.device_id)
        {
            Some(existing_device) => {
                *existing_device = device.clone();
            }
            None => service_data.inverse3.push(device.clone()),
        }
    }

    for device in &input_data.verse_grip {
        match service_data
            .verse_grip
            .iter_mut()
            .find(|d| d.device_id == device.device_id)
        {
            Some(existing_device) => {
                *existing_device = device.clone();
            }
            None => service_data.verse_grip.push(device.clone()),
        }
    }

    for device in &input_data.wireless_verse_grip {
        match service_data
            .wireless_verse_grip
            .iter_mut()
            .find(|d| d.device_id == device.device_id)
        {
            Some(existing_device) => {
                *existing_device = device.clone();
            }
            None => service_data.wireless_verse_grip.push(device.clone()),
        }
    }

    for device in &input_data.custom_verse_grip {
        match service_data
            .custom_verse_grip
            .iter_mut()
            .find(|d| d.device_id == device.device_id)
        {
            Some(existing_device) => {
                *existing_device = device.clone();
            }
            None => service_data.custom_verse_grip.push(device.clone()),
        }
    }

    service_data
        .inverse3
        .retain(|d| input_data.inverse3.iter().any(|nd| nd.device_id == d.device_id));
    service_data
        .verse_grip
        .retain(|d| input_data.verse_grip.iter().any(|nd| nd.device_id == d.device_id));
    service_data.wireless_verse_grip.retain(|d| {
        input_data
            .wireless_verse_grip
            .iter()
            .any(|nd| nd.device_id == d.device_id)
    });
    service_data.custom_verse_grip.retain(|d| {
        input_data
            .custom_verse_grip
            .iter()
            .any(|nd| nd.device_id == d.device_id)
    });

    apply_verse_grip_dedup(service_data, verse_grip_mode);

    if input_data.session_id != 0 {
        service_data.session_id = input_data.session_id;
    }
    if input_data.session.is_some() {
        service_data.session = input_data.session.clone();
    }

    Ok(())
}