#include <filesystem>
#include <linux/input-event-codes.h>
#include <thread>
#include "../../shared.hpp"
#include "../../hyprctlCompat.hpp"
#include "../shared.hpp"
#include "tests.hpp"

using namespace Hyprutils::OS;
using namespace Hyprutils::Memory;

static int         ret      = 0;
static std::string flagFile = "/tmp/hyprtester-keybinds.txt";

// Because i don't feel like changing someone elses code.
enum eKeyboardModifierIndex : uint8_t {
    MOD_SHIFT = 1,
    MOD_CAPS,
    MOD_CTRL,
    MOD_ALT,
    MOD_MOD2,
    MOD_MOD3,
    MOD_META,
    MOD_MOD5
};

static void clearFlag() {
    std::filesystem::remove(flagFile);
}

static bool checkFlag() {
    bool exists = std::filesystem::exists(flagFile);
    clearFlag();
    return exists;
}

static bool attemptCheckFlag(int attempts, int intervalMs) {
    for (int i = 0; i < attempts; i++) {
        if (checkFlag())
            return true;

        std::this_thread::sleep_for(std::chrono::milliseconds(intervalMs));
    }

    return false;
}

static std::string readKittyOutput() {
    std::string output = Tests::execAndGet("kitten @ --to unix:/tmp/hyprtester-kitty.sock get-text --extent all");
    // chop off shell prompt
    std::size_t pos = output.rfind("$");
    if (pos != std::string::npos) {
        pos += 1;
        if (pos < output.size())
            output.erase(0, pos);
    }
    // NLog::log("Kitty output: '{}'", output);
    return output;
}

static void awaitKittyPrompt() {
    // wait until we see the shell prompt, meaning it's ready for test inputs
    for (int i = 0; i < 10; i++) {
        std::string output = Tests::execAndGet("kitten @ --to unix:/tmp/hyprtester-kitty.sock get-text --extent all");
        if (output.rfind("$") == std::string::npos) {
            std::this_thread::sleep_for(std::chrono::milliseconds(200));
            continue;
        }
        return;
    }
    NLog::log("{}Error: timed out waiting for kitty prompt", Colors::RED);
}

static CUniquePointer<CProcess> spawnRemoteControlKitty() {
    auto kittyProc = Tests::spawnKitty("keybinds_test", {"-o", "allow_remote_control=yes", "--listen-on", "unix:/tmp/hyprtester-kitty.sock", "--config", "NONE", "/bin/sh"});
    // wait a bit to ensure shell prompt is sent, we are going to read the text after it
    std::this_thread::sleep_for(std::chrono::milliseconds(100));
    if (kittyProc)
        awaitKittyPrompt();
    return kittyProc;
}

static void testBind() {
    EXPECT(checkFlag(), false);
    EXPECT(getFromSocket("/keyword bind SUPER,Y,exec,touch " + flagFile), "ok");
    // press keybind
    OK(getFromSocket("/dispatch plugin:test:keybind 1,7,29"));
    // await flag
    EXPECT(attemptCheckFlag(20, 50), true);
    // release keybind
    OK(getFromSocket("/dispatch plugin:test:keybind 0,0,29"));
    EXPECT(getFromSocket("/keyword unbind SUPER,Y"), "ok");
}

static void testBindKey() {
    EXPECT(checkFlag(), false);
    EXPECT(getFromSocket("/keyword bind ,Y,exec,touch " + flagFile), "ok");
    // press keybind
    OK(getFromSocket("/dispatch plugin:test:keybind 1,0,29"));
    // await flag
    EXPECT(attemptCheckFlag(20, 50), true);
    // release keybind
    OK(getFromSocket("/dispatch plugin:test:keybind 0,0,29"));
    EXPECT(getFromSocket("/keyword unbind ,Y"), "ok");
}

