cge_nes 0.1.2

//! Input register management module
//!
//! Handles the NES input registers ($4016 and $4017) which are used to read
//! controller input states. Implements the serial reading protocol used by the NES
//! to communicate with controllers.

use super::gamepad::GamepadState;
use crate::input::host_input::HostInput;
use crate::ConnectedSocket;

/// Represents the state of an input device connected to a controller port
#[derive(Copy, Clone, PartialEq, Eq, Debug)]
pub enum InputDeviceState {
    /// A gamepad is connected with its current button states
    Gamepad(GamepadState),
    /// No device is connected to this port
    Disconnected,
}

impl Default for InputDeviceState {
    fn default() -> Self {
        Self::Disconnected
    }
}

/// Manages the state of a single input register
#[derive(Default)]
struct InputRegister {
    /// Cached state of the connected input device
    cached_state: InputDeviceState,
    /// Current shift register value for serial reading
    gamepad_bits_shift: u8,
    /// Which controller port this register represents
    connected_socket: ConnectedSocket,
    /// Whether strobe mode is currently active
    strobe_mode_on: bool,
}

/// Manages both input registers ($4016 and $4017)
pub struct InputRegisters {
    /// First player's input register ($4016)
    player1: InputRegister,
    /// Second player's input register ($4017)
    player2: InputRegister,
}

impl Default for InputRegisters {
    fn default() -> Self {
        Self::new()
    }
}

impl InputRegisters {
    /// Creates a new instance of InputRegisters
    pub fn new() -> Self {
        Self {
            player1: InputRegister::new(ConnectedSocket::Player1),
            player2: InputRegister::new(ConnectedSocket::Player2),
        }
    }

    /// Enables strobe mode for both input registers
    ///
    /// In strobe mode, the registers continuously read the current button states
    pub fn set_strobe_mode_on(&mut self) {
        self.player1.set_strobe_mode_on();
        self.player2.set_strobe_mode_on();
    }

    /// Disables strobe mode and latches the current input states
    ///
    /// # Arguments
    /// * `host_input` - Interface to query the current input states from the host system
    ///
    /// When strobe mode is disabled, the registers enter serial reading mode
    /// where button states are read one bit at a time
    pub fn set_strobe_mode_off(&mut self, host_input: &mut impl HostInput) {
        self.player1.set_strobe_mode_off(host_input);
        self.player2.set_strobe_mode_off(host_input);
    }

    /// Reads the next bit from the specified input register
    ///
    /// # Arguments
    /// * `host_input` - Interface to query input state if needed (when in strobe mode)
    /// * `player` - Which player's register to read from (Player1/Player2)
    ///
    /// # Returns
    /// The next bit from the register's shift register (0 or 1)
    pub fn read_register(
        &mut self,
        host_input: &mut impl HostInput,
        player: ConnectedSocket,
    ) -> u8 {
        match player {
            ConnectedSocket::Player1 => self.player1.read_register(host_input),
            ConnectedSocket::Player2 => self.player2.read_register(host_input),
        }
    }
}

impl InputRegister {
    /// Creates a new input register for the specified controller port
    ///
    /// # Arguments
    /// * `connected_socket` - Specifies which controller port (Player1/Player2) this register manages
    fn new(connected_socket: ConnectedSocket) -> Self {
        Self {
            connected_socket,
            ..Default::default()
        }
    }

    /// Enables strobe mode for this register
    fn set_strobe_mode_on(&mut self) {
        self.strobe_mode_on = true;
    }

    /// Disables strobe mode and latches the current input state
    ///
    /// # Arguments
    /// * `host_input` - Interface to query the current input state from the host system
    fn set_strobe_mode_off(&mut self, host_input: &mut impl HostInput) {
        if self.strobe_mode_on == true {
            self.strobe_mode_on = false;

            self.cached_state = host_input.input_device_state(self.connected_socket);

            if let InputDeviceState::Gamepad(state) = self.cached_state {
                self.gamepad_bits_shift = state.buttons_pressed.bits();
            }
        }
    }

