// 16550-compatible serial input driver. This file routes COM1 receive bytes into the canonical input core.

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

#include <ntlite/status.h>
#include <ntlite/ddi/hal_api.h>
#include <ntlite/ddi/input_device.h>
#include <kernel/ke_debug.h>
#include <kernel/ke_irq.h>

#define SERIAL_COM1_BASE 0x3F8u
#define SERIAL_COM1_IRQ  4u

#define SERIAL_REG_DATA       0u
#define SERIAL_REG_IER        1u
#define SERIAL_REG_IIR        2u
#define SERIAL_REG_FCR        2u
#define SERIAL_REG_LCR        3u
#define SERIAL_REG_MCR        4u
#define SERIAL_REG_LSR        5u

#define SERIAL_LSR_DATA_READY 0x01u
#define SERIAL_LSR_TX_EMPTY   0x20u
#define SERIAL_LCR_DLAB       0x80u
#define SERIAL_MCR_DTR        0x01u
#define SERIAL_MCR_RTS        0x02u
#define SERIAL_MCR_OUT2       0x08u

#define SERIAL_IER_RX_READY   0x01u
#define SERIAL_FCR_ENABLE     0x01u
#define SERIAL_FCR_CLEAR_RX   0x02u
#define SERIAL_FCR_CLEAR_TX   0x04u
#define SERIAL_FCR_TRIGGER_14 0xC0u

typedef enum SERIAL_ESCAPE_STATE {
    SerialEscapeNone = 0,
    SerialEscapeEsc,
    SerialEscapeCsi,
    SerialEscapeCsiNumber,
} SERIAL_ESCAPE_STATE;

static volatile SERIAL_ESCAPE_STATE gSerialEscapeState;
static volatile uint8_t gSerialCsiNumber;
static volatile bool gSerialDropNextLf;

// Reads one byte from a COM register.
static uint8_t SerialRead(uint16_t reg) {
    return __inbyte((unsigned short)(SERIAL_COM1_BASE + reg));
}

// Writes one byte to a COM register.
static void SerialWrite(uint16_t reg, uint8_t value) {
    __outbyte((unsigned short)(SERIAL_COM1_BASE + reg), value);
}

// Submits one input key transition.
static void SerialSubmitTransition(INPUT_KEY key, bool pressed) {
    INPUT_KEY_EVENT event;
    event.Key = key;
    event.Pressed = pressed;
    KeInputSubmitKey(&event);
}

// Submits one logical key press and release.
static void SerialSubmitKey(INPUT_KEY key) {
    if (key == INPUT_KEY_NONE) return;
    SerialSubmitTransition(key, true);
    SerialSubmitTransition(key, false);
}

// Submits one shifted key press and release.
static void SerialSubmitShiftedKey(INPUT_KEY key) {
    if (key == INPUT_KEY_NONE) return;
    SerialSubmitTransition(INPUT_KEY_LEFT_SHIFT, true);
    SerialSubmitKey(key);
    SerialSubmitTransition(INPUT_KEY_LEFT_SHIFT, false);
}

// Submits one control-modified key press and release.
static void SerialSubmitCtrlKey(INPUT_KEY key) {
    if (key == INPUT_KEY_NONE) return;
    SerialSubmitTransition(INPUT_KEY_LEFT_CTRL, true);
    SerialSubmitKey(key);
    SerialSubmitTransition(INPUT_KEY_LEFT_CTRL, false);
}

