//! PS/2 scancode set 2 → USB HID usage decoding, plus the keyboard state a driver
//! needs on top of it (pressed keys, modifier tracking, caps-lock toggle).
//!
//! Set 2 is what a keyboard sends when the 8042's legacy set-1 translation is off —
//! which is how ps2-bus.zig deliberately configures the controller. A key's **make**
//! code is one byte (two with an `E0` prefix for the "extended" keys added after the
//! original AT layout); its **break** code is the same code behind an `F0` prefix.
//! Pause alone is an eight-byte `E1` sequence with no break.
//!
//! The output vocabulary is USB HID keyboard-page usages (a=4, enter=40, ...), the
//! same numbering the input protocol's `Keycode` and the xkeyboard-config layout
//! tables use — so a decoded usage indexes a layout directly.
//!
//! Everything here is pure (no imports beyond `std`, no IO), so it is host-testable:
//! the tests at the bottom run under `zig build test`.

const std = @import("std");

// --- USB HID usages the state machine itself needs to recognize --------------

pub const usage_caps_lock: u8 = 0x39;
pub const usage_left_control: u8 = 0xE0;
pub const usage_left_shift: u8 = 0xE1;
pub const usage_left_alt: u8 = 0xE2;
pub const usage_right_control: u8 = 0xE4;
pub const usage_right_shift: u8 = 0xE5;
pub const usage_right_alt: u8 = 0xE6; // AltGr — selects XKB level 3

// --- scancode set 2 → HID usage tables ---------------------------------------

/// Single-byte (non-`E0`) make codes. Zero means "no key" — protocol bytes (ACK,
/// BAT results) and reserved codes land there and decode to nothing.
pub const set2_base: [256]u8 = blk: {
    var table = [_]u8{0} ** 256;
    // function row
    table[0x01] = 0x42; // F9
    table[0x03] = 0x3E; // F5
    table[0x04] = 0x3C; // F3
    table[0x05] = 0x3A; // F1
    table[0x06] = 0x3B; // F2
    table[0x07] = 0x45; // F12
    table[0x09] = 0x43; // F10
    table[0x0A] = 0x41; // F8
    table[0x0B] = 0x3F; // F6
    table[0x0C] = 0x3D; // F4
    table[0x78] = 0x44; // F11
    table[0x83] = 0x40; // F7
    // letters
    table[0x1C] = 0x04; // A
    table[0x32] = 0x05; // B
    table[0x21] = 0x06; // C
    table[0x23] = 0x07; // D
    table[0x24] = 0x08; // E
    table[0x2B] = 0x09; // F
    table[0x34] = 0x0A; // G
    table[0x33] = 0x0B; // H
    table[0x43] = 0x0C; // I
    table[0x3B] = 0x0D; // J
    table[0x42] = 0x0E; // K
    table[0x4B] = 0x0F; // L
    table[0x3A] = 0x10; // M
    table[0x31] = 0x11; // N
    table[0x44] = 0x12; // O
    table[0x4D] = 0x13; // P
    table[0x15] = 0x14; // Q
    table[0x2D] = 0x15; // R
    table[0x1B] = 0x16; // S
    table[0x2C] = 0x17; // T
    table[0x3C] = 0x18; // U
    table[0x2A] = 0x19; // V
    table[0x1D] = 0x1A; // W
    table[0x22] = 0x1B; // X
    table[0x35] = 0x1C; // Y
    table[0x1A] = 0x1D; // Z
    // digit row
    table[0x16] = 0x1E; // 1
    table[0x1E] = 0x1F; // 2
    table[0x26] = 0x20; // 3
    table[0x25] = 0x21; // 4
    table[0x2E] = 0x22; // 5
    table[0x36] = 0x23; // 6
    table[0x3D] = 0x24; // 7
    table[0x3E] = 0x25; // 8
    table[0x46] = 0x26; // 9
    table[0x45] = 0x27; // 0
    // control and whitespace
    table[0x5A] = 0x28; // Enter
    table[0x76] = 0x29; // Escape
    table[0x66] = 0x2A; // Backspace
    table[0x0D] = 0x2B; // Tab
    table[0x29] = 0x2C; // Space
    // punctuation
    table[0x4E] = 0x2D; // - _
    table[0x55] = 0x2E; // = +
    table[0x54] = 0x2F; // [ {
    table[0x5B] = 0x30; // ] }
    table[0x5D] = 0x31; // \ | (non-US hash on ISO boards, same position)
    table[0x4C] = 0x33; // ; :
    table[0x52] = 0x34; // ' "
    table[0x0E] = 0x35; // ` ~
    table[0x41] = 0x36; // , <
    table[0x49] = 0x37; // . >
    table[0x4A] = 0x38; // / ?
    table[0x61] = 0x64; // non-US backslash (the extra ISO key between shift and Z)
    // locks
    table[0x58] = usage_caps_lock;
    table[0x77] = 0x53; // Num Lock
    table[0x7E] = 0x47; // Scroll Lock
    // keypad
    table[0x7C] = 0x55; // keypad *
    table[0x7B] = 0x56; // keypad -
    table[0x79] = 0x57; // keypad +
    table[0x69] = 0x59; // keypad 1
    table[0x72] = 0x5A; // keypad 2
    table[0x7A] = 0x5B; // keypad 3
    table[0x6B] = 0x5C; // keypad 4
    table[0x73] = 0x5D; // keypad 5
    table[0x74] = 0x5E; // keypad 6
    table[0x6C] = 0x5F; // keypad 7
    table[0x75] = 0x60; // keypad 8
    table[0x7D] = 0x61; // keypad 9
    table[0x70] = 0x62; // keypad 0
    table[0x71] = 0x63; // keypad .
    // modifiers
    table[0x14] = usage_left_control;
    table[0x12] = usage_left_shift;
    table[0x11] = usage_left_alt;
    table[0x59] = usage_right_shift;
    break :blk table;
};

