#pragma once

// LiveKeystrokeEngine — device-side queue and dispatcher for the
// "host streams keystrokes via BLE during macro recording" feature.
//
// Design rules:
//
//   1. BLE write callback runs on the NimBLE host task. It must NOT
//      touch USB HID directly — radio contention with TinyUSB causes
//      panics on ESP32-S3 (the very issue this whole subsystem is
//      designed around). The callback only enqueues parsed events.
//
//   2. The main loop calls drainQueue() each iteration. The HID
//      critical-section gate in usb_hid.h means BLE polling is
//      already deferred while keystrokes emit, so by the time we
//      get here it's safe to take over the USB radio.
//
//   3. pressRaw / releaseRaw are wrapped in beginCritical/endCritical
//      so the main loop's gating logic observes the in-flight HID
//      operation and defers BLE work — the existing protection
//      pattern, reused.
//
//   4. Queue is small and fixed-size; if full we set a flag the
//      BLE manager exposes as an ERROR notify back to the host.

#include <Arduino.h>
#include <freertos/FreeRTOS.h>
#include "usb_hid.h"

class LiveKeystrokeEngine {
public:
    static constexpr int QUEUE_CAP = 256;
    // Fixed device-side replay buffer. The host sends each event with a
    // host-monotonic timestamp; we emit at host_t + this offset so BLE
    // jitter is absorbed and the cadence the user typed on the host is
    // reproduced exactly on the target USB HID side.
    static constexpr uint32_t REPLAY_BUFFER_MS = 100;

    struct Event {
        uint8_t  action;        // 0 = down, 1 = up
        uint8_t  hid;           // raw USB HID usage code
        uint32_t scheduledMs;   // device millis() when this event should
                                // emit. Computed at enqueue time from the
                                // host's relative timestamp and the
                                // session anchor.
    };

    void begin(HIDController* hid) { _hid = hid; }

    // Start/stop are called from the BLE manager when the host sends
    // START / STOP control frames (or from MacroPad.ino on emergency
    // exits). idempotent.
    void start() {
        portENTER_CRITICAL(&_mux);
        _head = _tail = _count = 0;
        _overflow = false;
        _anchorSet = false;
        _active = true;
        portEXIT_CRITICAL(&_mux);
    }

    void stop() {
        portENTER_CRITICAL(&_mux);
        _active = false;
        _head = _tail = _count = 0;
        _anchorSet = false;
        // Release any held mouse buttons defensively — queue a buttons=0
        // report for the main loop to emit (don't touch USB HID here, this
        // can run on the NimBLE callback task).
        _mouseButtons = 0;
        _mouseWheel = 0;
        _mousePending = true;
        portEXIT_CRITICAL(&_mux);
        // releaseAll is best-effort defensive cleanup. Wrap in
        // critical so the BLE poll gate observes it even though we
        // ourselves are already on the main task here.
        if (_hid) {
            _hid->beginCritical();
            _hid->keyboard.releaseAll();
            _hid->endCritical();
        }
    }

    bool isActive() const {
        // Volatile read, no mutex — _active is set under mutex but
        // read-only here. Worst case: one extra loop iteration.
        return _active;
    }

    // Enqueue one event with the host's relative timestamp. Called from
    // the NimBLE host task — keep it fast. Returns false if the queue
    // is full (caller should ERROR-notify the host).
    //
    // We anchor on the first event of the session: _anchorMs becomes
    // the device-millis() value that corresponds to host_t = 0. Each
    // event's scheduledMs is then _anchorMs + host_t_ms + REPLAY_BUFFER_MS,
    // which preserves the host's typing cadence and adds a small jitter
    // buffer so events that arrive slightly out of cadence still emit
    // smoothly.
    bool enqueue(uint8_t action, uint8_t hid_code, uint32_t host_t_ms) {
        bool ok = false;
        portENTER_CRITICAL(&_mux);
        if (_active && _count < QUEUE_CAP) {
            if (!_anchorSet) {
                // First event of the session: pin the anchor so this
                // event's scheduled time is exactly now + buffer.
                _anchorMs = millis() - host_t_ms;
                _anchorSet = true;
            }
            _q[_tail].action = action;
            _q[_tail].hid = hid_code;
            _q[_tail].scheduledMs = _anchorMs + host_t_ms + REPLAY_BUFFER_MS;
            _tail = (_tail + 1) % QUEUE_CAP;
            _count++;
            ok = true;
        } else if (!_active) {
            // not live; signal NOT_LIVE_MODE upstream
        } else {
            _overflow = true;
        }
        portEXIT_CRITICAL(&_mux);
        return ok;
    }

