// Kernel input core. This file owns canonical keyboard and mouse state shared by input drivers and syscalls.

#include <stdint.h>
#include <stdbool.h>
#include <stddef.h>
#include <string.h>

#include <kernel/ke_input.h>
#include <kernel/ke_debug.h>
#include <kernel/ke_display.h>
#include <kernel/ke_event.h>
#include <kernel/ke_dpc.h>
#include <kernel/ob_object.h>

#define KEY_QUEUE_SIZE   256u
#define KEY_QUEUE_MASK   (KEY_QUEUE_SIZE - 1u)
#define MOUSE_QUEUE_SIZE 128u
#define MOUSE_QUEUE_MASK (MOUSE_QUEUE_SIZE - 1u)
#define MAX_INPUT_KEYS   256u
#define MAX_INPUT_DEVICES 16u
#define KEY_REPEAT_DELAY_US 500000ull
#define KEY_REPEAT_PERIOD_US 33000ull

typedef struct KEY_CHAR_PAIR {
    INPUT_KEY key;
    char normal;
    char shifted;
} KEY_CHAR_PAIR;

static const HAL_API* gInputHal;

static volatile uint8_t gKeyDown[MAX_INPUT_KEYS];
static volatile uint32_t gModifiers;

typedef struct INPUT_KEY_QUEUE_ENTRY {
    INPUT_KEY key;
    char      ch;
    uint32_t  modifiers;
} INPUT_KEY_QUEUE_ENTRY;

static INPUT_KEY_QUEUE_ENTRY gKeyQueue[KEY_QUEUE_SIZE];
static volatile int gKeyHead;
static volatile int gKeyTail;
static volatile INPUT_KEY gRepeatKey;
static volatile uint64_t  gRepeatNextTick;
static KDPC gRepeatDpc;

static INPUT_MOUSE_EVENT gMouseQueue[MOUSE_QUEUE_SIZE];
static volatile int      gMouseHead;
static volatile int      gMouseTail;
static volatile int      gMouseX;
static volatile int      gMouseY;
static volatile int      gMouseBoundsW = 1920;
static volatile int      gMouseBoundsH = 1080;
static volatile uint8_t  gMouseButtons;

static volatile uint32_t gNextInputDeviceId = 1;
static EVENT* gKeyboardEvent;

static char KiInputKeyToChar(INPUT_KEY key);
static void KiInputServiceRepeat(void);

static const KEY_CHAR_PAIR gCharMap[] = {
    { INPUT_KEY_1, '1', '!' }, { INPUT_KEY_2, '2', '@' },
    { INPUT_KEY_3, '3', '#' }, { INPUT_KEY_4, '4', '$' },
    { INPUT_KEY_5, '5', '%' }, { INPUT_KEY_6, '6', '^' },
    { INPUT_KEY_7, '7', '&' }, { INPUT_KEY_8, '8', '*' },
    { INPUT_KEY_9, '9', '(' }, { INPUT_KEY_0, '0', ')' },
    { INPUT_KEY_MINUS, '-', '_' }, { INPUT_KEY_EQUAL, '=', '+' },
    { INPUT_KEY_LEFT_BRACKET, '[', '{' }, { INPUT_KEY_RIGHT_BRACKET, ']', '}' },
    { INPUT_KEY_BACKSLASH, '\\', '|' }, { INPUT_KEY_SEMICOLON, ';', ':' },
    { INPUT_KEY_APOSTROPHE, '\'', '"' }, { INPUT_KEY_GRAVE, '`', '~' },
    { INPUT_KEY_COMMA, ',', '<' }, { INPUT_KEY_PERIOD, '.', '>' },
    { INPUT_KEY_SLASH, '/', '?' },
    { INPUT_KEY_SPACE, ' ', ' ' }, { INPUT_KEY_TAB, '\t', '\t' },
    { INPUT_KEY_ENTER, '\n', '\n' }, { INPUT_KEY_BACKSPACE, '\b', '\b' },
    { INPUT_KEY_KP_DIVIDE, '/', '/' }, { INPUT_KEY_KP_MULTIPLY, '*', '*' },
    { INPUT_KEY_KP_MINUS, '-', '-' }, { INPUT_KEY_KP_PLUS, '+', '+' },
    { INPUT_KEY_KP_ENTER, '\n', '\n' }, { INPUT_KEY_KP_1, '1', '1' },
    { INPUT_KEY_KP_2, '2', '2' }, { INPUT_KEY_KP_3, '3', '3' },
    { INPUT_KEY_KP_4, '4', '4' }, { INPUT_KEY_KP_5, '5', '5' },
    { INPUT_KEY_KP_6, '6', '6' }, { INPUT_KEY_KP_7, '7', '7' },
    { INPUT_KEY_KP_8, '8', '8' }, { INPUT_KEY_KP_9, '9', '9' },
    { INPUT_KEY_KP_0, '0', '0' }, { INPUT_KEY_KP_PERIOD, '.', '.' },
};