// Converts an ASCII printable byte into input events.
static void SerialSubmitPrintable(uint8_t ch) {
    if (ch >= 'a' && ch <= 'z') {
        SerialSubmitKey((INPUT_KEY)(INPUT_KEY_A + (ch - 'a')));
        return;
    }
    if (ch >= 'A' && ch <= 'Z') {
        SerialSubmitShiftedKey((INPUT_KEY)(INPUT_KEY_A + (ch - 'A')));
        return;
    }
    if (ch >= '1' && ch <= '9') {
        SerialSubmitKey((INPUT_KEY)(INPUT_KEY_1 + (ch - '1')));
        return;
    }
    if (ch == '0') {
        SerialSubmitKey(INPUT_KEY_0);
        return;
    }

    switch (ch) {
        case ' ': SerialSubmitKey(INPUT_KEY_SPACE); return;
        case '-': SerialSubmitKey(INPUT_KEY_MINUS); return;
        case '=': SerialSubmitKey(INPUT_KEY_EQUAL); return;
        case '[': SerialSubmitKey(INPUT_KEY_LEFT_BRACKET); return;
        case ']': SerialSubmitKey(INPUT_KEY_RIGHT_BRACKET); return;
        case '\\': SerialSubmitKey(INPUT_KEY_BACKSLASH); return;
        case ';': SerialSubmitKey(INPUT_KEY_SEMICOLON); return;
        case '\'': SerialSubmitKey(INPUT_KEY_APOSTROPHE); return;
        case '`': SerialSubmitKey(INPUT_KEY_GRAVE); return;
        case ',': SerialSubmitKey(INPUT_KEY_COMMA); return;
        case '.': SerialSubmitKey(INPUT_KEY_PERIOD); return;
        case '/': SerialSubmitKey(INPUT_KEY_SLASH); return;

        case '!': SerialSubmitShiftedKey(INPUT_KEY_1); return;
        case '@': SerialSubmitShiftedKey(INPUT_KEY_2); return;
        case '#': SerialSubmitShiftedKey(INPUT_KEY_3); return;
        case '$': SerialSubmitShiftedKey(INPUT_KEY_4); return;
        case '%': SerialSubmitShiftedKey(INPUT_KEY_5); return;
        case '^': SerialSubmitShiftedKey(INPUT_KEY_6); return;
        case '&': SerialSubmitShiftedKey(INPUT_KEY_7); return;
        case '*': SerialSubmitShiftedKey(INPUT_KEY_8); return;
        case '(': SerialSubmitShiftedKey(INPUT_KEY_9); return;
        case ')': SerialSubmitShiftedKey(INPUT_KEY_0); return;
        case '_': SerialSubmitShiftedKey(INPUT_KEY_MINUS); return;
        case '+': SerialSubmitShiftedKey(INPUT_KEY_EQUAL); return;
        case '{': SerialSubmitShiftedKey(INPUT_KEY_LEFT_BRACKET); return;
        case '}': SerialSubmitShiftedKey(INPUT_KEY_RIGHT_BRACKET); return;
        case '|': SerialSubmitShiftedKey(INPUT_KEY_BACKSLASH); return;
        case ':': SerialSubmitShiftedKey(INPUT_KEY_SEMICOLON); return;
        case '"': SerialSubmitShiftedKey(INPUT_KEY_APOSTROPHE); return;
        case '~': SerialSubmitShiftedKey(INPUT_KEY_GRAVE); return;
        case '<': SerialSubmitShiftedKey(INPUT_KEY_COMMA); return;
        case '>': SerialSubmitShiftedKey(INPUT_KEY_PERIOD); return;
        case '?': SerialSubmitShiftedKey(INPUT_KEY_SLASH); return;
        default: return;
    }
}

// Submits a completed CSI escape sequence.
static void SerialSubmitCsi(uint8_t finalByte) {
    uint8_t number = gSerialCsiNumber;
    gSerialEscapeState = SerialEscapeNone;
    gSerialCsiNumber = 0;

    switch (finalByte) {
        case 'A': SerialSubmitKey(INPUT_KEY_UP); return;
        case 'B': SerialSubmitKey(INPUT_KEY_DOWN); return;
        case 'C': SerialSubmitKey(INPUT_KEY_RIGHT); return;
        case 'D': SerialSubmitKey(INPUT_KEY_LEFT); return;
        case 'H': SerialSubmitKey(INPUT_KEY_HOME); return;
        case 'F': SerialSubmitKey(INPUT_KEY_END); return;
        case '~':
            switch (number) {
                case 1: SerialSubmitKey(INPUT_KEY_HOME); return;
                case 3: SerialSubmitKey(INPUT_KEY_DELETE); return;
                case 4: SerialSubmitKey(INPUT_KEY_END); return;
                case 5: SerialSubmitKey(INPUT_KEY_PAGE_UP); return;
                case 6: SerialSubmitKey(INPUT_KEY_PAGE_DOWN); return;
                default: return;
            }
        default: return;
    }
}

