chuniio-yubideck/chuniio/src/lib.rs

#![allow(clippy::missing_safety_doc)]

mod configuration;

use std::{
    ffi::c_void,
    rc::Rc,
    sync::{
        atomic::Ordering,
        Mutex,
    },
    time::Duration,
};

use ::log::error;
use shared_memory::Shmem;
use winapi::{
    shared::{
        minwindef::{BOOL, DWORD, HINSTANCE, LPVOID, TRUE},
        winerror::{E_FAIL, S_OK},
    },
    um::winnt::{DLL_PROCESS_ATTACH, HRESULT},
};
use yubideck_common::{create_shared_memory, init_logger, INPUT_SHMEM_SIZE, OUTPUT_SHMEM_SIZE};

static mut INPUT_SHMEM: Option<Rc<Shmem>> = None;
static mut OUTPUT_SHMEM: Option<Rc<Mutex<Shmem>>> = None;
pub static DEVICE_POLLING_INTERVAL: Duration = Duration::from_millis(100);

cfg_if::cfg_if! {
    if #[cfg(any(chuni, chusanapp))] {
        use std::{sync::{atomic::AtomicBool, RwLock}, thread};

        use anyhow::Result;
        use ::log::info;
        use rusb::{DeviceHandle, GlobalContext};

        type SliderCallbackFn = unsafe extern "C" fn(data: *const u8);

        const TIMEOUT: Duration = Duration::from_millis(20);

        static DEVICE: RwLock<Option<DeviceHandle<GlobalContext>>> = RwLock::new(None);
        static SLIDER_ACTIVE: AtomicBool = AtomicBool::new(false);
    }
}

cfg_if::cfg_if! {
    if #[cfg(any(chuni, amdaemon))] {
        use std::{ffi::{c_int, CString}, sync::atomic::AtomicU16};

        use lazy_static::lazy_static;
        use winapi::um::{winuser::GetAsyncKeyState, winbase::GetPrivateProfileIntA};

        use crate::configuration::Configuration;

        static COIN_COUNT: AtomicU16 = AtomicU16::new(0);
        static mut COIN_PRESSED: bool = false;

        lazy_static! {
            static ref CONFIGURATION: Configuration = {
                let io3 = CString::new("io3").unwrap();
                let test = CString::new("test").unwrap();
                let service = CString::new("service").unwrap();
                let coin = CString::new("coin").unwrap();
                let cfg_file = CString::new(".\\segatools.ini").unwrap();

                unsafe {
                    Configuration {
                        test_key: GetPrivateProfileIntA(
                            io3.as_ptr(),
                            test.as_ptr(),
                            0x31,
                            cfg_file.as_ptr(),
                        ),
                        service_key: GetPrivateProfileIntA(
                            io3.as_ptr(),
                            service.as_ptr(),
                            0x32,
                            cfg_file.as_ptr(),
                        ),
                        coin_key: GetPrivateProfileIntA(
                            io3.as_ptr(),
                            coin.as_ptr(),
                            0x33,
                            cfg_file.as_ptr(),
                        ),
                    }
                }
            };
        }
    }
}

#[no_mangle]
extern "system" fn DllMain(_dll_module: HINSTANCE, call_reason: DWORD, _reserved: LPVOID) -> BOOL {
    if call_reason == DLL_PROCESS_ATTACH {
        init_logger()
    }

    TRUE
}

#[no_mangle]
pub extern "C" fn chuni_io_get_api_version() -> u16 {
    0x0102
}

#[no_mangle]
pub extern "C" fn chuni_io_jvs_init() -> HRESULT {
    match create_shared_memory("YubideckInput", INPUT_SHMEM_SIZE, true) {
        Ok(s) => unsafe { INPUT_SHMEM = Some(Rc::new(s)) },
        Err(e) => {
            error!("Could not acquire shared memory for YubiDeck input: {e:#?}");
            return E_FAIL;
        }
    }

    match create_shared_memory("YubideckOutput", OUTPUT_SHMEM_SIZE, false) {
        Ok(s) => unsafe { OUTPUT_SHMEM = Some(Rc::new(Mutex::new(s))) },
        Err(e) => {
            error!("Could not obtain shared memory for YubiDeck output: {e:#?}");
            return E_FAIL;
        }
    }

    #[cfg(not(amdaemon))]
    {
        thread::spawn(yubideck_init);
    }

    S_OK
}

