use core::mem::MaybeUninit;

use alloc::{boxed::Box, sync::Arc};
use async_trait::async_trait;
use yggdrasil_abi::io::{KeyboardKey, KeyboardKeyEvent};

use crate::{device::UsbDeviceAccess, error::UsbError, info::UsbDeviceClass};

use super::{UsbClassInfo, UsbDriver};

pub struct UsbHidKeyboardDriver;

const MODIFIER_MAP: &[KeyboardKey] = &[
    KeyboardKey::LControl,
    KeyboardKey::LShift,
    KeyboardKey::LAlt,
    KeyboardKey::Unknown,
    KeyboardKey::RControl,
    KeyboardKey::RShift,
    KeyboardKey::RAlt,
    KeyboardKey::Unknown,
];

#[derive(Default)]
struct KeyboardState {
    state: [u64; 4],
    mods: u8,
}

impl KeyboardState {
    pub fn new() -> Self {
        Self::default()
    }

    pub fn translate_key(k: u8) -> KeyboardKey {
        match k {
            4..=29 => KeyboardKey::Char(k - 4 + b'a'),
            30..=38 => KeyboardKey::Char(k - 30 + b'1'),
            39 => KeyboardKey::Char(b'0'),

            40 => KeyboardKey::Enter,
            41 => KeyboardKey::Escape,
            42 => KeyboardKey::Backspace,
            43 => KeyboardKey::Tab,
            44 => KeyboardKey::Char(b' '),
            45 => KeyboardKey::Char(b'-'),
            46 => KeyboardKey::Char(b'='),
            47 => KeyboardKey::Char(b'['),
            48 => KeyboardKey::Char(b']'),
            49 => KeyboardKey::Char(b'\\'),
            51 => KeyboardKey::Char(b';'),
            52 => KeyboardKey::Char(b'\''),
            53 => KeyboardKey::Char(b'`'),
            54 => KeyboardKey::Char(b','),
            55 => KeyboardKey::Char(b'.'),
            56 => KeyboardKey::Char(b'/'),

            58..=69 => KeyboardKey::F(k - 58),

            _ => {
                log::debug!("Unknown key: {}", k);
                KeyboardKey::Unknown
            }
        }
    }

    pub fn retain_modifiers(
        &mut self,
        m: u8,
        events: &mut [MaybeUninit<KeyboardKeyEvent>],
    ) -> usize {
        let mut count = 0;
        let released = self.mods & !m;
        for (i, modifier) in MODIFIER_MAP.iter().enumerate().take(8) {
            if released & (1 << i) != 0 {
                events[count].write(KeyboardKeyEvent::Released(*modifier));
                count += 1;
            }
        }
        self.mods &= m;
        count
    }

    pub fn press_modifiers(
        &mut self,
        m: u8,
        events: &mut [MaybeUninit<KeyboardKeyEvent>],
    ) -> usize {
        let mut count = 0;
        let pressed = m & !self.mods;
        for (i, modifier) in MODIFIER_MAP.iter().enumerate().take(8) {
            if pressed & (1 << i) != 0 {
                events[count].write(KeyboardKeyEvent::Pressed(*modifier));
                count += 1;
            }
        }
        self.mods = m;
        count
    }

    pub fn retain(&mut self, keys: &[u8], events: &mut [MaybeUninit<KeyboardKeyEvent>]) -> usize {
        let mut count = 0;
        for i in 1..256 {
            if self.state[i / 64] & (1 << (i % 64)) != 0 && !keys.contains(&(i as u8)) {
                events[count].write(KeyboardKeyEvent::Released(Self::translate_key(i as u8)));
                self.state[i / 64] &= !(1 << (i % 64));
                count += 1;
            }
        }
        count
    }

    pub fn press(&mut self, keys: &[u8], events: &mut [MaybeUninit<KeyboardKeyEvent>]) -> usize {
        let mut count = 0;
        for &k in keys {
            let index = (k as usize) / 64;
            if self.state[index] & (1 << (k % 64)) == 0 {
                self.state[index] |= 1 << (k % 64);
                events[count].write(KeyboardKeyEvent::Pressed(Self::translate_key(k)));
                count += 1;
            }
        }
        count
    }
}

#[async_trait]
impl UsbDriver for UsbHidKeyboardDriver {
    async fn run(self: Arc<Self>, device: Arc<UsbDeviceAccess>) -> Result<(), UsbError> {
        // TODO not sure whether to use boot protocol (easy) or GetReport
        let config = device.select_configuration(|_| true).await?.unwrap();

        log::info!("Setup HID keyboard");
        let pipe = device
            .open_interrupt_in_pipe(1, config.endpoints[0].max_packet_size as u16)
            .await?;

        let mut buffer = [0; 8];
        let mut state = KeyboardState::new();
        let mut events = [MaybeUninit::uninit(); 16];

        loop {
            let mut event_count = 0;

            let len = pipe.read(&mut buffer).await?;
            if len < 8 {
                continue;
            }

            event_count += state.retain_modifiers(buffer[0], &mut events);
            event_count += state.press_modifiers(buffer[0], &mut events[event_count..]);
            event_count += state.retain(&buffer[2..], &mut events[event_count..]);
            event_count += state.press(&buffer[2..], &mut events[event_count..]);

            let events = unsafe { events[..event_count].assume_init_ref() };

            for &event in events {
                log::trace!("Generic Keyboard: {:?}", event);
                ygg_driver_input::send_event(event);
            }
        }
    }

    fn name(&self) -> &'static str {
        "USB HID Keyboard"
    }

    fn probe(&self, class: &UsbClassInfo, _device: &UsbDeviceAccess) -> bool {
        log::info!(
            "class = {:?}, subclass = {:02x}",
            class.class,
            class.subclass
        );
        class.class == UsbDeviceClass::Hid && (class.subclass == 0x00 || class.subclass == 0x01)
    }
}
