/// PS/2 keyboard driver — polling, scancode set 1, US QWERTY layout.
// Port addresses
const DATA_PORT: u16 = 0x60;
const STATUS_PORT: u16 = 0x64;

// Status register bits
const STAT_OUTPUT_FULL: u8 = 0x01;

// Modifier bit constants
const MOD_LSHIFT: u8 = 0x01;
const MOD_RSHIFT: u8 = 0x02;
const MOD_LCTRL: u8 = 0x04;
const MOD_RCTRL: u8 = 0x08;
const MOD_LALT: u8 = 0x10;
const MOD_RALT: u8 = 0x20;
const MOD_CAPS: u8 = 0x40;

/// A decoded key event.
pub enum KeyEvent {
    Char(char),
    Enter,
    Backspace,
    Tab,
    Escape,
    ArrowUp,
    ArrowDown,
    ArrowLeft,
    ArrowRight,
    Home,
    End,
    PageUp,
    PageDown,
    Delete,
    F(u8),
    Unknown(u8),
}

// ----- scancode-set-1 translation tables (US QWERTY) -----

/// Normal (unshifted) scancode → ASCII mapping.  0 = unmapped.
static SCAN_NORMAL: [u8; 128] = {
    let mut t = [0u8; 128];
    t[0x02] = b'1';
    t[0x03] = b'2';
    t[0x04] = b'3';
    t[0x05] = b'4';
    t[0x06] = b'5';
    t[0x07] = b'6';
    t[0x08] = b'7';
    t[0x09] = b'8';
    t[0x0A] = b'9';
    t[0x0B] = b'0';
    t[0x0C] = b'-';
    t[0x0D] = b'=';
    t[0x10] = b'q';
    t[0x11] = b'w';
    t[0x12] = b'e';
    t[0x13] = b'r';
    t[0x14] = b't';
    t[0x15] = b'y';
    t[0x16] = b'u';
    t[0x17] = b'i';
    t[0x18] = b'o';
    t[0x19] = b'p';
    t[0x1A] = b'[';
    t[0x1B] = b']';
    t[0x1E] = b'a';
    t[0x1F] = b's';
    t[0x20] = b'd';
    t[0x21] = b'f';
    t[0x22] = b'g';
    t[0x23] = b'h';
    t[0x24] = b'j';
    t[0x25] = b'k';
    t[0x26] = b'l';
    t[0x27] = b';';
    t[0x28] = b'\'';
    t[0x29] = b'`';
    t[0x2B] = b'\\';
    t[0x2C] = b'z';
    t[0x2D] = b'x';
    t[0x2E] = b'c';
    t[0x2F] = b'v';
    t[0x30] = b'b';
    t[0x31] = b'n';
    t[0x32] = b'm';
    t[0x33] = b',';
    t[0x34] = b'.';
    t[0x35] = b'/';
    t[0x37] = b'*';
    t[0x39] = b' ';
    t
};

/// Shifted scancode → ASCII.
static SCAN_SHIFT: [u8; 128] = {
    let mut t = [0u8; 128];
    t[0x02] = b'!';
    t[0x03] = b'@';
    t[0x04] = b'#';
    t[0x05] = b'$';
    t[0x06] = b'%';
    t[0x07] = b'^';
    t[0x08] = b'&';
    t[0x09] = b'*';
    t[0x0A] = b'(';
    t[0x0B] = b')';
    t[0x0C] = b'_';
    t[0x0D] = b'+';
    t[0x10] = b'Q';
    t[0x11] = b'W';
    t[0x12] = b'E';
    t[0x13] = b'R';
    t[0x14] = b'T';
    t[0x15] = b'Y';
    t[0x16] = b'U';
    t[0x17] = b'I';
    t[0x18] = b'O';
    t[0x19] = b'P';
    t[0x1A] = b'{';
    t[0x1B] = b'}';
    t[0x1E] = b'A';
    t[0x1F] = b'S';
    t[0x20] = b'D';
    t[0x21] = b'F';
    t[0x22] = b'G';
    t[0x23] = b'H';
    t[0x24] = b'J';
    t[0x25] = b'K';
    t[0x26] = b'L';
    t[0x27] = b':';
    t[0x28] = b'"';
    t[0x29] = b'~';
    t[0x2B] = b'|';
    t[0x2C] = b'Z';
    t[0x2D] = b'X';
    t[0x2E] = b'C';
    t[0x2F] = b'V';
    t[0x30] = b'B';
    t[0x31] = b'N';
    t[0x32] = b'M';
    t[0x33] = b'<';
    t[0x34] = b'>';
    t[0x35] = b'?';
    t[0x37] = b'*';
    t[0x39] = b' ';
    t
};

