haply 0.8.2

use approx::{assert_relative_eq, relative_eq};
use haply::http::InverseHttpClient;
use haply::{device, device_model::*, HaplyDevice};
use std::collections::HashMap;
use std::time::Instant;
use std::net::TcpListener;
// use haply::device::{update_state_on_message, update_service_struct, create_cmd_msgs, create_entire_msg};
use log::debug;
use tokio;

use haply_mock::models::Device as MockDevice;
use haply_mock::models::DeviceInfo as MockDeviceInfo;
use haply_mock::models::Inverse3Info as MockInverse3Info;
use haply_mock::models::VerseGripWirelessInfo as MockVerseGripWirelessInfo;
// use haply_mock::models::I;
use haply_mock::models::Inverse3RTState as MockInverse3RTState;
use haply_mock::models::WirelessVerseGripRTState as MockWirelessVerseGripRTState;
use haply_mock::DeviceRTState as MockDeviceRTState;
use haply_mock::Service as MockService;
use haply_mock::ServiceConfig as MockServiceConfig;

fn inverse_device(id: &str) -> MockDevice {
    MockDevice {
        id: id.to_string(),
        specific_info: MockDeviceInfo::Inverse3(MockInverse3Info::default()),
        ..MockDevice::default()
    }
}

fn versegrip_device(id: &str) -> MockDevice {
    MockDevice {
        id: id.to_string(),
        specific_info: MockDeviceInfo::VerseGripWireless(MockVerseGripWirelessInfo::default()),
        ..MockDevice::default()
    }
}

/// Setup a mock service with given device IDs and default states
///
/// Once set up, the service can be started with `service_mock.start()`
async fn setup_mock_service(inverse_id: &str, versegrip_id: &str) -> (MockService, u16, u16, u16) {
    // Helper to get available port
    fn pick_free_port() -> std::io::Result<u16> {
        let listener = TcpListener::bind(("127.0.0.1", 0))?;
        let port = listener.local_addr()?.port();
        drop(listener);
        Ok(port)
    }

    let http_port = pick_free_port().expect("failed to pick free http port");
    let simulation_ws_port = pick_free_port().expect("failed to pick free simulation ws port");
    let events_ws_port = pick_free_port().expect("failed to pick free events ws port");

    let config = MockServiceConfig {
        http_port,
        simulation_ws_port,
        events_ws_port,
    };
    let service_mock = MockService::new(config);
    let inverse = inverse_device(inverse_id);
    let versegrip = versegrip_device(versegrip_id);
    let inverse_state = MockInverse3RTState::default();
    let versegrip_state = MockWirelessVerseGripRTState::default();
    service_mock
        .set_connected_devices(vec![inverse.clone(), versegrip.clone()])
        .await;
    let mut devices_state = HashMap::new();
    devices_state.insert(inverse_id.to_string(), inverse_state.into());
    devices_state.insert(versegrip_id.to_string(), versegrip_state.into());
    service_mock.set_devices_data(devices_state).await;
    (service_mock, http_port, simulation_ws_port, events_ws_port)
}

fn compare_inverse_leader_with_state(
    leader_state: &Inverse3Device,
    state: &MockInverse3RTState,
) -> bool {
    for i in 0..3 {
        if !relative_eq!(
            leader_state.state.angular_position.unwrap().0[i],
            state.angular_position[i] as f32,
            epsilon = 1e-6
        ) {
            return false;
        }
        if !relative_eq!(
            leader_state.state.angular_velocity.unwrap().0[i],
            state.angular_velocity[i] as f32,
            epsilon = 1e-6
        ) {
            return false;
        }
        if !relative_eq!(
            leader_state.state.cursor_position.unwrap()[i],
            state.cursor_position[i] as f32,
            epsilon = 1e-6
        ) {
            return false;
        }
        if !relative_eq!(
            leader_state.state.cursor_velocity.unwrap()[i],
            state.cursor_velocity[i] as f32,
            epsilon = 1e-6
        ) {
            return false;
        }
    }
    for i in 0..4 {
        if !relative_eq!(
            leader_state.state.body_orientation.unwrap()[i],
            state.body_orientation[i] as f32,
            epsilon = 1e-6
        ) {
            return false;
        }
    }
    true
}

