#include "arch/Keyboard.h"
#include <luna/CType.h>

// PS/2 keyboard decoding routine.

static bool is_key_released(u8& scancode)
{
    if (scancode > 0x80)
    {
        scancode -= 0x80;
        return true;
    }
    return false;
}

constexpr u8 EXTENDED_KEY_CODE = 0xe0;
constexpr u8 LEFT_SHIFT = 0x2a;
constexpr u8 RIGHT_SHIFT = 0x36;
constexpr u8 CAPS_LOCK = 0x3a;

constexpr u8 LEFT_CONTROL = 0x1D;
constexpr u8 TAB = 0x0F;
constexpr u8 LEFT_ALT = 0x38;
constexpr u8 F11 = 0x57;
constexpr u8 F12 = 0x58;

static bool should_ignore_key(u8 scancode)
{
    return (scancode > 0x3A && scancode < 0x47) || scancode == TAB || scancode == F11 || scancode == F12;
}

constexpr char key_table[] = {
    '\0',
    '\1', // escape
    '1',  '2', '3', '4', '5', '6', '7', '8', '9', '0', '-',  '=', '\b',
    '\0', // tab
    'q',  'w', 'e', 'r', 't', 'y', 'u', 'i', 'o', 'p', '[',  ']', '\n',
    '\0', // left ctrl
    'a',  's', 'd', 'f', 'g', 'h', 'j', 'k', 'l', ';', '\'', '`',
    '\0', // left shift
    '\\', 'z', 'x', 'c', 'v', 'b', 'n', 'm', ',', '.', '/',
    '\0', // right shift
    '*',  // keypad *
    '\0', // left alt
    ' ',
    '\0', // caps lock
    '\0', // f1
    '\0', // f2
    '\0', // f3
    '\0', // f4
    '\0', // f5
    '\0', // f6
    '\0', // f7
    '\0', // f8
    '\0', // f9
    '\0', // f10
    '\0', // num lock
    '\0', // scroll lock
    '7',  // keypad 7
    '8',  // keypad 8
    '9',  // keypad 9
    '-',  // keypad -
    '4',  // keypad 4
    '5',  // keypad 5
    '6',  // keypad 6
    '+',  // keypad +
    '1',  // keypad 1
    '2',  // keypad 2
    '3',  // keypad 3
    '0',  // keypad 0
    '.',  // keypad .
};

constexpr char shifted_key_table[] = {
    '\0',
    '\1', // escape
    '!',  '@', '#', '$', '%', '^', '&', '*', '(', ')', '_', '+', '\b',
    '\0', // tab
    'Q',  'W', 'E', 'R', 'T', 'Y', 'U', 'I', 'O', 'P', '{', '}', '\n',
    '\0', // left ctrl
    'A',  'S', 'D', 'F', 'G', 'H', 'J', 'K', 'L', ':', '"', '~',
    '\0', // left shift
    '|',  'Z', 'X', 'C', 'V', 'B', 'N', 'M', '<', '>', '?',
    '\0', // right shift
    '*',  // keypad *
    '\0', // left alt
    ' ',
    '\0', // caps lock
    '\0', // f1
    '\0', // f2
    '\0', // f3
    '\0', // f4
    '\0', // f5
    '\0', // f6
    '\0', // f7
    '\0', // f8
    '\0', // f9
    '\0', // f10
    '\0', // num lock
    '\0', // scroll lock
    '7',  // keypad 7
    '8',  // keypad 8
    '9',  // keypad 9
    '-',  // keypad -
    '4',  // keypad 4
    '5',  // keypad 5
    '6',  // keypad 6
    '+',  // keypad +
    '1',  // keypad 1
    '2',  // keypad 2
    '3',  // keypad 3
    '0',  // keypad 0
    '.',  // keypad .
};

using namespace moon;