/// `E0`-prefixed make codes. `E0 12` is the "fake shift" the keyboard wraps around
/// Print Screen and navigation keys when a real shift is involved; it maps to zero
/// here, so it decodes to nothing and only the real key comes through.
pub const set2_extended: [256]u8 = blk: {
    var table = [_]u8{0} ** 256;
    table[0x11] = usage_right_alt;
    table[0x14] = usage_right_control;
    table[0x1F] = 0xE3; // left GUI
    table[0x27] = 0xE7; // right GUI
    table[0x2F] = 0x65; // application (menu)
    table[0x7C] = 0x46; // Print Screen (arrives as E0 12 E0 7C; the E0 12 decodes to nothing)
    table[0x4A] = 0x54; // keypad /
    table[0x5A] = 0x58; // keypad Enter
    table[0x70] = 0x49; // Insert
    table[0x6C] = 0x4A; // Home
    table[0x7D] = 0x4B; // Page Up
    table[0x71] = 0x4C; // Delete
    table[0x69] = 0x4D; // End
    table[0x7A] = 0x4E; // Page Down
    table[0x74] = 0x4F; // right arrow
    table[0x6B] = 0x50; // left arrow
    table[0x72] = 0x51; // down arrow
    table[0x75] = 0x52; // up arrow
    break :blk table;
};

// --- the byte-stream decoder --------------------------------------------------

/// One decoded key transition: which key (as a USB HID usage) and whether this is
/// a make (press or typematic repeat) or a break (release).
pub const DecodedKey = struct {
    usage: u8,
    make: bool,
};