fn compare_versegrip_leader_with_state(
    leader_state: &WirelessVerseGripDevice,
    state: &MockWirelessVerseGripRTState,
) -> bool {
    for i in 0..4 {
        if !relative_eq!(
            leader_state.state.orientation.unwrap()[i],
            state.ori[i] as f32,
            epsilon = 1e-6
        ) {
            return false;
        }
    }
    leader_state.state.buttons.unwrap().a == state.buttons[0]
        && leader_state.state.buttons.unwrap().b == state.buttons[1]
        && leader_state.state.buttons.unwrap().c == state.buttons[2]
}

// http stuff
// test the get_version function
#[tokio::test]
async fn test_get_version() {
    let (mut service, http_port, _sim_ws_port, _events_ws_port) = setup_mock_service("inv1", "vg1").await;
    let src_version = haply_mock::models::Version {
        build_time: "2024-01-01T00:00:00Z".to_string(),
        git_branch: "main".to_string(),
        git_describe: "v1.0.0".to_string(),
        git_hash: "abcdef1234567890".to_string(),
        git_tag: "v1.0.0".to_string(),
        project_name: "haply-service-mock".to_string(),
        project_version: "1.0.0".to_string(),
    };
    service.set_service_version(src_version.clone()).await;
    service.start();

    let client = InverseHttpClient::new(&format!("http://localhost:{}", http_port));
    let version = client.get_version().await.expect("Failed to get version");

    assert_eq!(src_version.build_time, version.build_time);
    assert_eq!(src_version.git_branch, version.git_branch);
    assert_eq!(src_version.git_describe, version.git_describe);
    assert_eq!(src_version.git_hash, version.git_hash);
    assert_eq!(src_version.git_tag, version.git_tag);
    assert_eq!(src_version.project_name, version.project_name);
    assert_eq!(src_version.project_version, version.project_version);
}

// test the get_devices function
#[tokio::test]
async fn test_http_get_devices() {
    let (mut service, http_port, _sim_ws_port, _events_ws_port) = setup_mock_service("inv1", "vg1").await;
    service.start();

    let client = InverseHttpClient::new(&format!("http://localhost:{}", http_port));
    let devices = client.get_devices().await.expect("Failed to get devices");

    assert_eq!(devices.len(), 2);
    assert!(devices
        .iter()
        .any(|d| matches!(d, Config::DeviceConfig(_)) && d.id() == "inv1"));
    assert!(devices
        .iter()
        .any(|d| matches!(d, Config::WVGConfig(_)) && d.id() == "vg1"));
}

// websocket stuff
// test the list_devices function
#[tokio::test]
async fn test_device_list_devices() {
    let (mut service, http_port, sim_ws_port, _events_ws_port) = setup_mock_service("inv1", "vg1").await;
    service.start();

    let device = HaplyDevice::new(&format!("http://localhost:{}", http_port), &format!("ws://localhost:{}", sim_ws_port))
        .await
        .expect("Failed to create device");

    let devices = device.list_devices().await.expect("Failed to list devices");

    assert_eq!(devices.len(), 2);
    assert!(devices
        .iter()
        .any(|d| matches!(d, Config::DeviceConfig(_)) && d.id() == "inv1"));
    assert!(devices
        .iter()
        .any(|d| matches!(d, Config::WVGConfig(_)) && d.id() == "vg1"));
}

// test the update_force function
#[tokio::test]
async fn test_update_force() {
    let (mut service, http_port, sim_ws_port, _events_ws_port) = setup_mock_service("inv1", "vg1").await;
    service.start();
    let mut device = HaplyDevice::new(&format!("http://localhost:{}", http_port), &format!("ws://localhost:{}", sim_ws_port))
        .await
        .expect("Failed to create device");
    let forces = vec![ForceInput {
        device_id: "inv1".to_string(),
        forces: Force {
            x: 1.0,
            y: 2.0,
            z: 2.5,
        },
    }];
    device
        .update_force(forces.clone(), Some(true), Some(true))
        .await
        .expect("Failed to update force");

    let mut got_target = false;
    let start = std::time::Instant::now();
    while start.elapsed().as_millis() < 1000 {
        let targets = service.get_inverse_targets().await;
        if let Some(target) = targets.get("inv1") {
            if target.is_empty() {
                continue;
            }
            assert_eq!(target.len(), 1);
            assert_relative_eq!(target[0].0[0], 1.0, epsilon = 1e-6);
            assert_relative_eq!(target[0].0[1], 2.0, epsilon = 1e-6);
            assert_relative_eq!(target[0].0[2], 2.5, epsilon = 1e-6);
            got_target = true;
            break;
        }
        tokio::time::sleep(tokio::time::Duration::from_millis(10)).await;
    }
    assert!(got_target, "Did not receive inverse targets for inv1");
}