#[no_mangle]
#[cfg(any(chuni, amdaemon))]
pub unsafe extern "C" fn chuni_io_jvs_poll(opbtn: *mut u8, beams: *mut u8) {
    if opbtn.is_null() || beams.is_null() {
        return;
    }

    let Some(input_shmem) = (unsafe { &INPUT_SHMEM }) else {
        return;
    };
    let input = unsafe { input_shmem.as_slice() };
    let ir_value = input[0];
    let mut buttons = input[1] & 3; // Buttons are in order: coin, service, test. We take the last 2 bits

    if GetAsyncKeyState(CONFIGURATION.test_key as c_int) != 0 {
        buttons |= 1;
    }

    if GetAsyncKeyState(CONFIGURATION.service_key as c_int) != 0 {
        buttons |= 2;
    }

    *opbtn = buttons;

    // Swap adjacent bit pairs
    *beams = ((ir_value & 0xAA) >> 1) | ((ir_value & 0x55) << 1);
}

#[no_mangle]
#[cfg(chusanapp)]
pub extern "C" fn chuni_io_jvs_poll(_opbtn: *mut u8, _beams: *mut u8) {}

#[no_mangle]
#[cfg(any(chuni, amdaemon))]
pub unsafe extern "C" fn chuni_io_jvs_read_coin_counter(total: *mut u16) {
    if total.is_null() {
        return;
    }

    let Some(input_shmem) = &INPUT_SHMEM else {
        return;
    };
    let input = unsafe { input_shmem.as_slice() };
    let coin_pressed = (input[1] & 4) != 0;

    if coin_pressed || unsafe { GetAsyncKeyState(CONFIGURATION.coin_key as c_int) } != 0 {
        if !COIN_PRESSED {
            COIN_PRESSED = true;
            COIN_COUNT.fetch_add(1, Ordering::Relaxed);
        }
    } else {
        COIN_PRESSED = false;
    }

    *total = COIN_COUNT.load(Ordering::Relaxed);
}

#[no_mangle]
#[cfg(chusanapp)]
pub extern "C" fn chuni_io_jvs_read_coin_counter(_total: *mut u16) {}

#[no_mangle]
pub extern "C" fn chuni_io_slider_init() -> HRESULT {
    #[cfg(any(chuni, chusanapp))]
    {
        let Some(out_shmem) = (unsafe { &OUTPUT_SHMEM }) else {
            error!("OUTPUT_SHMEM is unset.");
            return E_FAIL;
        };

        let Ok(mut out_shmem) = out_shmem.lock() else {
            error!("Could not acquire mutex of output shared memory");
            return E_FAIL;
        };

        let buf = unsafe { out_shmem.as_slice_mut() };

        buf[0] = 0;
        buf[61] = 1;
    }

    S_OK
}

#[no_mangle]
#[cfg(any(chuni, chusanapp))]
pub unsafe extern "C" fn chuni_io_slider_start(callback: *const c_void) {
    if callback.is_null() {
        return;
    }

    if SLIDER_ACTIVE.load(Ordering::SeqCst) {
        return;
    }

    SLIDER_ACTIVE.store(true, Ordering::Relaxed);

    let callback = std::mem::transmute::<_, SliderCallbackFn>(callback);

    thread::spawn(move || {
        let Some(input_shmem) = (unsafe { &INPUT_SHMEM }) else {
            return;
        };
        let usb_in = input_shmem.as_slice();
        let mut pressure = [0u8; 32];

        while SLIDER_ACTIVE.load(Ordering::Relaxed) {
            for (i, p) in usb_in.iter().skip(2).take(32).enumerate() {
                pressure[if i % 2 == 0 { 30 - i } else { 32 - i }] = *p
            }

            callback(pressure.as_ptr());
            thread::sleep(Duration::from_nanos(1_000_000));
        }
    });
}

#[no_mangle]
#[cfg(amdaemon)]
pub extern "C" fn chuni_io_slider_start(_callback: *const c_void) {}

#[no_mangle]
#[cfg(any(chuni, chusanapp))]
pub extern "C" fn chuni_io_slider_stop() {
    SLIDER_ACTIVE.store(false, Ordering::Relaxed);
}

#[no_mangle]
#[cfg(amdaemon)]
pub extern "C" fn chuni_io_slider_stop() {}

#[no_mangle]
#[cfg(any(chuni, chusanapp))]
pub unsafe extern "C" fn chuni_io_slider_set_leds(rgb: *const u8) {
    if rgb.is_null() {
        return;
    }

    let device_rg = DEVICE.read().unwrap();
    let Some(device) = device_rg.as_ref().map(|z| z) else { return };
    let Some(out_shmem) = &OUTPUT_SHMEM else {
        error!("OUTPUT_SHMEM is unset.");
        return;
    };
    let Ok(mut out_shmem) = out_shmem.lock() else {
        error!("Could not acquire mutex of output shared memory");
        return;
    };
    let buf = out_shmem.as_slice_mut();
    let ground = std::slice::from_raw_parts(rgb, 93);

    for (buf_chunk, state_chunk) in buf[1..61]
        .chunks_mut(3)
        .zip(ground.chunks(3).take(20))
    {
        buf_chunk[0] = state_chunk[1];
        buf_chunk[1] = state_chunk[2];
        buf_chunk[2] = state_chunk[0];
    }

    for (buf_chunk, state_chunk) in buf[62..95]
        .chunks_mut(3)
        .zip(ground.chunks(3).skip(20).take(11))
    {
        buf_chunk[0] = state_chunk[1];
        buf_chunk[1] = state_chunk[2];
        buf_chunk[2] = state_chunk[0];
    }

    if let Err(e) = device.write_interrupt(0x02, &buf[0..61], TIMEOUT) {
        error!("Error writing first batch of output data: {e:#?}");
    }

    if let Err(e) = device.write_interrupt(0x02, &buf[61..122], TIMEOUT) {
        error!("Error writing second batch of output data: {e:#?}");
    }
}

