#pragma once

#include "USB.h"
#include "USBHIDKeyboard.h"
#include "USBHIDMouse.h"
#include "USBHIDConsumerControl.h"
#include "abs_mouse.h"
#include "config.h"

// USB HID timing floors. The underlying SendReport call is blocking — it
// returns only after the host has acknowledged the report — so most "wait
// for the bytes to arrive" timing concerns are already covered by the
// library. These values are the EXTRA delay we hold after a press or
// between events so the *host application* has a chance to observe and
// process each keystroke (BIOS prompts, installer wizards, and PE shells
// can miss reports that flip in too quickly even after the USB stack has
// delivered them).
//
//   USB_HID_HOLD_MIN_MS — minimum keydown hold time. Floor for the
//     between-press-and-release delay even if the user's typeDelay
//     setting is smaller. 8 ms = one boot-keyboard poll interval.
//   USB_HID_INTER_CHAR_MS — gap between releasing one char's keys and
//     pressing the next char's. Keeps key-repeat detection happy and
//     stops fast-typing apps from coalescing two characters into one.
//   USB_HID_PROLOGUE_SETTLE_MS — delay after the defensive entry
//     releaseAll() before the first character's press(). Must be long
//     enough that the releaseAll's xfer-complete callback has fired,
//     otherwise the TinyUSB SendReport semaphore can desync on the
//     first character (see press/release retry below).
//   USB_HID_RETRY_SETTLE_MS — pause before retrying a failed press or
//     release report. Gives the endpoint FIFO time to drain.
//   USB_HID_COMBO_MOD_SETTLE_MS — extra dwell, on top of the per-mod
//     preDelay, between pressing the last modifier of a combo and
//     pressing the main key. Some hosts (notably Windows shell hotkey
//     handlers and BIOS UIs) need a clear "modifier is steady-state
//     held" window before the keycode arrives or they treat the combo
//     as a plain keypress without the modifier. 25 ms is empirically
//     enough on every host we've tested without being noticeable to a
//     human watching the combo fire.
#define USB_HID_HOLD_MIN_MS         8
#define USB_HID_INTER_CHAR_MS       5
#define USB_HID_PROLOGUE_SETTLE_MS  10
#define USB_HID_RETRY_SETTLE_MS     2
#define USB_HID_COMBO_MOD_SETTLE_MS 25

class HIDController {
public:
    USBHIDKeyboard keyboard;
    USBHIDMouse mouse;
    USBHIDConsumerControl consumer;
    AbsoluteMouse absMouse;   // absolute-position pointer for the BT Keyboard trackpad

    // Keyboard LED state (updated via host reports)
    volatile bool numLockOn = false;
    volatile bool capsLockOn = false;
    volatile bool scrollLockOn = false;
    volatile bool ledStateReceived = false;  // true once we've received at least one LED report

    void begin() {
        USB.productName("ATOMS3 MacroPad");
        USB.manufacturerName("M5Stack");

        // Register LED event callback BEFORE begin() so we don't miss events
        _instance = this;
        keyboard.onEvent(_keyboardEventCB);

        // Register all HID interfaces BEFORE USB.begin()
        keyboard.begin();
        mouse.begin();
        consumer.begin();
        absMouse.begin();

        // Leave shiftKeyReports at the library default (false). With it
        // ON, a shifted character emits FOUR reports — shift-down alone,
        // then shift+key, then shift-only on key-up, then shift-up.
        // Theoretically this matches a physical keyboard more closely
        // and is "what BIOSes expect," but in practice on a normal
        // Windows host the intermediate "shift-alone" report has been
        // observed to leave shift latched on across subsequent keys,
        // causing "Coconuts4frodo" to come out as "COCONUTS$FRODO".
        // The single-report form (shift+key bundled) is what every
        // tested host actually wants. If a future BIOS/PE target needs
        // the split form, enable per-target rather than globally.
        // keyboard.setShiftKeyReports(true);   // DO NOT enable globally.

        // Start the TinyUSB stack LAST - this finalizes all descriptors
        USB.begin();

        // Disable DTR/RTS triggered reboot AFTER USB stack is running
        USBSerial.enableReboot(false);
    }