static void testLongPress() {
    EXPECT(checkFlag(), false);
    EXPECT(getFromSocket("/keyword bindo SUPER,Y,exec,touch " + flagFile), "ok");
    EXPECT(getFromSocket("/keyword input:repeat_delay 100"), "ok");
    // press keybind
    OK(getFromSocket("/dispatch plugin:test:keybind 1,7,29"));
    // check no flag on short press
    std::this_thread::sleep_for(std::chrono::milliseconds(50));
    EXPECT(checkFlag(), false);
    // await repeat delay
    std::this_thread::sleep_for(std::chrono::milliseconds(200));
    EXPECT(checkFlag(), true);
    // release keybind
    OK(getFromSocket("/dispatch plugin:test:keybind 0,0,29"));
    EXPECT(getFromSocket("/keyword unbind SUPER,Y"), "ok");
}

static void testKeyLongPress() {
    EXPECT(checkFlag(), false);
    EXPECT(getFromSocket("/keyword bindo ,Y,exec,touch " + flagFile), "ok");
    EXPECT(getFromSocket("/keyword input:repeat_delay 100"), "ok");
    // press keybind
    OK(getFromSocket("/dispatch plugin:test:keybind 1,0,29"));
    // check no flag on short press
    std::this_thread::sleep_for(std::chrono::milliseconds(50));
    EXPECT(checkFlag(), false);
    // await repeat delay
    std::this_thread::sleep_for(std::chrono::milliseconds(200));
    EXPECT(checkFlag(), true);
    // release keybind
    OK(getFromSocket("/dispatch plugin:test:keybind 0,0,29"));
    EXPECT(getFromSocket("/keyword unbind ,Y"), "ok");
}

static void testLongPressRelease() {
    EXPECT(checkFlag(), false);
    EXPECT(getFromSocket("/keyword bindo SUPER,Y,exec,touch " + flagFile), "ok");
    EXPECT(getFromSocket("/keyword input:repeat_delay 100"), "ok");
    // press keybind
    OK(getFromSocket("/dispatch plugin:test:keybind 1,7,29"));
    // check no flag on short press
    std::this_thread::sleep_for(std::chrono::milliseconds(50));
    EXPECT(checkFlag(), false);
    // release keybind
    OK(getFromSocket("/dispatch plugin:test:keybind 0,0,29"));
    // await repeat delay
    std::this_thread::sleep_for(std::chrono::milliseconds(200));
    EXPECT(checkFlag(), false);
    EXPECT(getFromSocket("/keyword unbind SUPER,Y"), "ok");
}

static void testLongPressOnlyKeyRelease() {
    EXPECT(checkFlag(), false);
    EXPECT(getFromSocket("/keyword bindo SUPER,Y,exec,touch " + flagFile), "ok");
    EXPECT(getFromSocket("/keyword input:repeat_delay 100"), "ok");
    // press keybind
    OK(getFromSocket("/dispatch plugin:test:keybind 1,7,29"));
    // check no flag on short press
    std::this_thread::sleep_for(std::chrono::milliseconds(50));
    EXPECT(checkFlag(), false);
    // release key, keep modifier
    OK(getFromSocket("/dispatch plugin:test:keybind 0,7,29"));
    // await repeat delay
    std::this_thread::sleep_for(std::chrono::milliseconds(200));
    EXPECT(checkFlag(), false);
    OK(getFromSocket("/dispatch plugin:test:keybind 0,0,29"));
    EXPECT(getFromSocket("/keyword unbind SUPER,Y"), "ok");
}

static void testRepeat() {
    EXPECT(checkFlag(), false);
    EXPECT(getFromSocket("/keyword binde SUPER,Y,exec,touch " + flagFile), "ok");
    EXPECT(getFromSocket("/keyword input:repeat_delay 100"), "ok");
    // press keybind
    OK(getFromSocket("/dispatch plugin:test:keybind 1,7,29"));
    // await flag
    std::this_thread::sleep_for(std::chrono::milliseconds(200));
    EXPECT(checkFlag(), true);
    // await repeat delay
    std::this_thread::sleep_for(std::chrono::milliseconds(200));
    EXPECT(checkFlag(), true);
    // check that it continues repeating
    std::this_thread::sleep_for(std::chrono::milliseconds(200));
    EXPECT(checkFlag(), true);
    // release keybind
    OK(getFromSocket("/dispatch plugin:test:keybind 0,0,29"));
    EXPECT(getFromSocket("/keyword unbind SUPER,Y"), "ok");
}

