#include <M5Unified.h>
#include "config.h"
#include "settings.h"
#include "usb_hid.h"
#include "led_ui.h"
#include "display_ui.h"
#include "macro_storage.h"
#include "ble_keystore.h"
#include "ble_manager.h"
#include "espnow_manager.h"
#include "macro_engine.h"
#include "serial_protocol.h"
#include "live_keystroke.h"

SettingsManager settingsManager;
HIDController hid;
LedUI ledUI;
DisplayUI display;
MacroStorage storage;
DebugLog debugLog;
BLEKeyStore bleKeyStore;
BLEManager bleManager;
EspNowManager espnowManager;
MacroEngine engine;
SerialProtocol protocol;
LiveKeystrokeEngine liveEngine;
HardwareSerial rs232Serial(1);  // UART1 for RS232 via Atomic RS232 Base

// Passed to SerialProtocol so that when the host-side terminal reconfigures
// the RS232 port, any cached baud/config in the engine gets invalidated.
void onRS232Reconfig() {
    engine.resetRS232Cache();
}

int currentMacroIdx = 0;

// Track the last status string we drew to showLiveMode so we only
// re-render when the BLE state actually changes (avoids flicker and
// SPI traffic during heavy keystroke streaming).
static const char* _liveLastStatus = nullptr;
// True while the full-screen live-mode display is up (a host is
// connected). Lets us detect the connected->idle transition and
// repaint the macro selector exactly once.
static bool _liveScreenShown = false;
// True while the Bluetooth identify logo is up (host asked us to
// identify ourselves so the user can label this device).
static bool _liveIdentifyShown = false;
// Last device-label version we rendered, so a new label repaints the
// live-mode screen even when the status string hasn't changed.
static uint32_t _liveLabelVerShown = 0;
// Set true if the user cancels the auto-reconnect (power-loss resume) with a
// button hold. Suppresses BLE for the rest of this boot; a power cycle
// re-enables it.
static bool _liveResumeCancelled = false;
// True while the "Reconnecting..." screen is up (resume pending, not yet
// connected). Lets us repaint the selector exactly once when it clears.
static bool _resumeScreenShown = false;
// True while the "MESH HUB" screen is up; repainted when the hub's node
// roster changes (count shown on screen).
static bool _hubScreenShown = false;
static uint32_t _hubRosterVerShown = 0;
// Latches the 5-second transport-toggle so it fires exactly once per press
// (pressedFor() stays true for every update past the threshold). Cleared on
// button release.
static bool _modeSwitchArmed = false;