    /// Reads the next bit from this register's shift register
    ///
    /// # Arguments
    /// * `host_input` - Interface to query input state if needed (when in strobe mode)
    fn read_register(&mut self, host_input: &mut impl HostInput) -> u8 {
        if self.strobe_mode_on == true {
            self.cached_state = host_input.input_device_state(self.connected_socket);

            if let InputDeviceState::Gamepad(state) = self.cached_state {
                self.gamepad_bits_shift = state.buttons_pressed.bits();
            }
        }

        let result;
        match self.cached_state {
            InputDeviceState::Gamepad(_) => {
                result = self.gamepad_bits_shift & 1;
            }
            InputDeviceState::Disconnected => result = 0,
        }

        if self.strobe_mode_on == false {
            self.gamepad_bits_shift = (self.gamepad_bits_shift >> 1) | (1 << 7);
        }
        result
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::input::gamepad::GamepadButtonsPressedFlags;

    struct MockHostInput {
        inner: InputDeviceState,
    }

    impl HostInput for MockHostInput {
        fn input_device_state(&mut self, _player: ConnectedSocket) -> InputDeviceState {
            self.inner
        }
    }

    #[test]
    fn inputreg_querystatestrobeon_alwaysreturnfirstbutton() {
        let mut input = InputRegister::new(ConnectedSocket::Player1);

        let state_flags = GamepadButtonsPressedFlags::A;
        let device_state = InputDeviceState::Gamepad(GamepadState::new(state_flags));
        let mut host_input = MockHostInput {
            inner: device_state,
        };
        input.set_strobe_mode_on();

        for _ in 0..16 {
            assert_eq!(input.read_register(&mut host_input), 1);
        }

        let device_state = InputDeviceState::Gamepad(GamepadState::new(state_flags.complement()));
        let mut host_input = MockHostInput {
            inner: device_state,
        };

        for _ in 0..16 {
            assert_eq!(input.read_register(&mut host_input), 0);
        }
    }

    #[test]
    fn inputreg_querystate_statechangesafterstrobe() {
        let mut input = InputRegister::new(ConnectedSocket::Player1);

        let device_state =
            InputDeviceState::Gamepad(GamepadState::new(GamepadButtonsPressedFlags::A));
        let mut host_input = MockHostInput {
            inner: device_state,
        };

        assert_eq!(input.read_register(&mut host_input), 0);

        input.set_strobe_mode_on();

        assert_eq!(input.read_register(&mut host_input), 1);
    }

    #[test]
    fn inputreg_querystate_buttonsshiftcorrect() {
        let buttons = [
            GamepadButtonsPressedFlags::A,
            GamepadButtonsPressedFlags::B,
            GamepadButtonsPressedFlags::SELECT,
            GamepadButtonsPressedFlags::START,
            GamepadButtonsPressedFlags::UP,
            GamepadButtonsPressedFlags::DOWN,
            GamepadButtonsPressedFlags::LEFT,
            GamepadButtonsPressedFlags::RIGHT,
        ];

        let mut input = InputRegister::new(ConnectedSocket::Player1);

        for state_bits in 0..u8::MAX {
            let state_flags = GamepadButtonsPressedFlags::from_bits(state_bits)
                .expect("Invalid bits for gamepad flags (all combinations should be valid!!)");

            let device_state = InputDeviceState::Gamepad(GamepadState::new(state_flags));
            let mut host_input = MockHostInput {
                inner: device_state,
            };
            input.set_strobe_mode_on();
            input.set_strobe_mode_off(&mut host_input);

            for &button_flag in buttons.iter() {
                if input.read_register(&mut host_input) == 0 {
                    assert!(!state_flags.contains(button_flag));
                } else {
                    assert!(state_flags.contains(button_flag));
                }
            }

            for _ in 0..16 {
                assert_eq!(input.read_register(&mut host_input), 1)
            }
        }
    }
}