ConsoleLib/ConsoleLib.cpp

#include "include/ConsoleLib.hpp"

#include <atomic>
#include <cerrno>
#include <chrono>
#include <cstdlib>
#include <mutex>
#include <string>
#include <thread>

#ifdef _WIN32
#include <windows.h>
#include <winsock2.h>
#include <ws2tcpip.h>
#endif

#ifdef __unix
#include <arpa/inet.h>
#include <fcntl.h>
#include <netinet/in.h>
#include <sys/socket.h>
#include <termios.h>
#include <unistd.h>
#endif

namespace ConsoleLib {
using ::Color;

namespace {

std::array<bool, KEYCODE_COUNT> g_keyTable{};
std::mutex g_keyMutex;
std::atomic<bool> g_inputRunning{false};
std::atomic<std::uint8_t> g_inputFrequency{DEFAULT_CONSOLE_INPUT_FREQUENCY};
std::thread g_inputThread;

std::chrono::milliseconds currentPollInterval() {
    std::uint8_t frequency = g_inputFrequency.load(std::memory_order_relaxed);
    if (frequency == 0) {
        frequency = 1;
    }
    return std::chrono::milliseconds(1000 / frequency);
}

void recordKey(Keycode keycode) {
    if (keycode <= KEY_UNKNOWN || keycode >= KEY_COUNT) {
        return;
    }

    std::lock_guard lock(g_keyMutex);
    g_keyTable[static_cast<size_t>(keycode)] = true;
}

Keycode parseEscapeSequence(const char* buffer, size_t length) {
    if (length < 2) {
        return KEY_UNKNOWN;
    }

    const auto first = static_cast<unsigned char>(buffer[0]);

    // Unix/Linux VT 序列：ESC [ X / ESC O X
    if (first == 0x1B) {
        if (length < 3) {
            return KEY_UNKNOWN;
        }

        if (buffer[1] == '[') {
            switch (buffer[2]) {
                case 'A':
                    return KEY_UP;
                case 'B':
                    return KEY_DOWN;
                case 'C':
                    return KEY_RIGHT;
                case 'D':
                    return KEY_LEFT;
                case 'H':
                    return KEY_HOME;
                case 'F':
                    return KEY_END;
                default:
                    return KEY_UNKNOWN;
            }
        }

        if (buffer[1] == 'O') {
            switch (buffer[2]) {
                case 'P':
                    return KEY_F1;
                case 'Q':
                    return KEY_F2;
                case 'R':
                    return KEY_F3;
                case 'S':
                    return KEY_F4;
                default:
                    return KEY_UNKNOWN;
            }
        }

        return KEY_UNKNOWN;
    }

    // Windows 扩展键序列：0x00 / 0xE0 + 扫描码。
    if (first == 0x00 || first == 0xE0) {
        switch (buffer[1]) {
            case 'H':
                return KEY_UP;
            case 'P':
                return KEY_DOWN;
            case 'K':
                return KEY_LEFT;
            case 'M':
                return KEY_RIGHT;
            case 'G':
                return KEY_HOME;
            case 'O':
                return KEY_END;
            case ';':
                return KEY_F1;
            case '<':
                return KEY_F2;
            case '=':
                return KEY_F3;
            case '>':
                return KEY_F4;
            default:
                return KEY_UNKNOWN;
        }
    }

    return KEY_UNKNOWN;
}

std::string makeAnsiColorPrefix(Color color) {
    return "\033[38;2;" + std::to_string((color >> 16) & 0xFFu) + ";" +
           std::to_string((color >> 8) & 0xFFu) + ";" + std::to_string(color & 0xFFu) + "m";
}

}  // namespace

void startInputThread();
void restoreConsoleMode();

#ifdef _WIN32

static auto g_hStdin = INVALID_HANDLE_VALUE;
static auto g_hStdout = INVALID_HANDLE_VALUE;
static DWORD g_dwOriginalInMode = 0;
static DWORD g_dwOriginalOutMode = 0;
static bool g_bInitialized = false;

void init() {
    if (g_bInitialized) {
        startInputThread();
        return;
    }

    g_hStdin = GetStdHandle(STD_INPUT_HANDLE);
    g_hStdout = GetStdHandle(STD_OUTPUT_HANDLE);

    if (g_hStdin == INVALID_HANDLE_VALUE || g_hStdout == INVALID_HANDLE_VALUE) {
        return;
    }

    if (!GetConsoleMode(g_hStdin, &g_dwOriginalInMode)) {
        return;
    }

    DWORD dwNewInMode = g_dwOriginalInMode;
    dwNewInMode &= ~(ENABLE_LINE_INPUT | ENABLE_ECHO_INPUT);
    dwNewInMode |= ENABLE_VIRTUAL_TERMINAL_INPUT;

    if (!SetConsoleMode(g_hStdin, dwNewInMode)) {
        return;
    }

    if (!GetConsoleMode(g_hStdout, &g_dwOriginalOutMode)) {
        return;
    }

    DWORD dwNewOutMode = g_dwOriginalOutMode;
    dwNewOutMode |= ENABLE_VIRTUAL_TERMINAL_PROCESSING;

    if (!SetConsoleMode(g_hStdout, dwNewOutMode)) {
        return;
    }

    g_bInitialized = true;
    startInputThread();
}

void clearConsole() {
    const HANDLE hOut = GetStdHandle(STD_OUTPUT_HANDLE);
    if (hOut == INVALID_HANDLE_VALUE) {
        return;
    }

    CONSOLE_SCREEN_BUFFER_INFO csbi;
    DWORD count;
    const COORD homeCoords = {0, 0};

    if (!GetConsoleScreenBufferInfo(hOut, &csbi)) {
        return;
    }

    count = csbi.dwSize.X * csbi.dwSize.Y;
    FillConsoleOutputCharacterA(hOut, ' ', count, homeCoords, &count);
    FillConsoleOutputAttribute(hOut, csbi.wAttributes, count, homeCoords, &count);
    SetConsoleCursorPosition(hOut, homeCoords);
}

void RawPrint(std::string_view text) {
    if (g_hStdout == INVALID_HANDLE_VALUE) {
        return;
    }

    const DWORD bytesToWrite = text.size();
    DWORD bytesWritten = 0;
    WriteFile(g_hStdout, text.data(), bytesToWrite, &bytesWritten, nullptr);
}

void printToConsole(std::string_view data) { RawPrint(data); }

void printToConsole(std::string_view data, Color color) {
    std::string prefix = makeAnsiColorPrefix(color);
    RawPrint(prefix);
    RawPrint(data);
    RawPrint("\033[0m");
}

char getinput() {
    if (!g_bInitialized || g_hStdin == INVALID_HANDLE_VALUE) {
        return 0;
    }

    char buffer[16] = {0};
    DWORD charsRead = 0;

    if (ReadFile(g_hStdin, buffer, sizeof(buffer) - 1, &charsRead, NULL) && charsRead > 0) {
        if (charsRead == 1) {
            return buffer[0];
        }

        if (const Keycode key = parseEscapeSequence(buffer, charsRead); key != KEY_UNKNOWN) {
            return 0;  // 特殊键请使用 getConsoleKeyInput() / isKeyPressed()
        }
        return buffer[0];
    }

    return 0;
}

void restoreConsoleMode() {
    if (!g_bInitialized) {
        return;
    }

    if (g_hStdin != INVALID_HANDLE_VALUE) {
        SetConsoleMode(g_hStdin, g_dwOriginalInMode);
    }
    if (g_hStdout != INVALID_HANDLE_VALUE) {
        SetConsoleMode(g_hStdout, g_dwOriginalOutMode);
    }

    g_bInitialized = false;
}

void startInputThread() {
    if (bool expected = false; !g_inputRunning.compare_exchange_strong(expected, true)) {
        return;
    }

    if (g_inputThread.joinable()) {
        g_inputThread.join();
    }

    g_inputThread = std::thread([] {
        while (g_inputRunning.load(std::memory_order_relaxed)) {
            if (g_bInitialized && g_hStdin != INVALID_HANDLE_VALUE) {
                DWORD available = 0;
                if (GetNumberOfConsoleInputEvents(g_hStdin, &available) && available > 0) {
                    char buffer[16] = {0};
                    DWORD charsRead = 0;

                    if (ReadFile(g_hStdin, buffer, sizeof(buffer) - 1, &charsRead, nullptr) &&
                        charsRead > 0) {
                        if (charsRead == 1) {
                            recordKey(static_cast<Keycode>(static_cast<unsigned char>(buffer[0])));
                        } else {
                            Keycode key = parseEscapeSequence(buffer, charsRead);
                            if (key != KEY_UNKNOWN) {
                                recordKey(key);
                            } else {
                                recordKey(KEY_ESCAPE);
                            }
                        }
                    }
                }
            }

            std::this_thread::sleep_for(currentPollInterval());
        }
    });
}

#endif  // _WIN32

#ifdef __unix

static struct termios original_termios;
static bool raw_mode_enabled = false;

void init() {
    if (!raw_mode_enabled) {
        struct termios raw;
        if (tcgetattr(STDIN_FILENO, &original_termios) == -1) {
            return;
        }

        raw = original_termios;

        raw.c_lflag &= ~(ECHO | ICANON | IEXTEN);
        raw.c_iflag &= ~(IXON | ICRNL | BRKINT | INPCK | ISTRIP);
        raw.c_oflag &= ~(OPOST);
        raw.c_cflag |= (CS8);

        raw.c_cc[VMIN] = 0;
        raw.c_cc[VTIME] = 0;

        if (tcsetattr(STDIN_FILENO, TCSAFLUSH, &raw) == -1) {
            return;
        }

        raw_mode_enabled = true;
    }

    startInputThread();
}

void RawPrint(std::string_view text) { write(STDOUT_FILENO, text.data(), text.size()); }

void printToConsole(std::string_view data) { RawPrint(data); }

void printToConsole(std::string_view data, Color color) {
    std::string prefix = makeAnsiColorPrefix(color);
    RawPrint(prefix);
    RawPrint(data);
    RawPrint("\033[0m");
}

void clearConsole() { RawPrint("\033[2J\033[H"); }

char getinput() {
    char ch = 0;
    return read(STDIN_FILENO, &ch, 1) == 1 ? ch : 0;
}

void restoreConsoleMode() {
    if (!raw_mode_enabled) {
        return;
    }

    tcsetattr(STDIN_FILENO, TCSAFLUSH, &original_termios);
    raw_mode_enabled = false;
}

void startInputThread() {
    bool expected = false;
    if (!g_inputRunning.compare_exchange_strong(expected, true)) {
        return;
    }

    if (g_inputThread.joinable()) {
        g_inputThread.join();
    }

    g_inputThread = std::thread([] {
        std::string pending;
        char buffer[32] = {0};

        while (g_inputRunning.load(std::memory_order_relaxed)) {
            ssize_t n = read(STDIN_FILENO, buffer, sizeof(buffer));
            if (n > 0) {
                pending.append(buffer, static_cast<size_t>(n));

                size_t pos = 0;
                while (pos < pending.size()) {
                    unsigned char c = static_cast<unsigned char>(pending[pos]);

                    if (c == 0x1B) {
                        if (pending.size() - pos >= 2 &&
                            (pending[pos + 1] == '[' || pending[pos + 1] == 'O')) {
                            if (pending.size() - pos >= 3) {
                                Keycode key = parseEscapeSequence(pending.data() + pos, 3);
                                recordKey(key);
                                pos += 3;
                                continue;
                            }
                            break;  // 序列还不完整，等待下一次读取
                        }

                        recordKey(KEY_ESCAPE);
                        ++pos;
                        continue;
                    }

                    if (c < ASCII_TABLE_SIZE) {
                        recordKey(static_cast<Keycode>(c));
                    }
                    ++pos;
                }

                pending.erase(0, pos);
            } else if (n < 0 && errno != EINTR) {
                // If stdin is no longer readable, stop polling instead of spinning
                // forever.
                g_inputRunning.store(false);
                break;
            }

            std::this_thread::sleep_for(currentPollInterval());
        }
    });
}

#endif  // __unix

void stopInputThread() {
    g_inputRunning.store(false);

    if (g_inputThread.joinable()) {
        g_inputThread.join();
    }
}

std::array<bool, KEYCODE_COUNT>& getConsoleKeyInput() {
    static thread_local std::array<bool, KEYCODE_COUNT> snapshot{};

    std::lock_guard<std::mutex> lock(g_keyMutex);
    snapshot = g_keyTable;
    g_keyTable.fill(false);

    return snapshot;
}

bool isKeyPressed(Keycode key) {
    if (key <= KEY_UNKNOWN || key >= KEY_COUNT) {
        return false;
    }

    std::lock_guard<std::mutex> lock(g_keyMutex);
    return g_keyTable[static_cast<size_t>(key)];
}
void sleep_ms(unsigned int ms) {
#ifdef _WIN32
    Sleep(static_cast<DWORD>(ms));
#else
    std::this_thread::sleep_for(std::chrono::milliseconds(ms));
#endif
}

void setInputFrequency(std::uint8_t frequency) {
    g_inputFrequency.store(frequency == 0 ? 1 : frequency, std::memory_order_relaxed);
}

std::uint8_t getInputFrequency() { return g_inputFrequency.load(std::memory_order_relaxed); }

void quit(int exitCode) {
    stopInputThread();
    restoreConsoleMode();
    exit(exitCode);
}
}  // namespace ConsoleLib