    // Called from the main loop when !hid->isCritical(). Drains events
    // whose scheduled time has passed; events with future scheduledMs
    // stay in the queue so the host's typing cadence is preserved on
    // emission. Up to maxPerTick events per call so a flood doesn't
    // starve other main-loop work.
    //
    // We must not delay() here — the main loop owns timing for the
    // 3-second button-hold exit, button polling, BLE callback flush.
    // Anything not ready yet stays queued until the next loop iter.
    // Returns the number of events emitted (drives the activity flicker
    // on the screenless AtomS3 Lite).
    int drainQueue(int maxPerTick = 32) {
        if (!_hid || !_active) return 0;
        if (_count == 0) return 0;

        uint32_t now = millis();

        // Peek head first — if it's not ready, nothing else is either
        // (events are enqueued in monotonic schedule order).
        bool headReady = false;
        portENTER_CRITICAL(&_mux);
        if (_count > 0) {
            headReady = (int32_t)(_q[_head].scheduledMs - now) <= 0;
        }
        portEXIT_CRITICAL(&_mux);
        if (!headReady) return 0;

        _hid->beginCritical();
        int emitted = 0;
        while (emitted < maxPerTick) {
            Event ev;
            bool got = false;
            portENTER_CRITICAL(&_mux);
            if (_count > 0 && (int32_t)(_q[_head].scheduledMs - now) <= 0) {
                ev = _q[_head];
                _head = (_head + 1) % QUEUE_CAP;
                _count--;
                got = true;
            }
            portEXIT_CRITICAL(&_mux);
            if (!got) break;
            if (ev.action == 0) {
                _hid->keyboard.pressRaw(ev.hid);
            } else {
                _hid->keyboard.releaseRaw(ev.hid);
            }
            emitted++;
        }
        _hid->endCritical();
        return emitted;
    }

    // ---- Absolute mouse (BT Keyboard trackpad) ----
    //
    // Mouse state is kept separate from the keystroke queue and is NOT
    // cadence-buffered — the trackpad wants low latency. We coalesce: only
    // the latest position/buttons matter, and wheel ticks accumulate. The
    // main loop drains the latest state with drainMouse(). Because absolute
    // positions are self-correcting, dropping intermediate moves is fine.
    // Callable from the NimBLE task — no USB HID here, just state under the
    // spinlock.
    void enqueueMouse(uint8_t buttons, uint16_t x, uint16_t y, int8_t wheel) {
        portENTER_CRITICAL(&_mux);
        _mouseButtons = buttons;
        _mouseX = x;
        _mouseY = y;
        _mouseWheel += wheel;
        _mousePending = true;
        portEXIT_CRITICAL(&_mux);
    }

    // Called from the main loop (outside the HID-critical window). Emits the
    // latest absolute pointer report if one is pending. No delay/scheduling.
    void drainMouse() {
        if (!_hid) return;
        bool pending;
        uint8_t buttons;
        uint16_t x, y;
        int wheel;
        portENTER_CRITICAL(&_mux);
        pending = _mousePending;
        buttons = _mouseButtons;
        x = _mouseX;
        y = _mouseY;
        wheel = _mouseWheel;
        _mousePending = false;
        _mouseWheel = 0;
        portEXIT_CRITICAL(&_mux);
        if (!pending) return;
        int8_t w = (wheel > 127) ? 127 : (wheel < -127 ? -127 : (int8_t)wheel);
        _hid->absMouseReport(buttons, x, y, w);
    }

    // Drained by the BLE manager when it builds the next status frame.
    bool takeOverflowFlag() {
        bool was;
        portENTER_CRITICAL(&_mux);
        was = _overflow;
        _overflow = false;
        portEXIT_CRITICAL(&_mux);
        return was;
    }

private:
    HIDController* _hid = nullptr;
    Event _q[QUEUE_CAP];
    int _head = 0;
    int _tail = 0;
    int _count = 0;
    volatile bool _active = false;
    volatile bool _overflow = false;
    // Anchor mapping host_t=0 to a specific device millis() value.
    // Set on the first event of each session so subsequent events can
    // schedule emissions relative to the host's typing cadence.
    bool _anchorSet = false;
    uint32_t _anchorMs = 0;

    // Absolute mouse state (coalesced; see enqueueMouse/drainMouse).
    volatile bool _mousePending = false;
    uint8_t _mouseButtons = 0;
    uint16_t _mouseX = 0;
    uint16_t _mouseY = 0;
    int _mouseWheel = 0;

    portMUX_TYPE _mux = portMUX_INITIALIZER_UNLOCKED;
};