#[tokio::test]
async fn test_update_force_multiple_devices() {
    // Create mock service with multiple devices
    let (mut service, http_port, sim_ws_port, _events_ws_port) = setup_mock_service("inverse_id", "versegrip_id").await;

    let inverse = inverse_device("inv1");
    let versegrip = versegrip_device("vg1");
    let inverse2 = inverse_device("inv2");
    let versegrip2 = versegrip_device("vg2");
    let inverse3 = inverse_device("inv3");

    let inverse_state = MockInverse3RTState::default();
    let versegrip_state = MockWirelessVerseGripRTState::default();
    service
        .set_connected_devices(vec![inverse, versegrip, inverse2, versegrip2, inverse3])
        .await;
    let mut devices_state = HashMap::new();
    devices_state.insert("inv1".to_string(), inverse_state.clone().into());
    devices_state.insert("vg1".to_string(), versegrip_state.clone().into());
    devices_state.insert("inv2".to_string(), inverse_state.clone().into());
    devices_state.insert("vg2".to_string(), versegrip_state.clone().into());
    devices_state.insert("inv3".to_string(), inverse_state.clone().into());
    service.set_devices_data(devices_state).await;

    service.start();

    let mut device = HaplyDevice::new(&format!("http://localhost:{}", http_port), &format!("ws://localhost:{}", sim_ws_port))
        .await
        .expect("Failed to create device");
    let forces = vec![
        ForceInput {
            device_id: "inv1".to_string(),
            forces: Force {
                x: 1.0,
                y: 2.0,
                z: 2.5,
            },
        },
        ForceInput {
            device_id: "inv2".to_string(),
            forces: Force {
                x: 0.5,
                y: 1.5,
                z: 2.0,
            },
        },
        ForceInput {
            device_id: "inv3".to_string(),
            forces: Force {
                x: 0.0,
                y: 1.0,
                z: 1.5,
            },
        },
    ];

    device
        .update_force(forces.clone(), Some(true), None)
        .await
        .expect("Failed to update force");

    // Check that forces are not sent to devices without send force flag set

    let mut got_target = false;
    let start = std::time::Instant::now();
    while start.elapsed().as_millis() < 200 {
        let targets = service.get_inverse_targets().await;
        if targets.iter().any(|el| !el.1.is_empty()) {
            got_target = true;
        }
        tokio::time::sleep(tokio::time::Duration::from_millis(10)).await;
    }
    assert!(!got_target, "Expected no inverse targets to be received");

    device.send_command().await.expect("Failed to send command");

    let mut got_targets = HashMap::new();
    let start = std::time::Instant::now();
    while start.elapsed().as_millis() < 200 {
        let targets = service.get_inverse_targets().await;
        for (device_id, target) in targets {
            if !target.is_empty() {
                got_targets.insert(device_id, target);
            }
        }
        tokio::time::sleep(tokio::time::Duration::from_millis(10)).await;
    }
    assert_eq!(got_targets.len(), 3);
    let target1 = got_targets
        .get("inv1")
        .expect("Did not receive target for inv1");
    assert_eq!(target1.len(), 1);
    assert_relative_eq!(target1[0].0[0], 1.0, epsilon = 1e-6);
    assert_relative_eq!(target1[0].0[1], 2.0, epsilon = 1e-6);
    assert_relative_eq!(target1[0].0[2], 2.5, epsilon = 1e-6);

    let target2 = got_targets
        .get("inv2")
        .expect("Did not receive target for inv2");
    assert_eq!(target2.len(), 1);
    assert_relative_eq!(target2[0].0[0], 0.5, epsilon = 1e-6);
    assert_relative_eq!(target2[0].0[1], 1.5, epsilon = 1e-6);
    assert_relative_eq!(target2[0].0[2], 2.0, epsilon = 1e-6);

    let target3 = got_targets
        .get("inv3")
        .expect("Did not receive target for inv3");
    assert_eq!(target3.len(), 1);
    assert_relative_eq!(target3[0].0[0], 0.0, epsilon = 1e-6);
    assert_relative_eq!(target3[0].0[1], 1.0, epsilon = 1e-6);
    assert_relative_eq!(target3[0].0[2], 1.5, epsilon = 1e-6);
}