static void testKeyRepeat() {
    EXPECT(checkFlag(), false);
    EXPECT(getFromSocket("/keyword binde ,Y,exec,touch " + flagFile), "ok");
    EXPECT(getFromSocket("/keyword input:repeat_delay 100"), "ok");
    // press keybind
    OK(getFromSocket("/dispatch plugin:test:keybind 1,0,29"));
    // await flag
    std::this_thread::sleep_for(std::chrono::milliseconds(50));
    EXPECT(checkFlag(), true);
    // await repeat delay
    std::this_thread::sleep_for(std::chrono::milliseconds(200));
    EXPECT(checkFlag(), true);
    // check that it continues repeating
    std::this_thread::sleep_for(std::chrono::milliseconds(200));
    EXPECT(checkFlag(), true);
    // release keybind
    OK(getFromSocket("/dispatch plugin:test:keybind 0,0,29"));
    EXPECT(getFromSocket("/keyword unbind ,Y"), "ok");
}

static void testRepeatRelease() {
    EXPECT(checkFlag(), false);
    EXPECT(getFromSocket("/keyword binde SUPER,Y,exec,touch " + flagFile), "ok");
    EXPECT(getFromSocket("/keyword input:repeat_delay 100"), "ok");
    // press keybind
    OK(getFromSocket("/dispatch plugin:test:keybind 1,7,29"));
    // await flag
    std::this_thread::sleep_for(std::chrono::milliseconds(50));
    EXPECT(checkFlag(), true);
    // release keybind
    OK(getFromSocket("/dispatch plugin:test:keybind 0,0,29"));
    // await repeat delay
    std::this_thread::sleep_for(std::chrono::milliseconds(200));
    clearFlag();
    std::this_thread::sleep_for(std::chrono::milliseconds(200));
    EXPECT(checkFlag(), false);
    // check that it is not repeating
    std::this_thread::sleep_for(std::chrono::milliseconds(200));
    EXPECT(checkFlag(), false);
    EXPECT(getFromSocket("/keyword unbind SUPER,Y"), "ok");
}

static void testRepeatOnlyKeyRelease() {
    EXPECT(checkFlag(), false);
    EXPECT(getFromSocket("/keyword binde SUPER,Y,exec,touch " + flagFile), "ok");
    EXPECT(getFromSocket("/keyword input:repeat_delay 100"), "ok");
    // press keybind
    OK(getFromSocket("/dispatch plugin:test:keybind 1,7,29"));
    // await flag
    std::this_thread::sleep_for(std::chrono::milliseconds(200));
    EXPECT(checkFlag(), true);
    // release key, keep modifier
    OK(getFromSocket("/dispatch plugin:test:keybind 0,7,29"));
    // await repeat delay
    std::this_thread::sleep_for(std::chrono::milliseconds(200));
    clearFlag();
    std::this_thread::sleep_for(std::chrono::milliseconds(200));
    EXPECT(checkFlag(), false);
    // check that it is not repeating
    std::this_thread::sleep_for(std::chrono::milliseconds(200));
    EXPECT(checkFlag(), false);
    OK(getFromSocket("/dispatch plugin:test:keybind 0,0,29"));
    EXPECT(getFromSocket("/keyword unbind SUPER,Y"), "ok");
}

static void testShortcutBind() {
    auto kittyProc = spawnRemoteControlKitty();
    if (!kittyProc) {
        NLog::log("{}Error: kitty did not spawn", Colors::RED);
        ret = 1;
        return;
    }
    EXPECT(getFromSocket("/dispatch focuswindow class:keybinds_test"), "ok");
    EXPECT(getFromSocket("/keyword bind SUPER,Y,sendshortcut,,q,"), "ok");
    // press keybind
    OK(getFromSocket("/dispatch plugin:test:keybind 1,7,29"));
    // release keybind
    std::this_thread::sleep_for(std::chrono::milliseconds(50));
    OK(getFromSocket("/dispatch plugin:test:keybind 0,0,29"));
    std::this_thread::sleep_for(std::chrono::milliseconds(50));
    const std::string output = readKittyOutput();
    EXPECT_COUNT_STRING(output, "y", 0);
    EXPECT_COUNT_STRING(output, "q", 1);
    EXPECT(getFromSocket("/keyword unbind SUPER,Y"), "ok");
    Tests::killAllWindows();
}