// Updates the waitable keyboard event to match queued keyboard input.
static void KiInputRefreshKeyboardEvent(void) {
    if (!gKeyboardEvent) return;
    if (gKeyHead != gKeyTail) KeSignalEvent(gKeyboardEvent);
    else KeClearEvent(gKeyboardEvent);
}

// Queues a keyboard event and wakes waiters when the event was accepted.
static void KiInputQueueKey(INPUT_KEY key, char ch) {
    if (key == INPUT_KEY_NONE) return;
    int next = (gKeyTail + 1) & (int)KEY_QUEUE_MASK;
    if (next == gKeyHead) return;
    gKeyQueue[gKeyTail].key = key;
    gKeyQueue[gKeyTail].ch = ch;
    gKeyQueue[gKeyTail].modifiers = gModifiers;
    gKeyTail = next;
    KiInputRefreshKeyboardEvent();
}

// Implements the input Queue Has Space operation.
static bool KiInputQueueHasSpace(void) {
    int next = (gKeyTail + 1) & (int)KEY_QUEUE_MASK;
    return next != gKeyHead;
}

// Implements the input Key Is Repeatable operation.
static bool KiInputKeyIsRepeatable(INPUT_KEY key) {
    if (key == INPUT_KEY_NONE) return false;

    switch (key) {
        case INPUT_KEY_LEFT_SHIFT:
        case INPUT_KEY_RIGHT_SHIFT:
        case INPUT_KEY_LEFT_CTRL:
        case INPUT_KEY_RIGHT_CTRL:
        case INPUT_KEY_LEFT_ALT:
        case INPUT_KEY_RIGHT_ALT:
        case INPUT_KEY_LEFT_GUI:
        case INPUT_KEY_RIGHT_GUI:
        case INPUT_KEY_CAPS_LOCK:
        case INPUT_KEY_NUM_LOCK:
            return false;
        default:
            return true;
    }
}

// Implements the input Now Ticks operation.
static uint64_t KiInputNowTicks(void) {
    if (gInputHal && gInputHal->Timer.QueryTicks)
        return gInputHal->Timer.QueryTicks();
    return 0;
}

// Implements the input Start Repeat operation.
static void KiInputStartRepeat(INPUT_KEY key, uint64_t now) {
    if (!KiInputKeyIsRepeatable(key)) return;
    gRepeatKey = key;
    gRepeatNextTick = now + KEY_REPEAT_DELAY_US;
}

// Implements the input Stop Repeat operation.
static void KiInputStopRepeat(INPUT_KEY key) {
    if (gRepeatKey == key) {
        gRepeatKey = INPUT_KEY_NONE;
        gRepeatNextTick = 0;
    }
}

// Implements the input Service Repeat operation.
static void KiInputServiceRepeat(void) {
    INPUT_KEY key = gRepeatKey;
    if (!KiInputKeyIsRepeatable(key) || key >= MAX_INPUT_KEYS || !gKeyDown[key]) {
        gRepeatKey = INPUT_KEY_NONE;
        gRepeatNextTick = 0;
        return;
    }

    uint64_t now = KiInputNowTicks();
    if (now == 0 || now < gRepeatNextTick)
        return;

    if (!KiInputQueueHasSpace()) {
        gRepeatNextTick = now + KEY_REPEAT_PERIOD_US;
        return;
    }

    KiInputQueueKey(key, KiInputKeyToChar(key));
    gRepeatNextTick = now + KEY_REPEAT_PERIOD_US;
}

// Implements the input Repeat Dpc Routine operation.
static void KiInputRepeatDpcRoutine(KDPC* dpc, void* context, void* systemArgument1, void* systemArgument2) {
    (void)dpc;
    (void)context;
    (void)systemArgument1;
    (void)systemArgument2;

    KiInputServiceRepeat();
}