    // True while the engine is in the middle of a sequence of HID writes
    // (typeText / keyCombo / mediaKey / macro playback / probeNumLock).
    // The main loop checks this and skips non-USB-HID polling work
    // (BLE log flushes, RS232 RX scraping) so the typing path runs as
    // uninterrupted as possible. Read-only from outside.
    bool isCritical() const { return _critical; }

    // External entry points for callers that emit HID reports directly
    // (e.g. ``macro`` node-type playback in the engine, which uses
    // pressRaw / releaseRaw against the keyboard object). Wrap the
    // sequence in beginCritical() / endCritical() so the main-loop
    // priority gate observes it the same as typeText / keyCombo.
    void beginCritical() { _critical = true; }
    void endCritical()   { _critical = false; }

    // Probe whether a host PC is alive by toggling Num Lock and checking
    // if the LED state changes.  Works regardless of the initial Num Lock
    // state.  Restores the original state if the host is alive.
    //
    // Returns true if the host responded (PC is alive), false otherwise.
    //
    // Algorithm:
    //   1. Read current Num Lock LED state  (before)
    //   2. Send Num Lock keypress           (toggle)
    //   3. Wait for host to report new LED  (after)
    //   4. Compare before vs after
    //      - Changed  → host is alive → toggle back to restore → return true
    //      - Same     → host is dead / not connected         → return false
    //
    // stepCallback is called at each phase so the display can show progress.
    typedef void (*ProbeStepCB)(const char* phase, void* userData);

    bool probeNumLock(uint16_t waitMs = 250, ProbeStepCB stepCB = nullptr, void* cbData = nullptr) {
        _critical = true;
        struct CriticalGuard {
            HIDController* h;
            ~CriticalGuard() { h->_critical = false; }
        } guard{this};

        // Step 1: Record the "before" state
        if (stepCB) stepCB("Read state...", cbData);
        bool before = numLockOn;
        delay(25);

        // Step 2: Toggle Num Lock
        if (stepCB) stepCB("Toggling...", cbData);
        ledStateReceived = false;
        keyboard.press(KEY_NUM_LOCK);
        delay(25);
        keyboard.releaseAll();
        delay(25);

        // Step 3: Wait for the host to send an LED report
        if (stepCB) stepCB("Waiting for host...", cbData);
        uint32_t deadline = millis() + waitMs;
        while (!ledStateReceived && millis() < deadline) {
            delay(5);
        }
        delay(25);

        bool after = numLockOn;

        // Step 4: Compare — if LED changed, host is alive; restore original state
        if (ledStateReceived && after != before) {
            if (stepCB) stepCB("Host alive! Restoring...", cbData);
            ledStateReceived = false;
            keyboard.press(KEY_NUM_LOCK);
            delay(25);
            keyboard.releaseAll();
            delay(25);

            // Wait for restore to register
            deadline = millis() + waitMs;
            while (!ledStateReceived && millis() < deadline) {
                delay(5);
            }
            delay(25);
            return true;
        }

        if (stepCB) stepCB("No response", cbData);
        delay(25);
        return false;
    }

    // Count the number of Scroll Lock LED transitions observed within a
    // listening window. The Get Variables node uses this to read a simple
    // signal channel from a host-side script (e.g. PowerShell calling
    // user32!keybd_event with VK_SCROLL). Scroll Lock is preferred over
    // Num Lock because most users have Num Lock toggling externally
    // (numeric keypads, BIOS settings) which would corrupt the count,
    // whereas Scroll Lock is virtually never touched by other software.
    //
    // We count BOTH directions (off->on AND on->off) so each press the
    // host script issues maps 1:1 to one count, regardless of starting
    // state. Earlier versions only counted off->on, which made the second
    // press of a "press twice" sequence invisible.
    int probeScrollLockSequence(uint32_t windowMs) {
        int count = 0;
        bool prev = scrollLockOn;
        uint32_t deadline = millis() + windowMs;
        while (millis() < deadline) {
            bool now = scrollLockOn;
            if (now != prev) {
                count++;
                prev = now;
            }
            delay(5);
        }
        return count;
    }