#[no_mangle]
#[cfg(amdaemon)]
pub extern "C" fn chuni_io_slider_set_leds(_rgb: *const u8) {}

#[no_mangle]
pub extern "C" fn chuni_io_led_init() -> HRESULT {
    S_OK
}

#[no_mangle]
#[cfg(any(chuni, chusanapp))]
pub unsafe extern "C" fn chuni_io_led_set_colors(board: u8, rgb: *const u8) {
    let device_rg = DEVICE.read().unwrap();
    let Some(device) = device_rg.as_ref().map(|z| z) else { return };
    let Some(out_shmem) = &OUTPUT_SHMEM else {
        error!("OUTPUT_SHMEM is unset.");
        return;
    };
    let Ok(mut out_shmem) = out_shmem.lock() else {
        error!("Could not acquire mutex of output shared memory");
        return;
    };
    let buf = out_shmem.as_slice_mut();
    let data = std::slice::from_raw_parts(rgb, 183);

    match board {
        0 => {
            // left air
            buf[95] = data[150];
            buf[96] = data[151];
            buf[97] = data[152];
        }
        1 => {
            // right air
            buf[98] = data[180];
            buf[99] = data[181];
            buf[100] = data[182];
        }
        _ => {}
    }

    if let Err(e) = device.write_interrupt(0x02, &buf[61..122], TIMEOUT) {
        error!("Error writing second batch of output data: {e:#?}");
    }
}

#[no_mangle]
#[cfg(amdaemon)]
pub extern "C" fn chuni_io_led_set_colors(_rgb: *const u8) {}

#[cfg(any(chuni, chusanapp))]
fn yubideck_init() -> Result<()> {
    {
        let mut global_device = DEVICE.write().unwrap();
        *global_device = None;
    }

    info!("Waiting for device...");

    loop {
        let Some(mut device) =
            rusb::open_device_with_vid_pid(0x1973, 0x2001)
        else {
            thread::sleep(DEVICE_POLLING_INTERVAL);
            continue;
        };

        device.set_active_configuration(1)?;
        device.claim_interface(0)?;

        let mut global_device = DEVICE.write().unwrap();
        *global_device = Some(device);

        thread::spawn(input_thread_proc);

        return Ok(())
    }
}

#[cfg(any(chuni, chusanapp))]
fn input_thread_proc() {    
    info!("Input thread started");

    let device_rg = DEVICE.read().unwrap();
    let Some(device) = device_rg.as_ref().map(|z| z) else { return };

    let mut input_shmem = match create_shared_memory("YubideckInput", INPUT_SHMEM_SIZE, false) {
        Ok(s) => s,
        Err(e) => {
            error!("Could not obtain shared memory for YubiDeck input: {e:#}");
            return;
        }
    };
    let usb_in = unsafe { input_shmem.as_slice_mut() };

    let mut output_shmem = match create_shared_memory("YubideckOutput", OUTPUT_SHMEM_SIZE, false) {
        Ok(s) => s,
        Err(e) => {
            error!("Could not obtain shared memory for YubiDeck output: {e:#?}");
            return;
        }
    };
    let usb_out = unsafe { output_shmem.as_slice_mut() };

    loop {
        if let Err(e) = device.read_interrupt(0x81, usb_in, TIMEOUT) {
            match e {
                rusb::Error::NoDevice | rusb::Error::Io => {
                    error!("Controller disconnected.");

                    usb_in.iter_mut().for_each(|m| *m = 0);

                    // Spawn a thread polling for a connection again
                    thread::spawn(yubideck_init);

                    return;
                }
                _ => error!("Failed to read data from YubiDeck: {e:#}"),
            }
        }

        // Update reader LED data
        if usb_out[122] == 1 {
            usb_out[122] = 0;

            if let Err(e) = device.write_interrupt(0x02, &usb_out[61..122], TIMEOUT) {
                error!("Error writing second batch of output data: {e:#?}");
            }
        }
    }
}