chuniio-yubideck/src/lib.rs

mod configuration;
mod log;

#[cfg(any(
    all(feature = "chusan", target_arch = "x86_64"),
    not(feature = "chusan")
))]
use std::{ffi::c_int, sync::atomic::AtomicU16};
use std::{
    ffi::{c_void, CString},
    fmt::Display,
    sync::{
        atomic::{AtomicBool, Ordering},
        Arc, OnceLock,
    },
    thread,
    time::Duration,
};
#[cfg(any(all(feature = "chusan", target_arch = "x86"), not(feature = "chusan")))]
use std::{sync::Mutex, thread::JoinHandle};

use ::log::{error, info};
use anyhow::{anyhow, Result};
use lazy_static::lazy_static;
use rusb::{DeviceHandle, GlobalContext};
use shared_memory::{Shmem, ShmemConf, ShmemError};
#[cfg(any(
    all(feature = "chusan", target_arch = "x86_64"),
    not(feature = "chusan")
))]
use winapi::um::winuser::GetAsyncKeyState;
use winapi::{
    shared::{
        minwindef::{BOOL, DWORD, HINSTANCE, LPVOID, TRUE},
        winerror::{E_FAIL, S_OK},
    },
    um::{
        winbase::GetPrivateProfileIntA,
        winnt::{DLL_PROCESS_ATTACH, HRESULT},
    },
};

use crate::{configuration::Configuration, log::init_logger};

#[cfg(any(all(feature = "chusan", target_arch = "x86"), not(feature = "chusan")))]
type SliderCallbackFn = unsafe extern "C" fn(data: *const u8);

static DEVICE: OnceLock<DeviceHandle<GlobalContext>> = OnceLock::new();
static mut INPUT_SHMEM: Option<Arc<Shmem>> = None;

#[cfg(any(all(feature = "chusan", target_arch = "x86"), not(feature = "chusan")))]
static mut SLIDER_THREAD: OnceLock<JoinHandle<()>> = OnceLock::new();

#[cfg(any(all(feature = "chusan", target_arch = "x86"), not(feature = "chusan")))]
static SLIDER_ACTIVE: AtomicBool = AtomicBool::new(false);

#[cfg(any(
    all(feature = "chusan", target_arch = "x86_64"),
    not(feature = "chusan")
))]
static COIN_COUNT: AtomicU16 = AtomicU16::new(0);

#[cfg(any(
    all(feature = "chusan", target_arch = "x86_64"),
    not(feature = "chusan")
))]
static COIN_PRESSED: AtomicBool = AtomicBool::new(false);

#[cfg(any(all(feature = "chusan", target_arch = "x86"), not(feature = "chusan")))]
static mut OUTPUT_SHMEM: Option<Arc<Mutex<Shmem>>> = None;

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 {
    match call_reason {
        DLL_PROCESS_ATTACH => init_logger(),
        _ => {}
    }

    TRUE
}

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

#[no_mangle]
#[cfg(any(
    all(feature = "chusan", target_arch = "x86_64"),
    not(feature = "chusan")
))]
pub extern "C" fn chuni_io_jvs_init() -> HRESULT {
    if cfg!(not(feature = "chusan")) {
        if let Err(e) = yubideck_init() {
            error!("Could not initialize YubiDeck: {e:#?}");
            return E_FAIL;
        }
    }

    return create_input_shared_memory();
}

#[no_mangle]
#[cfg(all(feature = "chusan", target_arch = "x86"))]
pub extern "C" fn chuni_io_jvs_init() -> HRESULT {
    S_OK
}

#[no_mangle]
#[cfg(any(
    all(feature = "chusan", target_arch = "x86_64"),
    not(feature = "chusan")
))]
pub 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

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

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

        *opbtn = buttons;
        *beams = ir_value;
    }
}

#[no_mangle]
#[cfg(all(feature = "chusan", target_arch = "x86"))]
pub extern "C" fn chuni_io_jvs_poll(_opbtn: *mut u8, _beams: *mut u8) {}