// Implements the input Set Modifier operation.
static void KiInputSetModifier(INPUT_KEY key, bool pressed) {
    uint32_t bit = 0;
    switch (key) {
        case INPUT_KEY_LEFT_SHIFT:  bit = INPUT_MOD_LSHIFT; break;
        case INPUT_KEY_RIGHT_SHIFT: bit = INPUT_MOD_RSHIFT; break;
        case INPUT_KEY_LEFT_CTRL:   bit = INPUT_MOD_LCTRL;  break;
        case INPUT_KEY_RIGHT_CTRL:  bit = INPUT_MOD_RCTRL;  break;
        case INPUT_KEY_LEFT_ALT:    bit = INPUT_MOD_LALT;   break;
        case INPUT_KEY_RIGHT_ALT:   bit = INPUT_MOD_RALT;   break;
        case INPUT_KEY_LEFT_GUI:    bit = INPUT_MOD_LGUI;   break;
        case INPUT_KEY_RIGHT_GUI:   bit = INPUT_MOD_RGUI;   break;
        default: break;
    }

    if (bit) {
        if (pressed) gModifiers |= bit;
        else         gModifiers &= ~bit;
    }

    if (pressed && key == INPUT_KEY_CAPS_LOCK) gModifiers ^= INPUT_MOD_CAPS_LOCK;
    if (pressed && key == INPUT_KEY_NUM_LOCK)  gModifiers ^= INPUT_MOD_NUM_LOCK;
}

// Implements the input Key To Char operation.
static char KiInputKeyToChar(INPUT_KEY key) {
    if (key >= INPUT_KEY_A && key <= INPUT_KEY_Z) {
        bool shifted = (gModifiers & INPUT_MOD_SHIFT) != 0;
        bool caps    = (gModifiers & INPUT_MOD_CAPS_LOCK) != 0;
        char c = (char)('a' + (key - INPUT_KEY_A));
        if (shifted ^ caps) c = (char)(c - 'a' + 'A');
        return c;
    }

    bool shifted = (gModifiers & INPUT_MOD_SHIFT) != 0;
    for (size_t i = 0; i < sizeof(gCharMap) / sizeof(gCharMap[0]); i++) {
        if (gCharMap[i].key == key)
            return shifted ? gCharMap[i].shifted : gCharMap[i].normal;
    }
    return 0;
}

// Implements the input Initialize operation.
void KeInputInitialize(const HAL_API* hal) {
    gInputHal = hal;
    memset((void*)gKeyDown, 0, sizeof(gKeyDown));
    gModifiers = 0;
    gKeyHead = gKeyTail = 0;
    gRepeatKey = INPUT_KEY_NONE;
    gRepeatNextTick = 0;
    gMouseHead = gMouseTail = 0;
    gMouseButtons = 0;
    KeInitializeDpc(&gRepeatDpc, KiInputRepeatDpcRoutine, (void*)0);

    if (!gKeyboardEvent) {
        KSTATUS s = KeEventCreate(EventManualReset, 0, &gKeyboardEvent);
        if (!KSUCCESS(s))
            DbgPrint("Input: keyboard wait event create failed st=0x%x\n", s);
    }
    KiInputRefreshKeyboardEvent();

    int w = 1920, h = 1080;
    DISPLAY_MODE mode;
    if (KSUCCESS(KeDisplayGetMode(&mode))) {
        w = (int)mode.Width;
        h = (int)mode.Height;
    }

    KeInputMouseSetBounds(w, h);
    KeInputMouseSetPosition(w / 2, h / 2);
    KeDisplayShowCursor(true);

    DbgPrint("Input: kernel input core initialized (%dx%d)\n", w, h);
}

// Implements the input Register Device operation.
__declspec(dllexport) KSTATUS KeInputRegisterDevice(const char* name, uint32_t flags, uint32_t* outDeviceId) {
    (void)flags;
    uint32_t id = gNextInputDeviceId++;
    if (id >= MAX_INPUT_DEVICES) return STATUS_NO_MEMORY;
    if (outDeviceId) *outDeviceId = id;
    DbgPrint("Input: registered device %u (%s)\n", id, name ? name : "unnamed");
    return STATUS_SUCCESS;
}

// Implements the input Submit Key operation.
__declspec(dllexport) void KeInputSubmitKey(const INPUT_KEY_EVENT* event) {
    if (!event) return;
    INPUT_KEY key = event->Key;
    if (key == INPUT_KEY_NONE || key >= MAX_INPUT_KEYS) return;

    bool wasDown = gKeyDown[key] != 0;
    bool pressed = event->Pressed;

    if (pressed) {
        // The input core owns typematic repeat. Treat repeated hardware make
        // events as idempotent so PS/2 typematic and USB HID share one policy.
        if (wasDown)
            return;

        gKeyDown[key] = 1u;
        KiInputSetModifier(key, true);
        KiInputQueueKey(key, KiInputKeyToChar(key));
        KiInputStartRepeat(key, KiInputNowTicks());
    } else {
        if (!wasDown)
            return;

        gKeyDown[key] = 0u;
        KiInputSetModifier(key, false);
        KiInputStopRepeat(key);
    }
}