void setup() {
    // Disable DTR/RTS bootloader reboot ASAP — CDC_ON_BOOT means TinyUSB
    // is already running before setup(). This prevents a crash-restart cycle
    // from entering bootloader when the Python app has the port open.
    USBSerial.enableReboot(false);

    // Initialize M5 WITHOUT touching USB serial
    auto cfg = M5.config();
    cfg.serial_baudrate = 0;  // Prevent M5.begin() from calling Serial.begin()
    M5.begin(cfg);

    // Initialize USB composite device (HID + CDC)
    hid.begin();

    // Short serial timeout so readStringUntil doesn't block the loop
    Serial.setTimeout(100);

    settingsManager.begin();

    M5.BtnA.setHoldThresh(settingsManager.settings.holdMs);
    M5.BtnA.setDebounceThresh(DEBOUNCE_MS);

    // Universal binary: M5GFX panel autodetect found an LCD on the AtomS3;
    // on the AtomS3 Lite there is none and getBoard() reports the Lite.
    // All screen output is gated inside DisplayUI; user feedback on the
    // Lite comes from the RGB LED (LedUI, via M5.Led).
    bool hasDisplay = (M5.getDisplayCount() > 0) &&
                      (M5.getBoard() != m5::board_t::board_M5AtomS3Lite);
    ledUI.begin(!hasDisplay);
    display.begin(settingsManager.settings.orientation, hasDisplay, &ledUI);
    display.showBoot(settingsManager.settings.liveTransport == LIVE_TX_BLE
                         ? "Mode: BLE" : "Mode: Mesh");

    if (!storage.begin()) {
        display.showMessage("FS Error!", TFT_RED);
        delay(2000);
    }
    storage.loadSubIndex();

    // BLE payload-encryption key — load from LittleFS or generate on first boot.
    // Must come after storage.begin() since LittleFS is mounted there.
    bleKeyStore.begin();

    debugLog.begin();
    debugLog.log("Device booted");

    engine.begin(&hid, &display, &storage, &settingsManager, &rs232Serial);

    // ESP-NOW mesh (live-keyboard transport). Radio stays OFF until the
    // idle loop brings up node listening (or the host app switches us
    // into hub mode over USB).
    espnowManager.begin(&settingsManager, &bleKeyStore, &liveEngine, &debugLog);

    protocol.begin(&settingsManager, &storage, &display, &debugLog,
                   &rs232Serial, onRS232Reconfig, &bleKeyStore, &bleManager,
                   &espnowManager);

    // Give USB time to fully enumerate
    delay(1500);
    showCurrentMacro();

    // BLE manager — only stores debug log pointer here.
    // NimBLE is NOT started at boot; it's started on-demand when a
    // bluetooth node is hit, then shut down after variables are received.
    // This avoids all BLE/USB radio contention during normal operation.
    bleManager.begin(&debugLog, &bleKeyStore);
    engine.setBLEManager(&bleManager);

    // Wire up the live-keystroke engine. The BLE manager pushes events
    // into it from its NimBLE write callback; we drain in the main loop.
    liveEngine.begin(&hid);
    bleManager.setLiveEngine(&liveEngine);

    // Check for saved execution state (power-loss recovery)
    int resumeSlot = 0, resumeNode = 0;
    if (engine.checkSavedState(resumeSlot, resumeNode)) {
        uint16_t delaySeconds = settingsManager.settings.resumeDelay;
        // resumeSlot is the actual LittleFS slot number (not display index).
        // Validate it exists by checking if the slot appears in the order array.
        bool slotValid = false;
        for (int i = 0; i < storage.macroCount; i++) {
            if (storage.order[i] == resumeSlot) { slotValid = true; break; }
        }

        Serial.printf("[BOOT] resume: slot=%d node=%d slotValid=%d delaySec=%u\n",
                      resumeSlot, resumeNode, (int)slotValid, (unsigned)delaySeconds);

        if (delaySeconds > 0 && slotValid) {
            // Drain stale button state. After a USB-power blip the M5.Btn
            // driver can latch a "wasClicked" on the first update — if we
            // peek at it during the countdown we'd cancel the resume the
            // user is depending on. Burn ~200ms of updates so anything
            // pending settles before we start watching for real input.
            for (int i = 0; i < 20; i++) { M5.update(); delay(10); }
            (void)M5.BtnA.wasClicked();
            (void)M5.BtnA.wasHold();

            // Countdown with cancel option. A HOLD (long press) cancels; a
            // single click is ignored — it's too easy to bump the button
            // accidentally while watching imaging, and that would silently
            // discard the resume.
            bool cancelled = false;
            uint16_t cancelledAt = delaySeconds;
            for (uint16_t remaining = delaySeconds; remaining > 0; remaining--) {
                char msg[64];
                snprintf(msg, sizeof(msg), "Resuming in %ds\nHold to cancel", remaining);
                display.showMessage(msg, TFT_YELLOW);
                // Poll button every 100ms during each second
                for (int i = 0; i < 10; i++) {
                    M5.update();
                    if (M5.BtnA.wasHold()) {
                        cancelled = true;
                        cancelledAt = remaining;
                        break;
                    }
                    delay(100);
                }
                if (cancelled) break;
            }

            if (!cancelled) {
                Serial.println("[BOOT] resume: countdown completed, resuming macro");
                if (!engine.resumeMacro(resumeSlot, resumeNode,
                                        storage.macros[resumeSlot].name)) {
                    Serial.println("[BOOT] resume: resumeMacro() failed, clearing state");
                    engine.clearExecutionState();
                    showCurrentMacro();
                }
            } else {
                Serial.printf("[BOOT] resume: cancelled by hold at %us remaining\n",
                              (unsigned)cancelledAt);
                engine.clearExecutionState();
                showCurrentMacro();
            }
        } else {
            Serial.printf("[BOOT] resume: skipped (delay=%u, slotValid=%d) — clearing state\n",
                          (unsigned)delaySeconds, (int)slotValid);
            engine.clearExecutionState();
        }
    } else {
        Serial.println("[BOOT] no resume state to load");
    }
}