    // Callback type: called before each character is typed.
    // Arguments: (fullText, charIndex, userData)
    typedef void (*CharCallback)(const char* fullText, int charIdx, void* userData);

    // Type a string one character at a time over the HID keyboard interface.
    //
    // Per-char timeline (with setShiftKeyReports(true) — see begin()):
    //   - keyboard.press(c)
    //       Unshifted: 1 report (key down). Blocks until host ack (~1-5 ms).
    //       Shifted:   2 reports (shift down, then key down). Blocks ~5-10 ms.
    //   - delay(holdMs)
    //       Host-side keydown processing time. Floored at USB_HID_HOLD_MIN_MS
    //       so very small typeDelay values don't starve apps that need to
    //       observe the keydown for at least one poll cycle.
    //   - keyboard.release(c)
    //       Unshifted: 1 report. Shifted: 2 reports (key up, then shift up).
    //   - delay(interCharMs)
    //       Host-side keyup processing + breathing room before the next
    //       keydown. Without this gap, fast key-repeat detection in some
    //       editors can elide every other keystroke.
    //   - For shifted chars, ``shiftExtraMs`` is added to the inter-char gap
    //       (NOT the hold) because the cost is on the release/next-press
    //       boundary: 4 reports must drain before the next char's shift
    //       state diverges.
    //
    // The original implementation surrounded each ``write()`` with two
    // ``releaseAll()`` calls — that's 4 extra HID reports per character
    // (~20 ms wasted) plus 15 ms of explicit delay. Functionally identical
    // to plain press/release because press() / release() are paired
    // already, so the releases were redundant. Removed here.
    //
    // The trailing settle is a final ``releaseAll()`` followed by
    // ``settleMs`` of quiet so any in-flight report fully drains before
    // the next macro node fires.
    void typeText(const char* text, uint8_t delayMs,
                  CharCallback onChar = nullptr, void* userData = nullptr,
                  uint16_t shiftExtraMs = 25, uint16_t settleMs = 150,
                  uint16_t holdMinMs = USB_HID_HOLD_MIN_MS,
                  uint16_t interCharMs = USB_HID_INTER_CHAR_MS) {
        _critical = true;
        int count = 0;
        const char* start = text;  // Keep pointer to full string for callback

        // Defensive prologue: clear any modifier state left held by a
        // preceding combo / macro-playback / failed earlier typeText.
        // Without this, a stuck shift bit from the previous node would
        // turn the whole text into shifted equivalents ("Coconuts4frodo"
        // → "COCONUTS$FRODO"). The cost is one report (~1-5 ms host
        // ack), cheap insurance.
        //
        // The post-prologue settle is bumped above interCharMs so the
        // releaseAll's xfer-complete callback has definitively fired
        // before the first character's press(). On TinyUSB / arduino-esp32
        // builds with the known semaphore-desync bug, a too-tight gap
        // here can cause the FIRST character of every typeText call to
        // drop silently.
        keyboard.releaseAll();
        uint16_t prologueSettle = interCharMs;
        if (prologueSettle < USB_HID_PROLOGUE_SETTLE_MS) prologueSettle = USB_HID_PROLOGUE_SETTLE_MS;
        delay(prologueSettle);

        // Hold time floor: ensure at least one USB poll cycle elapses with
        // the key down so the host always observes the press. Below this
        // floor, fast-typing apps can drop characters.
        uint16_t holdMs = delayMs;
        if (holdMs < holdMinMs) holdMs = holdMinMs;

        while (*text) {
            if (onChar) {
                onChar(start, count, userData);
            }

            char c = *text;
            bool shifted = _isShiftedChar(c);

            // press()/release() bundle shift+key into a single HID report
            // (since shiftKeyReports is left at the default `false`).
            // For an unshifted char: 1 report on press, 1 on release.
            // For a shifted char ('A', '!', etc.): same 2 reports, with
            // the shift modifier bit set in the press and cleared in the
            // release.
            //
            // SendReport blocks on a semaphore that is given by
            // tud_hid_report_complete_cb, so the happy path returns only
            // after the host has drained the endpoint. The known
            // arduino-esp32 / TinyUSB semaphore-desync bug occasionally
            // makes that semaphore-take time out, returning 0 from
            // press() / release() without the report ever reaching the
            // host. Without the retry, that one character drops
            // silently — the symptom is rare random misses like
            // "Start-Process" → "Start-Proess".
            _hidWriteWithRetry(true, (uint8_t)c, start, count);
            delay(holdMs);
            _hidWriteWithRetry(false, (uint8_t)c, start, count);

            // Inter-char gap. Shifted chars get an extra slice so the
            // shift bit has clearly cleared at the host before the next
            // unshifted character's keycode arrives — a few hosts have
            // been observed to apply a still-cached shift state to the
            // very next report.
            uint16_t gap = interCharMs;
            if (shifted && shiftExtraMs > 0) gap += shiftExtraMs;
            delay(gap);

            text++;
            count++;
        }

        // Final safety release in case something above failed mid-sequence
        // and left a key latched, then the configured settle.
        keyboard.releaseAll();
        delay(settleMs);
        _critical = false;
    }