/// Internal keyboard state
struct KbdState {
    modifiers: u8,
    got_e0: bool,
}

static mut KBD: KbdState = KbdState {
    modifiers: 0,
    got_e0: false,
};

/// Initialise the PS/2 port: drain stale bytes + enable keyboard interrupt.
pub fn init() {
    drain_input();

    // Enable the keyboard interrupt on the PS/2 controller.
    // Read Configuration Byte (command 0x20), set bit 0 (IRQ1 enable),
    // write back (command 0x60).
    unsafe {
        // Tell the keyboard controller we want to read the config byte.
        // Need to wait for the controller to be ready first.
        wait_write();
        core::arch::asm!("out dx, al", in("dx") 0x64u16, in("al") 0x20u8,
            options(nostack, preserves_flags));
        wait_read();
        let mut cfg: u8;
        core::arch::asm!("in al, dx", out("al") cfg, in("dx") 0x60u16,
            options(nostack, preserves_flags));

        cfg |= 0x01; // enable keyboard interrupt

        // Write config byte back
        wait_write();
        core::arch::asm!("out dx, al", in("dx") 0x64u16, in("al") 0x60u8,
            options(nostack, preserves_flags));
        wait_write();
        core::arch::asm!("out dx, al", in("dx") 0x60u16, in("al") cfg,
            options(nostack, preserves_flags));
    }
}

/// Wait until the PS/2 controller is ready to receive a command.
fn wait_write() {
    unsafe {
        loop {
            let status: u8;
            core::arch::asm!("in al, dx", out("al") status, in("dx") 0x64u16,
                options(nostack, preserves_flags));
            if status & 0x02 == 0 {
                // Input buffer empty — controller ready to receive
                break;
            }
        }
    }
}

/// Wait until the PS/2 controller has data available to read.
fn wait_read() {
    unsafe {
        loop {
            let status: u8;
            core::arch::asm!("in al, dx", out("al") status, in("dx") 0x64u16,
                options(nostack, preserves_flags));
            if status & 0x01 != 0 {
                // Output buffer full — data ready to read
                break;
            }
        }
    }
}

fn drain_input() {
    unsafe {
        loop {
            let status: u8;
            core::arch::asm!("in al, dx", out("al") status, in("dx") STATUS_PORT,
                options(nostack, preserves_flags));
            if status & STAT_OUTPUT_FULL == 0 {
                break;
            }
            let _data: u8;
            core::arch::asm!("in al, dx", out("al") _data, in("dx") DATA_PORT,
                options(nostack, preserves_flags));
        }
    }
}

/// Poll for a key event. Blocks via `hlt` until a key is pressed.
/// The keyboard interrupt handler pushes raw scancodes to a ring buffer;
/// this function pops and decodes them.
pub fn poll_key() -> KeyEvent {
    loop {
        if let Some(byte) = crate::interrupts::read_scancode() {
            if let Some(ev) = process_byte(byte) {
                return ev;
            }
        } else {
            unsafe {
                core::arch::asm!("sti; hlt", options(nostack, preserves_flags));
            }
        }
    }
}