// ---- Live-transport abstraction ---------------------------------------
// The device receives a live-keyboard session over exactly one radio,
// chosen by settings.liveTransport. These helpers hide which one is active
// so the idle loop's screen/reconnect logic is written once for both.

static bool liveIsBle() {
    return settingsManager.settings.liveTransport == LIVE_TX_BLE;
}

static bool liveSessionActive() {
    return liveIsBle()
        ? (bleManager.exchangeKind() == BLEManager::EX_LIVE &&
           bleManager.isClientConnected())
        : espnowManager.nodeInSession();
}

static bool liveIdentifyActive() {
    return liveIsBle() ? bleManager.liveIdentify()
                       : espnowManager.nodeIdentify();
}

static const char* liveStatusStr() {
    return liveIsBle() ? bleManager.liveStatusText()
                       : espnowManager.nodeStatusText();
}

static const char* liveLabelStr() {
    return liveIsBle() ? bleManager.liveLabel() : espnowManager.nodeLabel();
}

static uint32_t liveLabelVerNum() {
    return liveIsBle() ? bleManager.liveLabelVer()
                       : espnowManager.nodeLabelVer();
}

static bool liveResumePending() {
    return liveIsBle() ? bleManager.liveResumeRequested()
                       : espnowManager.resumeRequestedAtBoot();
}

static void liveConsumeResume() {
    if (liveIsBle()) bleManager.consumeLiveResume();
    else             espnowManager.consumeResume();
}

// Free both live radios. Safe to call when either/both are already down
// (each teardown is idempotent). Used before running a routine, on a
// transport toggle, and when a USB host claims the device.
static void liveShutdownRadios() {
    espnowManager.shutdown();
    bleManager.stopLive();
    bleManager.shutdown();
}

// Flip the persisted live transport (mesh <-> BLE), tear the current radio
// down so the idle loop brings the new one up, and confirm on-screen. Called
// from the idle selector when the button is held for MODE_SWITCH_HOLD_MS.
static void toggleLiveTransport() {
    uint8_t next = liveIsBle() ? LIVE_TX_MESH : LIVE_TX_BLE;
    settingsManager.set("live_tx", next);   // persists to NVS

    liveShutdownRadios();
    // Cancel any pending power-loss auto-reconnect for the old transport and
    // re-arm listening for the new one.
    espnowManager.consumeResume();
    bleManager.consumeLiveResume();
    _liveResumeCancelled = false;
    _liveScreenShown = false;
    _liveIdentifyShown = false;
    _resumeScreenShown = false;
    _liveLastStatus = nullptr;

    display.showModeSwitch(next == LIVE_TX_BLE);
    // Hold the confirmation ~1.2 s. Pump the LED engine (no-op on an LCD
    // board) so a screenless Lite actually animates its mode-switch burst.
    // Deliberately do NOT call M5.update() here — the button release must
    // be left for the main loop to observe so _modeSwitchArmed clears.
    uint32_t until = millis() + 1200;
    while ((int32_t)(millis() - until) < 0) {
        ledUI.tick();
        delay(20);
    }
    showCurrentMacro();
}