    static bool _isShiftedChar(char c) {
        if (c >= 'A' && c <= 'Z') return true;
        return c != 0 && strchr("!@#$%^&*()_+{}|:\"<>?~", c) != nullptr;
    }

    // Single press or release with one-shot retry. Returns true if the
    // report was acknowledged by the host on either the first try or
    // the retry. A logged failure means the character was lost — emit
    // a Serial line so the user can confirm in the field whether the
    // semaphore-desync failure mode actually fires for their hardware.
    //
    // ``isPress`` selects press vs release. ``c`` is the key. ``textCtx``
    // and ``idx`` are only used for the log line.
    bool _hidWriteWithRetry(bool isPress, uint8_t c, const char* textCtx, int idx) {
        size_t ok = isPress ? keyboard.press(c) : keyboard.release(c);
        if (ok) return true;

        // First attempt failed (semaphore-take timeout / FIFO not ready).
        // Brief settle then one retry.
        delay(USB_HID_RETRY_SETTLE_MS);
        ok = isPress ? keyboard.press(c) : keyboard.release(c);
        if (ok) {
            Serial.printf("[HID retry] %s '%c' idx=%d ok on retry\n",
                          isPress ? "press" : "release", (char)c, idx);
            return true;
        }

        Serial.printf("[HID retry] %s '%c' idx=%d FAILED twice — char dropped\n",
                      isPress ? "press" : "release", (char)c, idx);
        return false;
    }

    // Press a key combination (zero or more modifiers + zero or one main key).
    //
    // Timeline:
    //   1. Press each modifier, with ``preDelay`` after each press. The
    //      ``preDelay`` after the LAST modifier doubles as the
    //      modifier-to-key spacer — the host has one ack-plus-preDelay
    //      worth of time to observe each new modifier bit settled before
    //      anything else happens.
    //   2. When ``modCount > 0`` (combo has at least one modifier),
    //      apply an additional ``USB_HID_COMBO_MOD_SETTLE_MS`` (25 ms)
    //      dwell so the modifier is steady-state held for clearly more
    //      than one host poll cycle before the keycode arrives. Without
    //      this, fast-firing combos like Ctrl+R can race the host's
    //      modifier-state pipeline and register as a bare 'R' instead.
    //      Skipped for plain (modifier-less) key presses since there's
    //      nothing to settle.
    //   3. Press the main key.
    //   4. Hold the combo for ``postDelay`` so the host registers the
    //      shortcut as a tap, not a coalesced flicker.
    //   5. releaseAll() — single report clears everything. Floor of
    //      USB_HID_INTER_CHAR_MS afterwards so the next macro node doesn't
    //      race against the just-issued release.
    void keyCombo(const uint8_t* modifiers, uint8_t modCount, uint8_t key,
                  uint16_t preDelay = COMBO_KEY_PRE_DELAY,
                  uint16_t postDelay = COMBO_KEY_POST_DELAY) {
        _critical = true;
        // Caller controls the timings. Engine clamps to a 1ms floor before
        // calling us; we trust the inputs here.
        for (uint8_t i = 0; i < modCount; i++) {
            keyboard.press(modifiers[i]);
            delay(preDelay);
        }
        if (key != 0) {
            if (modCount > 0) {
                // Extra modifier-settle window so the host clearly sees
                // the modifier bits as held BEFORE the keycode flips on.
                delay(USB_HID_COMBO_MOD_SETTLE_MS);
            }
            keyboard.press(key);
        }
        delay(postDelay);
        keyboard.releaseAll();
        // Small trailing gap so the next node (often another combo or a
        // text node) doesn't race the just-emitted release report.
        delay(USB_HID_INTER_CHAR_MS);
        _critical = false;
    }

