#include "arch/Keyboard.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;
}

static bool g_ignore_next { false };
static bool g_left_shift { false };
static bool g_right_shift { false };
static bool g_capslock { false };

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 .
};

static bool is_shifted()
{
    if (g_capslock) return !(g_left_shift || g_right_shift);
    return g_left_shift || g_right_shift;
}

int g_modifiers = 0;

namespace Keyboard
{
    Option<char> decode_scancode(u8 scancode)
    {
        if (g_ignore_next)
        {
            g_ignore_next = false;
            return {};
        }

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

        bool released = is_key_released(scancode);

        if (scancode == LEFT_SHIFT)
        {
            g_left_shift = !released;
            return {};
        }

        if (scancode == RIGHT_SHIFT)
        {
            g_right_shift = !released;
            return {};
        }

        if (scancode == CAPS_LOCK)
        {
            if (!released) g_capslock = !g_capslock;
            return {};
        }

        if (scancode == LEFT_ALT)
        {
            if (released) g_modifiers &= ~LeftAlt;
            else
                g_modifiers |= LeftAlt;
            return {};
        }

        if (scancode == LEFT_CONTROL)
        {
            if (released) g_modifiers &= ~LeftControl;
            else
                g_modifiers |= LeftControl;
            return {};
        }

        if (should_ignore_key(scancode)) return {};

        if (released) return {};

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

    int modifiers()
    {
        return g_modifiers;
    }
}