use esp_hal::{
    Async,
    gpio::{Input, InputConfig, Pull},
    i2c::master::{BusTimeout, I2c, I2cAddress, SoftwareTimeout},
    time::Rate,
};
use log::info;
use rmk::{
    debounce::{DebounceState, DebouncerTrait},
    embassy_futures::yield_now,
    event::{Event, KeyboardEvent},
    input_device::InputDevice,
    matrix::{KeyState, MatrixTrait},
};

use crate::config::{MATRIX_AREA, MATRIX_COLS, MATRIX_ROWS};

pub struct RaiiGuard<F: FnOnce()> {
    on_drop: Option<F>,
}

impl<F: FnOnce()> RaiiGuard<F> {
    pub fn new(on_drop: F) -> Self {
        Self {
            on_drop: Some(on_drop),
        }
    }
}

impl<F: FnOnce()> Drop for RaiiGuard<F> {
    fn drop(&mut self) {
        if let Some(on_drop) = self.on_drop.take() {
            (on_drop)();
        }
    }
}

/// IO Expander Matrix
pub struct IoeMatrix<D>
where
    D: DebouncerTrait<MATRIX_ROWS, MATRIX_COLS>,
{
    /// Goes low when the matrix was changed and should be read.
    matrix_interrupt_low: Input<'static>,
    i2c: I2c<'static, Async>,
    /// IO expander addresses.
    ioe_addresses: [I2cAddress; 2],
    debouncer: D,
    /// States of key switches as detected by IO expanders.
    input_states: [[bool; MATRIX_ROWS]; MATRIX_COLS],
    /// Debounced key states.
    key_states: [[KeyState; MATRIX_ROWS]; MATRIX_COLS],
    /// Index of the current key switch being processed: row * columns + column
    scan_pos: usize,
    /// Re-scan needed flag
    #[cfg(feature = "async_matrix")]
    rescan_needed: bool,
}

impl<D> IoeMatrix<D>
where
    D: DebouncerTrait<MATRIX_ROWS, MATRIX_COLS>,
{
    pub async fn new(
        mut matrix_interrupt_low: Input<'static>,
        mut i2c: I2c<'static, Async>,
        debouncer: D,
        ioe_addresses: [I2cAddress; 2],
    ) -> Self {
        matrix_interrupt_low.apply_config(&InputConfig::default().with_pull(Pull::Up));
        i2c.apply_config(
            &esp_hal::i2c::master::Config::default()
                .with_frequency(Rate::from_khz(400))
                .with_timeout(BusTimeout::Disabled)
                .with_software_timeout(SoftwareTimeout::None),
        )
        .unwrap();

        for addr in ioe_addresses {
            loop {
                let Err(error) = i2c
                    .write_async(addr, &[0b10010000, 0xFF, 0xFF, 0xFF, 0xFF, 0xFF])
                    .await
                else {
                    break;
                };
                info!("I2C timed out when writing the inversion mask to an IO expander: {error:?}");
            }
        }

        Self {
            matrix_interrupt_low,
            i2c,
            ioe_addresses,
            debouncer,
            input_states: Default::default(),
            key_states: Default::default(),
            scan_pos: 0,
            #[cfg(feature = "async_matrix")]
            rescan_needed: false,
        }
    }
}

impl<D> InputDevice for IoeMatrix<D>
where
    D: DebouncerTrait<MATRIX_ROWS, MATRIX_COLS>,
{
    async fn read_event(&mut self) -> Event {
        loop {
            if self.scan_pos == 0 {
                // TODO: Only execute if interrupt received.
                // Timer::after(Duration::from_millis(100)).await;
                yield_now().await;

                // Load data from IO expanders.
                let mut input_register_banks = [0_u8; 10];

                loop {
                    let Err(esp_hal::i2c::master::Error::Timeout) = self
                        .i2c
                        .write_read_async(
                            self.ioe_addresses[0],
                            &[0x80],
                            &mut input_register_banks[..5],
                        )
                        .await
                    else {
                        break;
                    };
                    info!("I2C timed out when reading the left IO expander.");
                }

                loop {
                    let Err(esp_hal::i2c::master::Error::Timeout) = self
                        .i2c
                        .write_read_async(
                            self.ioe_addresses[1],
                            &[0x80],
                            &mut input_register_banks[5..],
                        )
                        .await
                    else {
                        break;
                    };
                    info!("I2C timed out when reading the right IO expander.");
                }

                for (row_idx, (left, right)) in (0..).zip(core::iter::zip(
                    &input_register_banks[..5],
                    &input_register_banks[5..],
                )) {
                    for col_idx_half in 0..(MATRIX_COLS / 2) {
                        self.input_states[col_idx_half][row_idx] = left & (1 << col_idx_half) != 0;
                        self.input_states[col_idx_half + MATRIX_COLS / 2][row_idx] =
                            right & (1 << col_idx_half) != 0;
                    }
                }
                // info!("rows: {rows:x?}")
            }

            let (row_idx, col_idx) = (self.scan_pos / MATRIX_COLS, self.scan_pos % MATRIX_COLS);
            let on_iteration_end = RaiiGuard::new(|| {
                // Process the next key switch.
                self.scan_pos = (self.scan_pos + 1) % MATRIX_AREA;
            });
            let debounce_state = self.debouncer.detect_change_with_debounce(
                row_idx,
                col_idx,
                self.input_states[col_idx][row_idx],
                &self.key_states[col_idx][row_idx],
            );

            if let DebounceState::Debounced = debounce_state {
                self.key_states[col_idx][row_idx].toggle_pressed();
                #[cfg(feature = "async_matrix")]
                {
                    self.rescan_needed = true;
                }
                return Event::Key(KeyboardEvent::key(
                    row_idx as u8,
                    col_idx as u8,
                    self.key_states[col_idx][row_idx].pressed,
                ));
            }

            // // Process the next key switch.
            // self.scan_pos = (self.scan_pos + 1) % MATRIX_AREA;

            // If there's key still pressed, always refresh the self.scan_start
            #[cfg(feature = "async_matrix")]
            if self.key_states[col_idx][row_idx].pressed {
                self.rescan_needed = true;
            }

            #[cfg(feature = "async_matrix")]
            {
                if !self.rescan_needed {
                    self.wait_for_key().await;
                }
                self.rescan_needed = false;
            }

            drop(on_iteration_end);
        }
    }
}

impl<D> MatrixTrait<MATRIX_ROWS, MATRIX_COLS> for IoeMatrix<D>
where
    D: DebouncerTrait<MATRIX_ROWS, MATRIX_COLS>,
{
    #[cfg(feature = "async_matrix")]
    async fn wait_for_key(&mut self) {
        // TODO

        // First, set all output pins to high
        for out in self.get_output_pins_mut().iter_mut() {
            out.set_high().ok();
        }

        // Wait for any key press
        self.wait_input_pins().await;

        // Set all output pins back to low
        for out in self.get_output_pins_mut().iter_mut() {
            out.set_low().ok();
        }
    }
}