// test the read_state function
#[tokio::test]
async fn test_read_state() {
    let (mut service, http_port, sim_ws_port, _events_ws_port) = setup_mock_service("inv1", "vg1").await;
    let inverse_state = MockDeviceRTState::Inverse3(MockInverse3RTState::new_random());
    let versegrip_state =
        MockDeviceRTState::WirelessVerseGrip(MockWirelessVerseGripRTState::new_random());

    service
        .update_device_states("inv1", inverse_state.clone())
        .await;
    service
        .update_device_states("vg1", versegrip_state.clone())
        .await;

    service.start();

    let device = HaplyDevice::new(&format!("http://localhost:{}", http_port), &format!("ws://localhost:{}", sim_ws_port))
        .await
        .expect("Failed to create device");
    let mut state = None;
    let start = Instant::now();
    while start.elapsed().as_millis() < 200 {
        let st = device.read_state().await.unwrap();
        if !st.inverse3.is_empty() && !st.wireless_verse_grip.is_empty() {
            state = Some(st);
            break;
        }
        tokio::time::sleep(tokio::time::Duration::from_millis(10)).await;
    }

    let state = state.expect("Did not receive full device state");

    assert!(state.verse_grip.is_empty());
    assert!(state.custom_verse_grip.is_empty());
    assert_eq!(state.inverse3.len(), 1);
    assert_eq!(state.inverse3[0].device_id, "inv1");
    let inv_state = match inverse_state {
        MockDeviceRTState::Inverse3(s) => s,
        _ => panic!("Expected Inverse3 state"),
    };

    assert!(compare_inverse_leader_with_state(
        &state.inverse3[0],
        &inv_state
    ));

    assert_eq!(state.wireless_verse_grip.len(), 1);
    assert_eq!(state.wireless_verse_grip[0].device_id, "vg1");
    let vg_state = match versegrip_state {
        MockDeviceRTState::WirelessVerseGrip(s) => s,
        _ => panic!("Expected WirelessVerseGrip state"),
    };

    assert!(compare_versegrip_leader_with_state(
        &state.wireless_verse_grip[0],
        &vg_state
    ));
}