constexpr KeyCode keycode_table[] = {
    K_Unknown,    K_Esc,
    K_CH01, // 1
    K_CH02, // 2
    K_CH03, // 3
    K_CH04, // 4
    K_CH05, // 5
    K_CH06, // 6
    K_CH07, // 7
    K_CH08, // 8
    K_CH09, // 9
    K_CH10, // 0
    K_CH11, // -
    K_CH12, // =
    K_Backspace,  K_Tab,
    K_CH13, // Q
    K_CH14, // W
    K_CH15, // E
    K_CH16, // R
    K_CH17, // T
    K_CH18, // Y
    K_CH19, // U
    K_CH20, // I
    K_CH21, // O
    K_CH22, // P
    K_CH23, // [
    K_CH24, // ]
    K_Enter,      K_LeftControl,
    K_CH25, // A
    K_CH26, // S
    K_CH27, // D
    K_CH28, // F
    K_CH29, // G
    K_CH30, // H
    K_CH31, // J
    K_CH32, // K
    K_CH33, // L
    K_CH34, // ;
    K_CH35, // '
    K_CH00, // `
    K_LeftShift,
    K_CH36, // #
    K_CH38, // Z
    K_CH39, // X
    K_CH40, // C
    K_CH41, // V
    K_CH42, // B
    K_CH43, // N
    K_CH44, // M
    K_CH45, // ,
    K_CH46, // .
    K_CH47, // /
    K_RightShift, K_KeypadMul,   K_LeftAlt,
    K_CH48, // Space
    K_CapsLock,   K_F1,          K_F2,          K_F3,      K_F4,        K_F5,         K_F6,
    K_F7,         K_F8,          K_F9,          K_F10,     K_NumLock,   K_ScrollLock, K_Keypad7,
    K_Keypad8,    K_Keypad9,     K_KeypadMinus, K_Keypad4, K_Keypad5,   K_Keypad6,    K_KeypadPlus,
    K_Keypad1,    K_Keypad2,     K_Keypad3,     K_Keypad0, K_KeypadDot, K_Unknown,    K_Unknown,
    K_CH37, // Backslash
    K_F11,        K_F12,
};

static bool is_shifted(const Keyboard::TTYKeyboardState& state)
{
    if (state.capslock) return !(state.left_shift || state.right_shift);
    return state.left_shift || state.right_shift;
}

namespace Keyboard
{
    Option<char> decode_scancode_tty(u8 scancode, TTYKeyboardState& state)
    {
        if (state.ignore_next)
        {
            state.ignore_next = false;
            return {};
        }

        // FIXME: Support extended scancodes.
        if (scancode == EXTENDED_KEY_CODE)
        {
            state.ignore_next = true;
            return {};
        }

        bool released = is_key_released(scancode);

        if (scancode == LEFT_SHIFT)
        {
            state.left_shift = !released;
            return {};
        }

        if (scancode == RIGHT_SHIFT)
        {
            state.right_shift = !released;
            return {};
        }

        if (scancode == CAPS_LOCK)
        {
            if (!released) state.capslock = !state.capslock;
            return {};
        }

        if (scancode == LEFT_CONTROL)
        {
            if (released) state.left_control = false;
            else
                state.left_control = true;
            return {};
        }

        if (should_ignore_key(scancode)) return {};

        if (released) return {};

        if (state.left_control)
        {
            char key;
            if (is_shifted(state)) key = shifted_key_table[scancode];
            else
                key = key_table[scancode];
            if (_islower(key)) key = (char)_toupper(key);
            if (_isupper(key)) return key - 0x40;
            if (key == '@') return key - 0x40;
            if (key > 'Z' && key < '`') return key - 0x40;
            if (key == '?') return 0x7f;
        }

        if (is_shifted(state)) return shifted_key_table[scancode];
        return key_table[scancode];
    }

    Option<KeyboardPacket> decode_scancode(u8 scancode, KeyboardState& state)
    {
        if (state.ignore_next)
        {
            state.ignore_next = false;
            return {};
        }

        // FIXME: Support extended scancodes.
        if (scancode == EXTENDED_KEY_CODE)
        {
            state.ignore_next = true;
            return {};
        }

        bool released = is_key_released(scancode);

        return KeyboardPacket { keycode_table[scancode], released };
    }
}