// Handles one byte after an ESC prefix.
static bool SerialHandleEscapeByte(uint8_t ch) {
    if (gSerialEscapeState == SerialEscapeNone) return false;

    if (gSerialEscapeState == SerialEscapeEsc) {
        if (ch == '[') {
            gSerialEscapeState = SerialEscapeCsi;
            gSerialCsiNumber = 0;
            return true;
        }

        gSerialEscapeState = SerialEscapeNone;
        gSerialCsiNumber = 0;
        SerialSubmitKey(INPUT_KEY_ESCAPE);
        return false;
    }

    if (gSerialEscapeState == SerialEscapeCsi || gSerialEscapeState == SerialEscapeCsiNumber) {
        if (ch >= '0' && ch <= '9') {
            gSerialEscapeState = SerialEscapeCsiNumber;
            if (gSerialCsiNumber < 25)
                gSerialCsiNumber = (uint8_t)(gSerialCsiNumber * 10u + (ch - '0'));
            return true;
        }

        if (ch == ';')
            return true;

        SerialSubmitCsi(ch);
        return true;
    }

    gSerialEscapeState = SerialEscapeNone;
    gSerialCsiNumber = 0;
    return false;
}

// Routes one received serial byte into the input core.
static void SerialReceiveByte(uint8_t ch) {
    if (SerialHandleEscapeByte(ch)) return;

    if (ch == 0x1Bu) {
        gSerialEscapeState = SerialEscapeEsc;
        gSerialCsiNumber = 0;
        return;
    }

    if (ch == '\r') {
        SerialSubmitKey(INPUT_KEY_ENTER);
        gSerialDropNextLf = true;
        return;
    }

    if (ch == '\n') {
        if (gSerialDropNextLf) {
            gSerialDropNextLf = false;
            return;
        }
        SerialSubmitKey(INPUT_KEY_ENTER);
        return;
    }

    gSerialDropNextLf = false;

    if (ch == '\t') {
        SerialSubmitKey(INPUT_KEY_TAB);
        return;
    }

    if (ch == '\b' || ch == 0x7Fu) {
        SerialSubmitKey(INPUT_KEY_BACKSPACE);
        return;
    }

    if (ch == 0x03u) {
        SerialSubmitCtrlKey(INPUT_KEY_C);
        return;
    }

    if (ch >= 0x20u && ch <= 0x7Eu)
        SerialSubmitPrintable(ch);
}

// Handles receive interrupts for COM1.
static void SerialIsr(void) {
    for (uint32_t i = 0; i < 256u; i++) {
        if ((SerialRead(SERIAL_REG_LSR) & SERIAL_LSR_DATA_READY) == 0)
            break;
        SerialReceiveByte(SerialRead(SERIAL_REG_DATA));
    }
}

// Initializes COM1 as an 8N1 115200 baud input device.
static void SerialInitializeCom1(void) {
    SerialWrite(SERIAL_REG_IER, 0x00u);
    SerialWrite(SERIAL_REG_LCR, SERIAL_LCR_DLAB);
    SerialWrite(SERIAL_REG_DATA, 0x01u);
    SerialWrite(SERIAL_REG_IER, 0x00u);
    SerialWrite(SERIAL_REG_LCR, 0x03u);
    SerialWrite(SERIAL_REG_FCR, SERIAL_FCR_ENABLE | SERIAL_FCR_CLEAR_RX |
                            SERIAL_FCR_CLEAR_TX | SERIAL_FCR_TRIGGER_14);
    SerialWrite(SERIAL_REG_MCR, SERIAL_MCR_DTR | SERIAL_MCR_RTS | SERIAL_MCR_OUT2);

    while (SerialRead(SERIAL_REG_LSR) & SERIAL_LSR_DATA_READY)
        (void)SerialRead(SERIAL_REG_DATA);

    SerialWrite(SERIAL_REG_IER, SERIAL_IER_RX_READY);
}

// DriverEntry registers COM1 as an input source.
KSTATUS DriverEntry(const HAL_API* hal) {
    (void)hal;

    uint32_t deviceId = 0;
    KSTATUS st = KeInputRegisterDevice("serial-com1", 0, &deviceId);
    if (!KSUCCESS(st))
        return st;

    gSerialEscapeState = SerialEscapeNone;
    gSerialCsiNumber = 0;
    gSerialDropNextLf = false;

    SerialInitializeCom1();
    KeIrqRegister(SERIAL_COM1_IRQ, SerialIsr);

    DbgPrint("serial: initialized COM1 input device=%u irq=%u\n", deviceId, SERIAL_COM1_IRQ);
    return STATUS_SUCCESS;
}