// Implements the input Commit Mouse Position operation.
static void KiInputCommitMousePosition(int nx, int ny, uint8_t buttons) {
    int bw = gMouseBoundsW;
    int bh = gMouseBoundsH;

    if (bw < 1) bw = 1;
    if (bh < 1) bh = 1;
    if (nx < 0) nx = 0;
    if (ny < 0) ny = 0;
    if (nx >= bw) nx = bw - 1;
    if (ny >= bh) ny = bh - 1;

    gMouseX = nx;
    gMouseY = ny;
    gMouseButtons = buttons;
    KeDisplaySetCursorPosition(nx, ny);
}

// Implements the input Submit Mouse operation.
__declspec(dllexport) void KeInputSubmitMouse(const INPUT_MOUSE_EVENT* event) {
    if (!event) return;

    int oldX = gMouseX;
    int oldY = gMouseY;
    int nx = oldX + event->dx;
    int ny = oldY + event->dy;

    KiInputCommitMousePosition(nx, ny, event->buttons);

    int next = (gMouseTail + 1) & (int)MOUSE_QUEUE_MASK;
    if (next != gMouseHead) {
        gMouseQueue[gMouseTail] = *event;
        gMouseTail = next;
    }
}

// Implements the input Submit Mouse Absolute operation.
__declspec(dllexport) void KeInputSubmitMouseAbsolute(const INPUT_MOUSE_ABSOLUTE_EVENT* event) {
    if (!event) return;

    int oldX = gMouseX;
    int oldY = gMouseY;
    KiInputCommitMousePosition(event->x, event->y, event->buttons);

    INPUT_MOUSE_EVENT rel;
    rel.dx = gMouseX - oldX;
    rel.dy = gMouseY - oldY;
    rel.wheel = event->wheel;
    rel.buttons = event->buttons;

    int next = (gMouseTail + 1) & (int)MOUSE_QUEUE_MASK;
    if (next != gMouseHead) {
        gMouseQueue[gMouseTail] = rel;
        gMouseTail = next;
    }
}

// Implements the input Submit Mouse Absolute Scaled operation.
__declspec(dllexport) void KeInputSubmitMouseAbsoluteScaled(uint32_t rawX, uint32_t rawY, uint32_t maxX, uint32_t maxY, int wheel, uint8_t buttons) {
    int bw = gMouseBoundsW;
    int bh = gMouseBoundsH;
    if (bw < 1) bw = 1;
    if (bh < 1) bh = 1;

    INPUT_MOUSE_ABSOLUTE_EVENT ev;
    ev.x = maxX ? (int)(((uint64_t)rawX * (uint64_t)(bw - 1)) / (uint64_t)maxX) : 0;
    ev.y = maxY ? (int)(((uint64_t)rawY * (uint64_t)(bh - 1)) / (uint64_t)maxY) : 0;
    ev.wheel = wheel;
    ev.buttons = buttons;
    KeInputSubmitMouseAbsolute(&ev);
}

INPUT_KEY KeInputKbPoll(void) {
    if (gKeyHead == gKeyTail) KiInputServiceRepeat();
    return (gKeyHead == gKeyTail) ? INPUT_KEY_NONE : gKeyQueue[gKeyHead].key;
}

// Waits until at least one keyboard event is available or the timeout expires.
KSTATUS KeInputWaitKey(uint8_t alertable, uint64_t timeoutMs) {
    if (KeInputKbPoll() != INPUT_KEY_NONE)
        return STATUS_SUCCESS;
    if (!gKeyboardEvent)
        return STATUS_FAILURE;

    for (;;) {
        EVENT* event = gKeyboardEvent;
        uint32_t selected = 0;
        KSTATUS s = KeWaitForMultipleEvents(&event, 1, alertable, timeoutMs, &selected);
        if (!KSUCCESS(s))
            return s;
        if (KeInputKbPoll() != INPUT_KEY_NONE)
            return STATUS_SUCCESS;
    }
}

// Services typematic repeat from the scheduler timer bottom half.
void KeInputTimerTick(void) {
    INPUT_KEY key = gRepeatKey;
    if (!KiInputKeyIsRepeatable(key) || key >= MAX_INPUT_KEYS || !gKeyDown[key])
        return;

    uint64_t now = KiInputNowTicks();
    if (now != 0 && now >= gRepeatNextTick)
        KeInsertQueueDpc(&gRepeatDpc, (void*)0, (void*)0);
}