static void testShortcutBindKey() {
    auto kittyProc = spawnRemoteControlKitty();
    if (!kittyProc) {
        NLog::log("{}Error: kitty did not spawn", Colors::RED);
        ret = 1;
        return;
    }
    EXPECT(getFromSocket("/dispatch focuswindow class:keybinds_test"), "ok");
    EXPECT(getFromSocket("/keyword bind ,Y,sendshortcut,,q,"), "ok");
    // press keybind
    OK(getFromSocket("/dispatch plugin:test:keybind 1,0,29"));
    // release keybind
    std::this_thread::sleep_for(std::chrono::milliseconds(50));
    OK(getFromSocket("/dispatch plugin:test:keybind 0,0,29"));
    std::this_thread::sleep_for(std::chrono::milliseconds(50));
    const std::string output = readKittyOutput();
    EXPECT_COUNT_STRING(output, "y", 0);
    // disabled: doesn't work in CI
    // EXPECT_COUNT_STRING(output, "q", 1);
    EXPECT(getFromSocket("/keyword unbind ,Y"), "ok");
    Tests::killAllWindows();
}

static void testShortcutLongPress() {
    auto kittyProc = spawnRemoteControlKitty();
    if (!kittyProc) {
        NLog::log("{}Error: kitty did not spawn", Colors::RED);
        ret = 1;
        return;
    }
    EXPECT(getFromSocket("/dispatch focuswindow class:keybinds_test"), "ok");
    EXPECT(getFromSocket("/keyword bindo SUPER,Y,sendshortcut,,q,"), "ok");
    EXPECT(getFromSocket("/keyword input:repeat_delay 100"), "ok");
    EXPECT(getFromSocket("/keyword input:repeat_rate 10"), "ok");
    // press keybind
    OK(getFromSocket("/dispatch plugin:test:keybind 1,7,29"));
    // await repeat delay
    std::this_thread::sleep_for(std::chrono::milliseconds(200));
    OK(getFromSocket("/dispatch plugin:test:keybind 0,0,29"));
    std::this_thread::sleep_for(std::chrono::milliseconds(200));
    const std::string output = readKittyOutput();
    int               yCount = Tests::countOccurrences(output, "y");
    // sometimes 1, sometimes 2, not sure why
    // keybind press sends 1 y immediately
    // then repeat triggers, sending 1 y
    // final release stop repeats, and shouldn't send any more
    EXPECT(true, yCount == 1 || yCount == 2);
    EXPECT_COUNT_STRING(output, "q", 1);
    EXPECT(getFromSocket("/keyword unbind SUPER,Y"), "ok");
    Tests::killAllWindows();
}

static void testShortcutLongPressKeyRelease() {
    auto kittyProc = spawnRemoteControlKitty();
    if (!kittyProc) {
        NLog::log("{}Error: kitty did not spawn", Colors::RED);
        ret = 1;
        return;
    }
    EXPECT(getFromSocket("/dispatch focuswindow class:keybinds_test"), "ok");
    EXPECT(getFromSocket("/keyword bindo SUPER,Y,sendshortcut,,q,"), "ok");
    EXPECT(getFromSocket("/keyword input:repeat_delay 100"), "ok");
    EXPECT(getFromSocket("/keyword input:repeat_rate 10"), "ok");
    // press keybind
    OK(getFromSocket("/dispatch plugin:test:keybind 1,7,29"));
    std::this_thread::sleep_for(std::chrono::milliseconds(10));
    // release key, keep modifier
    OK(getFromSocket("/dispatch plugin:test:keybind 0,7,29"));
    // await repeat delay
    std::this_thread::sleep_for(std::chrono::milliseconds(200));
    const std::string output = readKittyOutput();
    // disabled: doesn't work on CI
    // EXPECT_COUNT_STRING(output, "y", 1);
    EXPECT_COUNT_STRING(output, "q", 0);
    OK(getFromSocket("/dispatch plugin:test:keybind 0,0,29"));
    EXPECT(getFromSocket("/keyword unbind SUPER,Y"), "ok");
    Tests::killAllWindows();
}