/// Non-blocking poll — read from the interrupt ring buffer if available.
pub fn try_poll_key() -> Option<KeyEvent> {
    let byte = crate::interrupts::read_scancode()?;
    process_byte(byte)
}

/// Process a single scancode byte. Returns Some(KeyEvent) when a visible
/// event happens; returns None for modifier changes, releases, and prefixes.
fn process_byte(byte: u8) -> Option<KeyEvent> {
    let state = unsafe { &mut KBD };

    if byte == 0xE0 {
        state.got_e0 = true;
        return None;
    }

    if state.got_e0 {
        state.got_e0 = false;
        return process_e0(byte);
    }

    // Break code (key release) — bit 7 set
    if byte & 0x80 != 0 {
        update_modifiers(byte & 0x7F, false, false);
        return None;
    }

    // Make code (key press)
    update_modifiers(byte, true, false);

    match byte {
        0x01 => Some(KeyEvent::Escape),
        0x0E => Some(KeyEvent::Backspace),
        0x0F => Some(KeyEvent::Tab),
        0x1C => Some(KeyEvent::Enter),
        0x37 => Some(KeyEvent::Char('*')),
        0x3B..=0x44 => Some(KeyEvent::F(byte - 0x3B + 1)),
        _ => {
            let shifted = (state.modifiers & (MOD_LSHIFT | MOD_RSHIFT)) != 0;
            // CapsLock inverts case for letters
            let caps = state.modifiers & MOD_CAPS != 0;
            let table = if shifted { SCAN_SHIFT } else { SCAN_NORMAL };
            let mut ch = table[byte as usize];
            // Handle CapsLock for letters
            if caps && ch.is_ascii_alphabetic() {
                // Swap case: if table gives lowercase, make uppercase; etc.
                // Since normal table has lowercase, shift table has uppercase:
                ch = if shifted {
                    ch.to_ascii_lowercase()
                } else {
                    ch.to_ascii_uppercase()
                };
            }
            if ch != 0 {
                Some(KeyEvent::Char(ch as char))
            } else {
                Some(KeyEvent::Unknown(byte))
            }
        }
    }
}

/// Process an E0-prefixed scancode.
fn process_e0(byte: u8) -> Option<KeyEvent> {
    if byte & 0x80 != 0 {
        update_modifiers(byte & 0x7F, false, true);
        return None;
    }

    update_modifiers(byte, true, true);

    match byte {
        0x48 => Some(KeyEvent::ArrowUp),
        0x50 => Some(KeyEvent::ArrowDown),
        0x4B => Some(KeyEvent::ArrowLeft),
        0x4D => Some(KeyEvent::ArrowRight),
        0x47 => Some(KeyEvent::Home),
        0x4F => Some(KeyEvent::End),
        0x49 => Some(KeyEvent::PageUp),
        0x51 => Some(KeyEvent::PageDown),
        0x53 => Some(KeyEvent::Delete),
        0x35 => Some(KeyEvent::Char('/')),
        _ => Some(KeyEvent::Unknown(byte)),
    }
}

fn update_modifiers(scancode: u8, pressed: bool, e0: bool) {
    let state = unsafe { &mut KBD };
    match (scancode, e0) {
        (0x2A, false) => set_mod(state, MOD_LSHIFT, pressed),
        (0x36, false) => set_mod(state, MOD_RSHIFT, pressed),
        (0x1D, false) => set_mod(state, MOD_LCTRL, pressed),
        (0x1D, true) => set_mod(state, MOD_RCTRL, pressed),
        (0x38, false) => set_mod(state, MOD_LALT, pressed),
        (0x38, true) => set_mod(state, MOD_RALT, pressed),
        (0x3A, false) if pressed => {
            state.modifiers ^= MOD_CAPS;
        }
        _ => {}
    }
}

fn set_mod(state: &mut KbdState, bit: u8, pressed: bool) {
    if pressed {
        state.modifiers |= bit;
    } else {
        state.modifiers &= !bit;
    }
}