#[no_mangle]
#[cfg(any(
    all(feature = "chusan", target_arch = "x86_64"),
    not(feature = "chusan")
))]
pub extern "C" fn chuni_io_jvs_read_coin_counter(total: *mut u16) {
    if total.is_null() {
        return;
    }

    let Some(input_shmem) = (unsafe { &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 {
        let coin_previously_pressed = COIN_PRESSED.fetch_or(true, Ordering::Relaxed);
        if !coin_previously_pressed {
            COIN_COUNT.fetch_add(1, Ordering::Relaxed);
        }
    } else {
        COIN_PRESSED.store(false, Ordering::Relaxed);
    }

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

#[no_mangle]
#[cfg(all(feature = "chusan", target_arch = "x86"))]
pub extern "C" fn chuni_io_jvs_read_coin_counter(_total: *mut u16) {}

#[no_mangle]
#[cfg(any(all(feature = "chusan", target_arch = "x86"), not(feature = "chusan")))]
pub unsafe extern "C" fn chuni_io_slider_init() -> HRESULT {
    match create_shared_memory("Local\\YubideckOutput", 122, true) {
        Ok(s) => OUTPUT_SHMEM = Some(Arc::new(Mutex::new(s))),
        Err(e) => {
            error!("Could not obtain shared memory: {e:#?}");
            return E_FAIL;
        }
    }

    if cfg!(not(feature = "chusan")) {
        // Already initialized in chuni_io_jvs_init()
        return S_OK;
    }

    if let Err(e) = yubideck_init() {
        error!("Failed to initialize YubiDeck: {e:#?}");
        return E_FAIL;
    }

    return create_input_shared_memory();
}

#[no_mangle]
#[cfg(all(feature = "chusan", target_arch = "x86_64"))]
pub extern "C" fn chuni_io_slider_init(_callback: *const c_void) -> HRESULT {
    S_OK
}

#[no_mangle]
#[cfg(any(all(feature = "chusan", target_arch = "x86"), not(feature = "chusan")))]
pub unsafe extern "C" fn chuni_io_slider_start(callback: *const c_void) {
    if callback.is_null() {
        return;
    }

    if SLIDER_THREAD.get().is_some() {
        return;
    }

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

    let callback = std::mem::transmute::<_, SliderCallbackFn>(callback);
    let thread = 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) {
            pressure.copy_from_slice(&usb_in[2..34]);

            for i in 0..16 {
                pressure.swap(i * 2, i * 2 + 1);
            }

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

    SLIDER_THREAD.set(thread).unwrap();
}

#[no_mangle]
#[cfg(all(feature = "chusan", target_arch = "x86_64"))]
pub extern "C" fn chuni_io_slider_start(_callback: *const c_void) {}

#[no_mangle]
#[cfg(any(all(feature = "chusan", target_arch = "x86"), not(feature = "chusan")))]
pub extern "C" fn chuni_io_slider_stop() {
    let Some(thread) = (unsafe { SLIDER_THREAD.take() }) else {
        return;
    };

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

    thread.join().expect("Couldn't join slider input thread");
}

#[no_mangle]
#[cfg(all(feature = "chusan", target_arch = "x86_64"))]
pub extern "C" fn chuni_io_slider_stop() {}

#[no_mangle]
#[cfg(any(all(feature = "chusan", target_arch = "x86"), not(feature = "chusan")))]
pub unsafe extern "C" fn chuni_io_slider_set_leds(rgb: *const u8) {
    if rgb.is_null() {
        return;
    }

    let Some(device) = DEVICE.get() 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);

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

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

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

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

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

#[no_mangle]
#[cfg(all(feature = "chusan", target_arch = "x86_64"))]
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(all(feature = "chusan", target_arch = "x86"), not(feature = "chusan")))]
pub unsafe extern "C" fn chuni_io_led_set_colors(board: u8, rgb: *const u8) {
    let Some(device) = DEVICE.get() 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);

    buf[61] = 1;

    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..], Duration::from_millis(20)) {
        error!("Error writing second batch of output data: {e:#?}");
    }
}

#[no_mangle]
#[cfg(all(feature = "chusan", target_arch = "x86_64"))]
pub extern "C" fn chuni_io_led_set_colors(_rgb: *const u8) {}

fn create_input_shared_memory() -> HRESULT {
    match create_shared_memory("Local\\YubideckInput", 45, false) {
        Ok(s) => {
            unsafe { INPUT_SHMEM = Some(Arc::new(s)) };
            S_OK
        }
        Err(e) => {
            error!("Could not acquire shared memory: {e:#?}");
            E_FAIL
        }
    }
}

fn create_shared_memory<S>(os_id: S, size: usize, is_owner: bool) -> Result<Shmem>
where
    S: AsRef<str> + Display + Copy,
{
    let shmem_conf = ShmemConf::new().size(size).os_id(os_id);
    let shmem_result = match shmem_conf.clone().create() {
        Ok(s) => Ok(s),
        Err(ShmemError::MappingIdExists) => shmem_conf.open(),
        Err(e) => {
            return Err(anyhow!(
                "Failed to create/open shared memory {os_id}: {e:#}"
            ));
        }
    };

    shmem_result
        .map(|mut m| {
            m.set_owner(is_owner);
            m
        })
        .map_err(|e| anyhow!("Failed to create/open shared memory {os_id}: {e:#}"))
}

fn yubideck_init() -> Result<()> {
    let Some(mut device) = rusb::open_device_with_vid_pid(0x1973, 0x2001) else {
        return Err(anyhow!("YubiDeck not found."));
    };

    device.set_active_configuration(1)?;
    device.claim_interface(0)?;
    DEVICE
        .set(device)
        .map_err(|e| anyhow!("Cannot store device handle: {e:#?}"))?;

    thread::spawn(input_thread_proc);

    Ok(())
}

fn input_thread_proc() {
    info!("Input thread started");

    let mut shmem = match create_shared_memory("Local\\YubideckInput", 45, true) {
        Ok(s) => s,
        Err(e) => {
            error!("Could not obtain shared memory for YubiDeck input: {e:#}");
            return;
        }
    };
    let usb_in = unsafe { shmem.as_slice_mut() };
    let device = DEVICE.get().unwrap();

    loop {
        if let Err(e) = device.read_interrupt(0x81, usb_in, Duration::from_millis(20)) {
            error!("Failed to read data from YubiDeck: {e:#}");
        }
    }
}