static void testShortcutRepeat() {
    auto kittyProc = spawnRemoteControlKitty();
    if (!kittyProc) {
        NLog::log("{}Error: kitty did not spawn", Colors::RED);
        ret = 1;
        return;
    }
    EXPECT(getFromSocket("/dispatch focuswindow class:keybinds_test"), "ok");
    EXPECT(getFromSocket("/keyword binde SUPER,Y,sendshortcut,,q,"), "ok");
    EXPECT(getFromSocket("/keyword input:repeat_rate 5"), "ok");
    EXPECT(getFromSocket("/keyword input:repeat_delay 200"), "ok");
    // press keybind
    OK(getFromSocket("/dispatch plugin:test:keybind 1,7,29"));
    // await repeat
    std::this_thread::sleep_for(std::chrono::milliseconds(210));
    // release keybind
    OK(getFromSocket("/dispatch plugin:test:keybind 0,0,29"));
    std::this_thread::sleep_for(std::chrono::milliseconds(450));
    const std::string output = readKittyOutput();
    EXPECT_COUNT_STRING(output, "y", 0);
    int qCount = Tests::countOccurrences(output, "q");
    // sometimes 2, sometimes 3, not sure why
    // keybind press sends 1 q immediately
    // then repeat triggers, sending 1 q
    // final release stop repeats, and shouldn't send any more
    EXPECT(true, qCount == 2 || qCount == 3);
    EXPECT(getFromSocket("/keyword unbind SUPER,Y"), "ok");
    Tests::killAllWindows();
}

static void testShortcutRepeatKeyRelease() {
    auto kittyProc = spawnRemoteControlKitty();
    if (!kittyProc) {
        NLog::log("{}Error: kitty did not spawn", Colors::RED);
        ret = 1;
        return;
    }
    EXPECT(getFromSocket("/dispatch focuswindow class:keybinds_test"), "ok");
    EXPECT(getFromSocket("/keyword binde SUPER,Y,sendshortcut,,q,"), "ok");
    EXPECT(getFromSocket("/keyword input:repeat_rate 5"), "ok");
    EXPECT(getFromSocket("/keyword input:repeat_delay 200"), "ok");
    // press keybind
    OK(getFromSocket("/dispatch plugin:test:keybind 1,7,29"));
    std::this_thread::sleep_for(std::chrono::milliseconds(210));
    // release key, keep modifier
    OK(getFromSocket("/dispatch plugin:test:keybind 0,7,29"));
    // if repeat was still active, we'd get 2 more q's here
    std::this_thread::sleep_for(std::chrono::milliseconds(450));
    // release modifier
    const std::string output = readKittyOutput();
    EXPECT_COUNT_STRING(output, "y", 0);
    int qCount = Tests::countOccurrences(output, "q");
    // sometimes 2, sometimes 3, not sure why
    // keybind press sends 1 q immediately
    // then repeat triggers, sending 1 q
    // final release stop repeats, and shouldn't send any more
    EXPECT(true, qCount == 2 || qCount == 3);
    OK(getFromSocket("/dispatch plugin:test:keybind 0,0,29"));
    EXPECT(getFromSocket("/keyword unbind SUPER,Y"), "ok");
    Tests::killAllWindows();
}

static void testSubmap() {
    const auto press = [](const uint32_t key, const uint32_t mod = 0) {
        // +8 because udev -> XKB keycode.
        getFromSocket("/dispatch plugin:test:keybind 1," + std::to_string(mod) + "," + std::to_string(key + 8));
        getFromSocket("/dispatch plugin:test:keybind 0," + std::to_string(mod) + "," + std::to_string(key + 8));
    };

    NLog::log("{}Testing submaps", Colors::GREEN);
    // submap 1 no resets
    press(KEY_U, MOD_META);
    EXPECT_CONTAINS(getFromSocket("/submap"), "submap1");
    press(KEY_O);
    Tests::waitUntilWindowsN(1);
    EXPECT_CONTAINS(getFromSocket("/submap"), "submap1");
    // submap 2 resets to submap 1
    press(KEY_U);
    EXPECT_CONTAINS(getFromSocket("/submap"), "submap2");
    press(KEY_O);
    Tests::waitUntilWindowsN(2);
    EXPECT_CONTAINS(getFromSocket("/submap"), "submap1");
    // submap 3 resets to default
    press(KEY_I);
    EXPECT_CONTAINS(getFromSocket("/submap"), "submap3");
    press(KEY_O);
    Tests::waitUntilWindowsN(3);
    EXPECT_CONTAINS(getFromSocket("/submap"), "default");
    // submap 1 reset via keybind
    press(KEY_U, MOD_META);
    EXPECT_CONTAINS(getFromSocket("/submap"), "submap1");
    press(KEY_P);
    EXPECT_CONTAINS(getFromSocket("/submap"), "default");

    Tests::killAllWindows();
}