// Implements the input Is Key Down operation.
bool KeInputIsKeyDown(INPUT_KEY key) {
    if (key == INPUT_KEY_NONE || key >= MAX_INPUT_KEYS) return false;
    return gKeyDown[key] != 0;
}

// Implements the input Peek Event operation.
bool KeInputPeekEvent(INPUT_KEY_RECORD* record) {
    if (!record) return false;
    if (gKeyHead == gKeyTail) KiInputServiceRepeat();
    if (gKeyHead == gKeyTail) return false;

    INPUT_KEY_QUEUE_ENTRY entry = gKeyQueue[gKeyHead];
    record->Key = entry.key;
    record->Char = entry.ch;
    record->Modifiers = entry.modifiers;
    return true;
}

// Implements the input Read Event operation.
bool KeInputReadEvent(INPUT_KEY_RECORD* record) {
    if (!KeInputPeekEvent(record)) return false;
    gKeyHead = (gKeyHead + 1) & (int)KEY_QUEUE_MASK;
    KiInputRefreshKeyboardEvent();
    return true;
}

// Implements the input Read Key operation.
INPUT_KEY KeInputReadKey(void) {
    INPUT_KEY_RECORD record;
    return KeInputReadEvent(&record) ? record.Key : INPUT_KEY_NONE;
}

// Implements the input Read Char operation.
char KeInputReadChar(void) {
    INPUT_KEY_RECORD record;
    return KeInputReadEvent(&record) ? record.Char : 0;
}

// Implements the input Get Modifiers operation.
uint32_t KeInputGetModifiers(void) {
    return gModifiers;
}

// Implements the input Mouse Peek operation.
bool KeInputMousePeek(INPUT_MOUSE_EVENT* event) {
    if (!event) return false;
    if (gMouseHead == gMouseTail) return false;

    *event = gMouseQueue[gMouseHead];
    return true;
}

// Implements the input Mouse Read operation.
bool KeInputMouseRead(INPUT_MOUSE_EVENT* event) {
    if (!KeInputMousePeek(event)) return false;
    gMouseHead = (gMouseHead + 1) & (int)MOUSE_QUEUE_MASK;
    return true;
}

// Implements the input Mouse Get Position operation.
void KeInputMouseGetPosition(int* x, int* y) {
    if (x) *x = gMouseX;
    if (y) *y = gMouseY;
}

// Implements the input Mouse Set Position operation.
void KeInputMouseSetPosition(int x, int y) {
    int bw = gMouseBoundsW;
    int bh = gMouseBoundsH;
    if (bw < 1) bw = 1;
    if (bh < 1) bh = 1;
    if (x < 0) x = 0;
    if (y < 0) y = 0;
    if (x >= bw) x = bw - 1;
    if (y >= bh) y = bh - 1;
    gMouseX = x;
    gMouseY = y;
    KeDisplaySetCursorPosition(x, y);
}

// Implements the input Mouse Set Bounds operation.
void KeInputMouseSetBounds(int w, int h) {
    if (w < 1) w = 1;
    if (h < 1) h = 1;
    gMouseBoundsW = w;
    gMouseBoundsH = h;
    KeInputMouseSetPosition(gMouseX, gMouseY);
}

// Destroys an input queue handle.
static void KiInputQueueDestroy(OBJECT* object) {
    INPUT_QUEUE* queue = CONTAINING_RECORD(object, INPUT_QUEUE, Header);
    if (queue->WaitEvent) {
        ObDereferenceObject(&queue->WaitEvent->Header);
        queue->WaitEvent = (EVENT*)0;
    }
}

// Opens the default keyboard input queue as a waitable object.
KSTATUS KeInputOpenQueue(INPUT_QUEUE** out) {
    if (!out) return STATUS_INVALID_PARAMETER;
    *out = (INPUT_QUEUE*)0;
    if (!gKeyboardEvent) return STATUS_FAILURE;

    INPUT_QUEUE* queue = (INPUT_QUEUE*)ObAllocateObject(ObjectTypeInputQueue, sizeof(INPUT_QUEUE));
    if (!queue) return STATUS_NO_MEMORY;
    queue->Header.Destroy = KiInputQueueDestroy;
    ObReferenceObject(&gKeyboardEvent->Header);
    queue->WaitEvent = gKeyboardEvent;
    *out = queue;
    return STATUS_SUCCESS;
}

// Returns the event that wakes when keyboard input can be read.
EVENT* KeInputQueueWaitEvent(INPUT_QUEUE* queue) {
    return queue ? queue->WaitEvent : (EVENT*)0;
}