    // Emit one absolute-position pointer report (BT Keyboard trackpad).
    // Wrapped in the critical guard like all other live HID emission so the
    // main-loop priority gate defers other work while it's in flight.
    void absMouseReport(uint8_t buttons, uint16_t x, uint16_t y, int8_t wheel) {
        _critical = true;
        absMouse.report(buttons, x, y, wheel);
        _critical = false;
    }

    void mouseClick(uint8_t button) {
        mouse.click(button);
    }

    void mouseDoubleClick(uint8_t button) {
        mouse.click(button);
        delay(80);
        mouse.click(button);
    }

    void mousePress(uint8_t button) {
        mouse.press(button);
    }

    void mouseRelease(uint8_t button) {
        mouse.release(button);
    }

    void mediaKey(uint16_t key, uint16_t holdMs = 100) {
        _critical = true;
        consumer.press(key);
        delay(holdMs);
        consumer.release();
        _critical = false;
    }

    uint8_t resolveModifier(const char* mod) {
        if (strcmp(mod, "ctrl") == 0 || strcmp(mod, "control") == 0) return KEY_LEFT_CTRL;
        if (strcmp(mod, "shift") == 0) return KEY_LEFT_SHIFT;
        if (strcmp(mod, "alt") == 0) return KEY_LEFT_ALT;
        if (strcmp(mod, "gui") == 0 || strcmp(mod, "win") == 0 || strcmp(mod, "meta") == 0) return KEY_LEFT_GUI;
        if (strcmp(mod, "rctrl") == 0) return KEY_RIGHT_CTRL;
        if (strcmp(mod, "rshift") == 0) return KEY_RIGHT_SHIFT;
        if (strcmp(mod, "ralt") == 0 || strcmp(mod, "altgr") == 0) return KEY_RIGHT_ALT;
        if (strcmp(mod, "rgui") == 0) return KEY_RIGHT_GUI;
        return 0;
    }

    uint8_t resolveKey(const char* key) {
        if (strlen(key) == 1) return (uint8_t)key[0];
        if (strcmp(key, "enter") == 0 || strcmp(key, "return") == 0) return KEY_RETURN;
        if (strcmp(key, "esc") == 0 || strcmp(key, "escape") == 0) return KEY_ESC;
        if (strcmp(key, "backspace") == 0) return KEY_BACKSPACE;
        if (strcmp(key, "tab") == 0) return KEY_TAB;
        if (strcmp(key, "space") == 0) return KEY_SPACE;
        if (strcmp(key, "delete") == 0) return KEY_DELETE;
        if (strcmp(key, "insert") == 0) return KEY_INSERT;
        if (strcmp(key, "home") == 0) return KEY_HOME;
        if (strcmp(key, "end") == 0) return KEY_END;
        if (strcmp(key, "pageup") == 0) return KEY_PAGE_UP;
        if (strcmp(key, "pagedown") == 0) return KEY_PAGE_DOWN;
        if (strcmp(key, "up") == 0) return KEY_UP_ARROW;
        if (strcmp(key, "down") == 0) return KEY_DOWN_ARROW;
        if (strcmp(key, "left") == 0) return KEY_LEFT_ARROW;
        if (strcmp(key, "right") == 0) return KEY_RIGHT_ARROW;
        if (strcmp(key, "capslock") == 0) return KEY_CAPS_LOCK;
        if (strcmp(key, "numlock") == 0) return KEY_NUM_LOCK;
        if (strcmp(key, "scrolllock") == 0) return KEY_SCROLL_LOCK;
        if (strcmp(key, "printscreen") == 0) return KEY_PRINT_SCREEN;
        if (strcmp(key, "pause") == 0) return KEY_PAUSE;
        if (strcmp(key, "menu") == 0) return KEY_MENU;
        if (key[0] == 'f' || key[0] == 'F') {
            int num = atoi(key + 1);
            if (num >= 1 && num <= 12) return KEY_F1 + (num - 1);
            if (num >= 13 && num <= 24) return KEY_F13 + (num - 13);
        }
        return 0;
    }