void loop() {
    M5.update();

    // Safety net: re-assert reboot disable every loop iteration
    USBSerial.enableReboot(false);

    // LED pattern engine (AtomS3 Lite only; no-op with a display). Cheap:
    // recomputes the current pattern color and writes only on change.
    ledUI.tick();

    // Keyboard-priority gate: when the HID controller is in the middle
    // of a synchronous USB-emitting operation (typeText, keyCombo,
    // mediaKey, probe), defer non-USB-HID housekeeping. These polls
    // don't touch the keyboard interface themselves, but they share the
    // main-loop task with engine.tick() — and skipping them keeps the
    // FreeRTOS scheduling latency on the typing path as low as possible
    // for any future change that pumps the loop while typing.
    //
    // In current code the main loop is already blocked inside tick()
    // for the duration of a keyboard op, so the flag is normally only
    // observed BETWEEN ops. We still gate here so the contract holds.
    bool hidCritical = hid.isCritical();

    if (!hidCritical) {
        // Print deferred BLE status on the main task (safe for TinyUSB CDC)
        bleManager.pollStatus();
    }

    // Drain queued live-mode keystrokes outside the HID-critical window.
    // The drain itself wraps in beginCritical/endCritical so any other
    // gated work observes the new in-flight state — but we only START
    // draining when the previous gate has lifted, so emissions don't
    // stack on top of each other.
    if (!hidCritical && liveEngine.isActive()) {
        if (liveEngine.drainQueue() > 0) {
            // Screenless boards flicker the LED so the user can see
            // keystrokes flowing (no-op when a display is present).
            ledUI.liveActivity();
        }
    }

    // Drain the latest absolute-mouse report (BT Keyboard trackpad). Not
    // cadence-buffered — emitted as soon as the HID-critical window is clear.
    if (!hidCritical) {
        liveEngine.drainMouse();
    }

    bool settingsChanged = protocol.handleSerial();

    if (!hidCritical) {
        // Drain RS232 RX into the terminal buffer when the host-side terminal is open
        protocol.pollRS232();
    }

    if (settingsChanged) {
        M5.BtnA.setHoldThresh(settingsManager.settings.holdMs);
    }

    if (protocol.needsRefresh()) {
        if (currentMacroIdx >= storage.macroCount) {
            currentMacroIdx = 0;
        }
        _liveScreenShown = false;
        _liveIdentifyShown = false;
        showCurrentMacro();
    }

    // ---- ESP-NOW mesh (live-keyboard transport) ----
    //
    // The mesh replaced the old per-device live-BLE channel. While idle
    // the device passively listens on the mesh channel so the hub (the
    // USB-attached unit the host app drives) can discover and JOIN it
    // WITHOUT any on-device gesture. Policy mirrors the old BLE one:
    //
    //   * Profile uploads over USB (isUploadActive) and a USB-connected
    //     host app keep the radio OFF — unless the host explicitly made
    //     us the hub (espnow_hub command), which overrides isHostConnected.
    //   * A running routine owns the device for HID, so the radio is torn
    //     down the moment a routine starts (below). This also guarantees
    //     the on-demand BLE variables exchange never coexists with WiFi.
    //   * 10-second boot grace, skipped when the power-loss resume flag is
    //     set so a host session reconnects without delay. A button hold
    //     cancels the resume for this boot (idle UI below).
    //
    // Security: JOIN must decrypt under this device's AES key and the
    // keystroke stream under the session group key — a radio-local
    // attacker can never inject input.
    espnowManager.tick();

    // Bring up whichever live radio settings.liveTransport selects (never
    // both — BLE and WiFi contend on the S3). The hub path is unaffected:
    // it's entered only over USB and owns the radio while active.
    static const uint32_t LIVE_LISTEN_BOOT_GRACE_MS = 10000;
    bool bleLiveUp = bleManager.isBLEActive() &&
                     bleManager.exchangeKind() == BLEManager::EX_LIVE;
    bool resumeReq = liveResumePending() && !_liveResumeCancelled;
    if (espnowManager.isHub()) {
        // Host-driven bridge. It auto-reverts to OFF if the host app goes
        // quiet (see espnow_manager.h), so nothing to police here.
    } else if (protocol.isHostConnected() || protocol.isUploadActive()) {
        // A USB host app is talking to us — this is the configuring
        // computer, not a remote-resume scenario. Drop the resume request
        // and keep both live radios off for the rest of the boot.
        espnowManager.consumeResume();
        bleManager.consumeLiveResume();
        if (espnowManager.isRadioActive() || bleLiveUp) {
            liveShutdownRadios();
            _liveScreenShown = false;
            _liveIdentifyShown = false;
            _resumeScreenShown = false;
        }
    } else if (!engine.isRunning() && !_liveResumeCancelled &&
               (resumeReq || millis() > LIVE_LISTEN_BOOT_GRACE_MS)) {
        if (liveIsBle()) {
            // Ensure the mesh radio is down, then advertise BLE live.
            if (espnowManager.isRadioActive()) espnowManager.shutdown();
            bleManager.startLive();          // idempotent
        } else {
            // Ensure BLE live is down, then idle-listen on the mesh.
            if (bleLiveUp) bleManager.stopLive();
            espnowManager.startNodeListen();  // idempotent
        }
    }

    // Don't process button for macro selection while receiving data
    if (protocol.isBusy()) return;

    if (engine.isRunning()) {
        engine.tick(settingsManager.settings.typeDelay);

        if (M5.BtnA.wasClicked()) {
            engine.onButtonClick();
        }
        if (M5.BtnA.wasHold()) {
            engine.onButtonHold();
        }

        // If engine just finished, return to macro selector
        if (!engine.isRunning()) {
            _liveScreenShown = false;
            _liveIdentifyShown = false;
            showCurrentMacro();
        }
    } else {
        // Idle. Hub mode first: the host app drives everything over USB;
        // we only show which unit is the hub and how many nodes it sees.
        if (espnowManager.isHub()) {
            if (!_hubScreenShown ||
                espnowManager.hubNodesVer() != _hubRosterVerShown) {
                display.showHubMode(espnowManager.hubNodeCount());
                _hubRosterVerShown = espnowManager.hubNodesVer();
                _hubScreenShown = true;
                _liveScreenShown = false;
                _liveIdentifyShown = false;
            }
            (void)M5.BtnA.wasClicked();
            (void)M5.BtnA.wasHold();
            return;
        }
        if (_hubScreenShown) {
            _hubScreenShown = false;
            showCurrentMacro();
        }

        // Joined a live session (on whichever transport is active): show
        // the live-mode screen and let the host drive. The only on-device
        // action is a long-press, which overrides into running the
        // currently-selected routine.
        bool liveConnected = liveSessionActive();
        if (liveConnected) {
            // Joined — the power-loss resume is satisfied; clear the
            // one-shot boot request so later grace logic is normal.
            liveConsumeResume();
            _resumeScreenShown = false;
            if (liveIdentifyActive()) {
                // Host is asking us to identify ourselves so the user can
                // label this specific device — draw the Bluetooth logo
                // (blue/white LED flash on a Lite).
                if (!_liveIdentifyShown) {
                    display.showBluetoothLogo();
                    _liveIdentifyShown = true;
                    _liveScreenShown = false;     // force a status repaint after
                    _liveLastStatus = nullptr;
                }
            } else {
                _liveIdentifyShown = false;
                const char* status = liveStatusStr();
                uint32_t lblVer = liveLabelVerNum();
                if (!_liveScreenShown || status != _liveLastStatus ||
                    lblVer != _liveLabelVerShown) {
                    display.showLiveMode(status, liveLabelStr());
                    _liveLastStatus = status;
                    _liveLabelVerShown = lblVer;
                    _liveScreenShown = true;
                    // Screenless mesh nodes: hub silence shows as the
                    // cyan/red "lost hub" pattern instead of live-idle.
                    if (!liveIsBle() && espnowManager.nodeLagging()) {
                        ledUI.lagging();
                    }
                }
            }
            (void)M5.BtnA.wasClicked();  // consumed, ignored while a host drives
            if (M5.BtnA.wasHold() && storage.macroCount > 0) {
                // Engine activity overrides a live session: free the radio,
                // then run the selected routine.
                liveShutdownRadios();
                _liveScreenShown = false;
                _liveIdentifyShown = false;
                int idx = currentMacroIdx % storage.macroCount;
                int slot = storage.order[idx];
                engine.startMacro(slot, storage.macros[slot].name);
            }
            return;
        }

        // No session yet. If we're auto-reconnecting after a power-loss
        // (resume flag set, not cancelled), show the "Reconnecting..."
        // screen and let a hold cancel it.
        if (liveResumePending() && !_liveResumeCancelled) {
            if (!_resumeScreenShown) {
                display.showReconnecting();
                _resumeScreenShown = true;
                _liveScreenShown = false;
                _liveIdentifyShown = false;
            }
            (void)M5.BtnA.wasClicked();
            if (M5.BtnA.wasHold()) {
                // Cancel the auto-reconnect: radios off for the rest of
                // this boot. A power cycle re-enables it.
                _liveResumeCancelled = true;
                liveConsumeResume();
                liveShutdownRadios();
                _resumeScreenShown = false;
                showCurrentMacro();
            }
            return;
        }

        // No session. Repaint the selector once if we were just showing
        // the live screen, then handle the normal gestures.
        if (_liveScreenShown || _liveIdentifyShown || _resumeScreenShown) {
            _liveScreenShown = false;
            _liveIdentifyShown = false;
            _resumeScreenShown = false;
            _liveLastStatus = nullptr;
            showCurrentMacro();
        }

        // Screen-button gestures (idle selector / listening):
        //   click               -> next macro
        //   hold 0.5s .. <5s     -> run the selected routine
        //   hold >= 5s           -> toggle the live transport (mesh <-> BLE)
        // The run gesture is classified on RELEASE so a long transport-toggle
        // hold has room to complete without the routine firing at ~0.5 s.
        if (M5.BtnA.pressedFor(MODE_SWITCH_HOLD_MS)) {
            if (!_modeSwitchArmed) {
                _modeSwitchArmed = true;   // fire once for this press
                toggleLiveTransport();
            }
            return;                        // hold still in progress
        }
        if (M5.BtnA.wasReleased()) {
            _modeSwitchArmed = false;      // ready for the next press
        }

        if (M5.BtnA.wasClicked()) {
            if (storage.macroCount > 0) {
                currentMacroIdx = (currentMacroIdx + 1) % storage.macroCount;
            }
            showCurrentMacro();
        }

        if (M5.BtnA.wasReleasedAfterHold() &&
            !M5.BtnA.wasReleaseFor(MODE_SWITCH_HOLD_MS) &&
            storage.macroCount > 0) {
            // Short hold released (not the 5 s toggle): tear both live radios
            // down so WiFi/BLE is fully off for HID emission (and for any BLE
            // variables node the routine may hit), then run.
            liveShutdownRadios();
            _liveScreenShown = false;
            int slot = storage.order[currentMacroIdx];
            engine.startMacro(slot, storage.macros[slot].name);
        }
    }
}

void showCurrentMacro() {
    if (storage.macroCount == 0) {
        display.showMessage("No Macros");
        return;
    }
    if (currentMacroIdx >= storage.macroCount) currentMacroIdx = 0;
    int slot = storage.order[currentMacroIdx];
    display.showMacroSelector(slot, storage.macros[slot].name, currentMacroIdx, storage.macroCount,
                              resolveColor(storage.macros[slot].labelColor),
                              settingsManager.settings.liveTransport == LIVE_TX_BLE);
}