static void testBindsAfterScroll() {
    NLog::log("{}Testing binds after scroll", Colors::GREEN);

    clearFlag();
    OK(getFromSocket("/keyword binds Alt_R,w,exec,touch " + flagFile));

    // press keybind before scroll
    OK(getFromSocket("/dispatch plugin:test:keybind 1,0,108")); // Alt_R press
    OK(getFromSocket("/dispatch plugin:test:keybind 1,4,25"));  // w press
    EXPECT(attemptCheckFlag(20, 50), true);
    OK(getFromSocket("/dispatch plugin:test:keybind 0,4,25"));  // w release
    OK(getFromSocket("/dispatch plugin:test:keybind 0,0,108")); // Alt_R release

    // scroll
    OK(getFromSocket("/dispatch plugin:test:scroll 120"));
    OK(getFromSocket("/dispatch plugin:test:scroll -120"));
    OK(getFromSocket("/dispatch plugin:test:scroll 120"));

    // press keybind after scroll
    OK(getFromSocket("/dispatch plugin:test:keybind 1,0,108")); // Alt_R press
    OK(getFromSocket("/dispatch plugin:test:keybind 1,4,25"));  // w press
    EXPECT(attemptCheckFlag(20, 50), true);
    OK(getFromSocket("/dispatch plugin:test:keybind 0,4,25"));  // w release
    OK(getFromSocket("/dispatch plugin:test:keybind 0,0,108")); // Alt_R release

    clearFlag();
    OK(getFromSocket("/keyword unbind Alt_R,w"));
}

static void testSubmapUniversal() {
    NLog::log("{}Testing submap universal", Colors::GREEN);

    EXPECT(checkFlag(), false);
    EXPECT(getFromSocket("/keyword bindu SUPER,Y,exec,touch " + flagFile), "ok");
    EXPECT_CONTAINS(getFromSocket("/submap"), "default");

    // keybind works on default submap
    OK(getFromSocket("/dispatch plugin:test:keybind 1,7,29"));
    OK(getFromSocket("/dispatch plugin:test:keybind 0,7,29"));
    EXPECT(attemptCheckFlag(30, 5), true);

    // keybind works on submap1
    getFromSocket("/dispatch plugin:test:keybind 1,7,30");
    getFromSocket("/dispatch plugin:test:keybind 0,7,30");
    EXPECT_CONTAINS(getFromSocket("/submap"), "submap1");
    OK(getFromSocket("/dispatch plugin:test:keybind 1,7,29"));
    OK(getFromSocket("/dispatch plugin:test:keybind 0,7,29"));
    EXPECT(attemptCheckFlag(30, 5), true);

    // reset to default submap
    getFromSocket("/dispatch plugin:test:keybind 1,0,33");
    getFromSocket("/dispatch plugin:test:keybind 0,0,33");
    EXPECT_CONTAINS(getFromSocket("/submap"), "default");

    EXPECT(getFromSocket("/keyword unbind SUPER,Y"), "ok");
}

static bool test() {
    NLog::log("{}Testing keybinds", Colors::GREEN);

    clearFlag();

    testBind();
    testBindKey();
    testLongPress();
    testKeyLongPress();
    testLongPressRelease();
    testLongPressOnlyKeyRelease();
    testRepeat();
    testKeyRepeat();
    testRepeatRelease();
    testRepeatOnlyKeyRelease();
    testShortcutBind();
    testShortcutBindKey();
    testShortcutLongPress();
    testShortcutLongPressKeyRelease();
    testShortcutRepeat();
    testShortcutRepeatKeyRelease();
    testSubmap();
    testSubmapUniversal();
    testBindsAfterScroll();

    clearFlag();
    return !ret;
}

REGISTER_TEST_FN(test)