#[tokio::test]
async fn test_read_state_updates() {
    let (mut service, http_port, sim_ws_port, _events_ws_port) = setup_mock_service("inv1", "vg1").await;

    service.start();

    let mut device = HaplyDevice::new(&format!("http://localhost:{}", http_port), &format!("ws://localhost:{}", sim_ws_port))
        .await
        .expect("Failed to create device");
    let mut state = None;
    let start = Instant::now();
    while start.elapsed().as_millis() < 200 {
        let st = device.read_state().await.unwrap();
        if !st.inverse3.is_empty() && !st.wireless_verse_grip.is_empty() {
            state = Some(st);
            break;
        }
        tokio::time::sleep(tokio::time::Duration::from_millis(10)).await;
    }

    // We got initial state, now update should happen only on command send
    let _ = state.expect("Did not receive full device state");

    // Updating state to new random values
    let inverse_state = MockDeviceRTState::Inverse3(MockInverse3RTState::new_random());
    let versegrip_state =
        MockDeviceRTState::WirelessVerseGrip(MockWirelessVerseGripRTState::new_random());

    service
        .update_device_states("inv1", inverse_state.clone())
        .await;
    service
        .update_device_states("vg1", versegrip_state.clone())
        .await;

    // Reading all state updates for 200 ms
    let mut state = None;
    let start = Instant::now();
    while start.elapsed().as_millis() < 200 {
        let st = device.read_state().await.unwrap();
        if !st.inverse3.is_empty() && !st.wireless_verse_grip.is_empty() {
            state = Some(st);
        }
        tokio::time::sleep(tokio::time::Duration::from_millis(10)).await;
    }

    let state = state.expect("Did not receive full device state");
    
    // Check that state matches default values (update didn't happen yet)
    assert!(compare_inverse_leader_with_state(&state.inverse3[0], &MockInverse3RTState::default()));
    assert!(compare_versegrip_leader_with_state(&state.wireless_verse_grip[0], &MockWirelessVerseGripRTState::default()));

    // Sending force render command to trigger state message
    device.send_command().await.expect("Failed to send command");

        let mut state = None;
    let start = Instant::now();
    while start.elapsed().as_millis() < 200 {
        let st = device.read_state().await.unwrap();
        if !st.inverse3.is_empty() && !st.wireless_verse_grip.is_empty() {
            state = Some(st);
        }
        tokio::time::sleep(tokio::time::Duration::from_millis(10)).await;
    }

    let state = state.expect("Did not receive full device state");

    let inv_state = match inverse_state {
        MockDeviceRTState::Inverse3(s) => s,
        _ => panic!("Expected Inverse3 state"),
    };
    assert!(compare_inverse_leader_with_state(&state.inverse3[0], &inv_state));

    let wg_state = match versegrip_state {
        MockDeviceRTState::WirelessVerseGrip(s) => s,
        _ => panic!("Expected WirelessVerseGrip state"),
    };
    assert!(compare_versegrip_leader_with_state(&state.wireless_verse_grip[0], &wg_state));
}

//todo: add more tests for other functions wihen mock is extended
// test the update_position function
// #[tokio::test]
// async fn test_update_position() {
//     let mut device = HaplyDevice::new("http://localhost:10000", "ws://localhost:10001")
//         .await
//         .expect("Failed to create device");
//     let positions = vec![PositionInput {
//         device_id: "device1".to_string(),
//         positions: Linear3D {
//             x: 1.0,
//             y: 2.0,
//             z: 3.0,
//         },
//     }];
//     device
//         .update_position(positions, Some(true))
//         .await
//         .expect("Failed to update position");
// }

// // test the update_extension_data function
// #[tokio::test]
// async fn test_update_extension_data() {
//     let mut device = HaplyDevice::new("http://localhost:10000", "ws://localhost:10001")
//         .await
//         .expect("Failed to create device");
//     let ext_data = vec![ExtensionDataInput {
//         device_id: "device1".to_string(),
//         extension_data: vec![0; 20],
//     }];
//     device
//         .update_extension_data(ext_data, Some(true))
//         .await
//         .expect("Failed to update extension data");
// }

// // test the probe_cursor_position function
// #[tokio::test]
// async fn test_probe_cursor_position() {
//     let mut device = HaplyDevice::new("http://localhost:10000", "ws://localhost:10001")
//         .await
//         .expect("Failed to create device");
//     let device_ids = vec!["device1".to_string()];
//     let cursor_position = device
//         .probe_cursor_position(device_ids, Some(true))
//         .await
//         .expect("Failed to probe cursor position");
//     debug!("Cursor position: {:?}", cursor_position);
// }

// // test the probe_angular_position function
// #[tokio::test]
// async fn test_probe_angular_position() {
//     let mut device = HaplyDevice::new("http://localhost:10000", "ws://localhost:10001")
//         .await
//         .expect("Failed to create device");
//     let device_ids = vec!["device1".to_string()];
//     let angular_position = device
//         .probe_angular_position(device_ids, Some(true))
//         .await
//         .expect("Failed to probe angular position");
//     debug!("Angular position: {:?}", angular_position);
// }

// // test the probe_orientation function
// #[tokio::test]
// async fn test_probe_orientation() {
//     let mut device = HaplyDevice::new("http://localhost:10000", "ws://localhost:10001")
//         .await
//         .expect("Failed to create device");
//     let device_ids = vec!["device1".to_string()];
//     let orientation = device
//         .probe_orientation(device_ids, Some(true))
//         .await
//         .expect("Failed to probe orientation");
//     debug!("Orientation: {:?}", orientation);
// }