    uint8_t resolveMouseButton(const char* btn) {
        if (strcmp(btn, "left") == 0) return MOUSE_LEFT;
        if (strcmp(btn, "right") == 0) return MOUSE_RIGHT;
        if (strcmp(btn, "middle") == 0) return MOUSE_MIDDLE;
        return MOUSE_LEFT;
    }

    uint16_t resolveMediaKey(const char* action) {
        if (strcmp(action, "vol_up") == 0) return CONSUMER_CONTROL_VOLUME_INCREMENT;
        if (strcmp(action, "vol_down") == 0) return CONSUMER_CONTROL_VOLUME_DECREMENT;
        if (strcmp(action, "mute") == 0) return CONSUMER_CONTROL_MUTE;
        if (strcmp(action, "play_pause") == 0) return CONSUMER_CONTROL_PLAY_PAUSE;
        if (strcmp(action, "next") == 0) return CONSUMER_CONTROL_SCAN_NEXT;
        if (strcmp(action, "prev") == 0) return CONSUMER_CONTROL_SCAN_PREVIOUS;
        if (strcmp(action, "stop") == 0) return CONSUMER_CONTROL_STOP;
        if (strcmp(action, "brightness_up") == 0) return CONSUMER_CONTROL_BRIGHTNESS_INCREMENT;
        if (strcmp(action, "brightness_down") == 0) return CONSUMER_CONTROL_BRIGHTNESS_DECREMENT;
        return 0;
    }

private:
    static HIDController* _instance;

    // Set true while a HID-emitting method (typeText, keyCombo, mediaKey,
    // probeNumLock) is in progress. The main loop reads via isCritical()
    // and defers non-USB-HID polling work (BLE log flush, RS232 RX
    // buffering, M5 button events) for the duration. In practice the
    // main loop is already blocked inside engine.tick() while these
    // methods run, so the flag mainly serves to document the intent and
    // protect any future caller that might pump the main loop from a
    // nested context.
    volatile bool _critical = false;

    static void _keyboardEventCB(void* arg, esp_event_base_t event_base,
                                  int32_t event_id, void* event_data) {
        if (!_instance) return;
        if (event_base == ARDUINO_USB_HID_KEYBOARD_EVENTS &&
            event_id == ARDUINO_USB_HID_KEYBOARD_LED_EVENT) {
            arduino_usb_hid_keyboard_event_data_t* data =
                (arduino_usb_hid_keyboard_event_data_t*)event_data;
            // Log every LED report so we can verify whether a host-side
            // toggle actually propagated to this USB HID device. If the
            // host's keybd_event(VK_SCROLL,...) updates the OS state but
            // doesn't trigger a Set Report to the keyboard, we'd see no
            // log lines here and that's the signaling channel's bug.
            bool prevN = _instance->numLockOn;
            bool prevC = _instance->capsLockOn;
            bool prevS = _instance->scrollLockOn;
            _instance->numLockOn = data->numlock;
            _instance->capsLockOn = data->capslock;
            _instance->scrollLockOn = data->scrolllock;
            _instance->ledStateReceived = true;
            if (prevN != data->numlock || prevC != data->capslock ||
                prevS != data->scrolllock) {
                Serial.printf("[LED] num=%d caps=%d scroll=%d (was %d/%d/%d)\n",
                              (int)data->numlock, (int)data->capslock,
                              (int)data->scrolllock,
                              (int)prevN, (int)prevC, (int)prevS);
            }
        }
    }
};

// Static member definition
HIDController* HIDController::_instance = nullptr;