// 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/ddk/module.h>
#include <ntlite/ddi/input_device.h>
#include <ntlite/ddk/debug.h>
#include <ntlite/ddk/dpc.h>
#include <ntlite/ddk/interrupt.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_REG_SCR        7u

#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

#define SERIAL_RX_QUEUE_SIZE 1024u
#define SERIAL_RX_DRAIN_LIMIT 256u
#define SERIAL_RX_QUEUE_MASK (SERIAL_RX_QUEUE_SIZE - 1u)

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

static KDPC gSerialRxDpc;
static volatile SERIAL_ESCAPE_STATE gSerialEscapeState;
static volatile uint8_t gSerialCsiNumber;
static volatile bool gSerialDropNextLf;
static uint8_t gSerialRxQueue[SERIAL_RX_QUEUE_SIZE];
static volatile uint32_t gSerialRxHead;
static volatile uint32_t gSerialRxTail;
static volatile uint32_t gSerialRxDropped;

// 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);
}

// Probes the 16550 scratch register so an absent COM1 device is rejected before any receive loop is entered.
static bool SerialProbeCom1(void) {
    uint8_t original = SerialRead(SERIAL_REG_SCR);

    SerialWrite(SERIAL_REG_SCR, 0x5Au);
    uint8_t first = SerialRead(SERIAL_REG_SCR);
    SerialWrite(SERIAL_REG_SCR, 0xA5u);
    uint8_t second = SerialRead(SERIAL_REG_SCR);
    SerialWrite(SERIAL_REG_SCR, original);

    return first == 0x5Au && second == 0xA5u;
}

// 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 queued receive bytes at DISPATCH_LEVEL.
static void SerialRxDpcRoutine(KDPC* dpc, void* context, void* systemArgument1, void* systemArgument2) {
    (void)dpc;
    (void)context;
    (void)systemArgument1;
    (void)systemArgument2;

    while (gSerialRxHead != gSerialRxTail) {
        uint32_t head = gSerialRxHead;
        uint8_t ch = gSerialRxQueue[head];
        gSerialRxHead = (head + 1u) & SERIAL_RX_QUEUE_MASK;
        SerialReceiveByte(ch);
    }
}

// Queues a received byte for bottom-half processing.
static void SerialQueueRxByte(uint8_t ch) {
    uint32_t tail = gSerialRxTail;
    uint32_t next = (tail + 1u) & SERIAL_RX_QUEUE_MASK;
    if (next == gSerialRxHead) {
        gSerialRxDropped++;
        return;
    }

    gSerialRxQueue[tail] = ch;
    gSerialRxTail = next;
    KeInsertQueueDpc(&gSerialRxDpc, 0, 0);
}

// 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;
        SerialQueueRxByte(SerialRead(SERIAL_REG_DATA));
    }
}

// Configures COM1 as an 8N1 115200 baud input device while leaving receive interrupts disabled.
static KSTATUS SerialConfigureCom1(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);

    for (uint32_t i = 0; i < SERIAL_RX_DRAIN_LIMIT; i++) {
        uint8_t lsr = SerialRead(SERIAL_REG_LSR);
        if (lsr == 0xFFu)
            return NTLSTATUS_NOT_FOUND;
        if ((lsr & SERIAL_LSR_DATA_READY) == 0)
            return NTLSTATUS_SUCCESS;
        (void)SerialRead(SERIAL_REG_DATA);
    }

    return NTLSTATUS_DEVICE_ERROR;
}

// Initializes the optional COM1 input source.
static KSTATUS SerialInitialize(NTL_MODULE* module) {
    (void)module;

    if (!SerialProbeCom1()) {
        DbgPrint("serial: COM1 not present; input driver disabled\n");
        return NTLSTATUS_NOT_FOUND;
    }

    KSTATUS st = SerialConfigureCom1();
    if (!KSUCCESS(st)) {
        DbgPrint("serial: COM1 did not become idle during initialization, status=0x%08X\n", st);
        return st;
    }

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

    KeInitializeDpc(&gSerialRxDpc, SerialRxDpcRoutine, 0);
    gSerialEscapeState = SerialEscapeNone;
    gSerialCsiNumber = 0;
    gSerialDropNextLf = false;
    gSerialRxHead = 0;
    gSerialRxTail = 0;
    gSerialRxDropped = 0;

    KeIrqRegister(SERIAL_COM1_IRQ, SerialIsr);
    SerialWrite(SERIAL_REG_IER, SERIAL_IER_RX_READY);

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

NTL_DECLARE_MODULE(SerialInitialize);