/// Turns the raw set-2 byte stream into `DecodedKey`s. Feed it every byte the
/// keyboard sends; most bytes complete a key and return one, prefix bytes return
/// null and arm the state machine for the next byte.
pub const Decoder = struct {
    const State = enum {
        idle,
        extended, // saw E0
        break_prefix, // saw F0
        extended_break, // saw E0 F0
        pause_skip, // inside the 8-byte E1 Pause sequence
    };

    state: State = .idle,
    /// Bytes still to swallow in `pause_skip`.
    skip: u8 = 0,

    /// The whole Pause make sequence is `E1 14 77 E1 F0 14 F0 77` — seven bytes
    /// after the leading `E1`, and no break sequence ever follows.
    const pause_bytes_after_e1: u8 = 7;

    pub fn feed(self: *Decoder, byte: u8) ?DecodedKey {
        switch (self.state) {
            .idle => switch (byte) {
                0xE0 => self.state = .extended,
                0xF0 => self.state = .break_prefix,
                0xE1 => {
                    self.state = .pause_skip;
                    self.skip = pause_bytes_after_e1;
                },
                // Anything else is a make code — or a protocol byte (0xFA ACK,
                // 0xAA BAT-passed, 0xEE echo, ...), which the tables map to zero.
                else => return decoded(set2_base[byte], true),
            },
            .extended => switch (byte) {
                0xF0 => self.state = .extended_break,
                else => {
                    self.state = .idle;
                    return decoded(set2_extended[byte], true);
                },
            },
            .break_prefix => {
                self.state = .idle;
                return decoded(set2_base[byte], false);
            },
            .extended_break => {
                self.state = .idle;
                return decoded(set2_extended[byte], false);
            },
            .pause_skip => {
                self.skip -= 1;
                if (self.skip == 0) self.state = .idle;
            },
        }
        return null;
    }

    fn decoded(usage: u8, make: bool) ?DecodedKey {
        if (usage == 0) return null; // unmapped or a protocol byte
        return .{ .usage = usage, .make = make };
    }
};

// --- driver-side keyboard state -----------------------------------------------

/// What a key transition did, plus the modifier state to stamp on the resulting
/// events (snapshotted after the transition was applied).
pub const Transition = struct {
    pub const Action = enum {
        pressed, // physical make of a key that was up
        repeated, // typematic make of a key already down — no new key_down
        released, // physical break
    };
    action: Action,
    modifiers: ModifierSnapshot,
};

/// The modifier state at one instant, in both vocabularies a driver needs: the
/// input protocol's coarse bits (shift/control/alt) and the level-selection
/// inputs xkeyboard-config takes (shift, caps_lock, AltGr as level3).
pub const ModifierSnapshot = struct {
    shift: bool, // either shift held
    control: bool, // either control held
    alt: bool, // either alt held (including AltGr)
    right_alt: bool, // AltGr specifically — the XKB level-3 selector
    caps_lock: bool, // the toggle, not the key
};

/// Tracks which keys are physically down and the caps-lock toggle, and classifies
/// each decoded transition. Pure state — no IO — so repeat detection and modifier
/// snapshots are host-testable.
pub const KeyboardState = struct {
    /// One bit per HID usage: set while the key is physically down.
    pressed: [32]u8 = [_]u8{0} ** 32,
    caps_lock: bool = false,

    pub fn apply(self: *KeyboardState, key: DecodedKey) Transition {
        const already_down = self.isPressed(key.usage);
        if (key.make) {
            if (!already_down) {
                self.setPressed(key.usage, true);
                if (key.usage == usage_caps_lock) self.caps_lock = !self.caps_lock;
            }
            return .{
                .action = if (already_down) .repeated else .pressed,
                .modifiers = self.snapshot(),
            };
        }
        self.setPressed(key.usage, false);
        return .{ .action = .released, .modifiers = self.snapshot() };
    }

    pub fn isPressed(self: *const KeyboardState, usage: u8) bool {
        return self.pressed[usage / 8] & (@as(u8, 1) << @intCast(usage % 8)) != 0;
    }

    fn setPressed(self: *KeyboardState, usage: u8, down: bool) void {
        const bit = @as(u8, 1) << @intCast(usage % 8);
        if (down) {
            self.pressed[usage / 8] |= bit;
        } else {
            self.pressed[usage / 8] &= ~bit;
        }
    }

    fn snapshot(self: *const KeyboardState) ModifierSnapshot {
        const right_alt = self.isPressed(usage_right_alt);
        return .{
            .shift = self.isPressed(usage_left_shift) or self.isPressed(usage_right_shift),
            .control = self.isPressed(usage_left_control) or self.isPressed(usage_right_control),
            .alt = self.isPressed(usage_left_alt) or right_alt,
            .right_alt = right_alt,
            .caps_lock = self.caps_lock,
        };
    }
};

// --- tests (host-run via `zig build test`) ------------------------------------

const testing = std.testing;

/// Feed `bytes` and return the single DecodedKey they should produce (fails the
/// test if they produce none or more than one).
fn feedOne(decoder: *Decoder, bytes: []const u8) !DecodedKey {
    var result: ?DecodedKey = null;
    for (bytes) |byte| {
        if (decoder.feed(byte)) |key| {
            try testing.expect(result == null);
            result = key;
        }
    }
    return result orelse error.TestExpectedResult;
}

fn feedNone(decoder: *Decoder, bytes: []const u8) !void {
    for (bytes) |byte| try testing.expectEqual(@as(?DecodedKey, null), decoder.feed(byte));
}

test "base make and break: A" {
    var decoder = Decoder{};
    try testing.expectEqual(DecodedKey{ .usage = 0x04, .make = true }, try feedOne(&decoder, &.{0x1C}));
    try testing.expectEqual(DecodedKey{ .usage = 0x04, .make = false }, try feedOne(&decoder, &.{ 0xF0, 0x1C }));
}

test "extended make and break: right arrow" {
    var decoder = Decoder{};
    try testing.expectEqual(DecodedKey{ .usage = 0x4F, .make = true }, try feedOne(&decoder, &.{ 0xE0, 0x74 }));
    try testing.expectEqual(DecodedKey{ .usage = 0x4F, .make = false }, try feedOne(&decoder, &.{ 0xE0, 0xF0, 0x74 }));
}

test "pause: the E1 sequence is consumed silently" {
    var decoder = Decoder{};
    try feedNone(&decoder, &.{ 0xE1, 0x14, 0x77, 0xE1, 0xF0, 0x14, 0xF0, 0x77 });
    // The decoder is back in idle: an ordinary key still decodes.
    try testing.expectEqual(DecodedKey{ .usage = 0x04, .make = true }, try feedOne(&decoder, &.{0x1C}));
}

test "protocol bytes decode to nothing" {
    var decoder = Decoder{};
    try feedNone(&decoder, &.{ 0xFA, 0xAA, 0xEE }); // ACK, BAT-passed, echo
}

test "print screen: the fake-shift E0 12 decodes to nothing" {
    var decoder = Decoder{};
    try feedNone(&decoder, &.{ 0xE0, 0x12 });
    try testing.expectEqual(DecodedKey{ .usage = 0x46, .make = true }, try feedOne(&decoder, &.{ 0xE0, 0x7C }));
}

test "typematic repeat is classified, not re-pressed" {
    var state = KeyboardState{};
    const a = DecodedKey{ .usage = 0x04, .make = true };
    try testing.expectEqual(Transition.Action.pressed, state.apply(a).action);
    try testing.expectEqual(Transition.Action.repeated, state.apply(a).action);
    try testing.expectEqual(Transition.Action.repeated, state.apply(a).action);
    try testing.expectEqual(Transition.Action.released, state.apply(.{ .usage = 0x04, .make = false }).action);
    try testing.expectEqual(Transition.Action.pressed, state.apply(a).action);
}

test "shift held shows in the snapshot of other keys" {
    var state = KeyboardState{};
    _ = state.apply(.{ .usage = usage_left_shift, .make = true });
    const transition = state.apply(.{ .usage = 0x04, .make = true });
    try testing.expect(transition.modifiers.shift);
    try testing.expect(!transition.modifiers.control);
    _ = state.apply(.{ .usage = usage_left_shift, .make = false });
    _ = state.apply(.{ .usage = 0x04, .make = false });
    try testing.expect(!state.apply(.{ .usage = 0x04, .make = true }).modifiers.shift);
}

test "right alt reports both alt and the level-3 selector" {
    var state = KeyboardState{};
    _ = state.apply(.{ .usage = usage_right_alt, .make = true });
    const transition = state.apply(.{ .usage = 0x04, .make = true });
    try testing.expect(transition.modifiers.alt);
    try testing.expect(transition.modifiers.right_alt);
}

test "caps lock toggles on make, not on repeat or break" {
    var state = KeyboardState{};
    try testing.expect(state.apply(.{ .usage = usage_caps_lock, .make = true }).modifiers.caps_lock);
    try testing.expect(state.apply(.{ .usage = usage_caps_lock, .make = true }).modifiers.caps_lock); // repeat
    try testing.expect(state.apply(.{ .usage = usage_caps_lock, .make = false }).modifiers.caps_lock);
    try testing.expect(!state.apply(.{ .usage = usage_caps_lock, .make = true }).modifiers.caps_lock); // second press: off
}