#include <ESP32-TWAI-CAN.hpp>
#include <Preferences.h>

#define CAN_TX 17
#define CAN_RX 16
#define STATUS_LED 2
#define LONG_PRESS_MS 500

// #define USE_CAN_BUS 

Preferences prefs;
enum Mode { MOM = 0, TOG = 1, OFF = 2, PUL = 3, STA = 4 };

struct StateStep {
    uint32_t canID;
    uint8_t byteIdx;
    uint8_t bitMask;
    bool isInverted;
    bool isLatching;
    uint8_t nextShort; 
    uint8_t nextLong;  
};

struct TxBtnConfig {
    char name[16];      
    Mode mode;         
    uint32_t heartbeatMs; 
    uint32_t pulseMs;
    StateStep states[10]; 
};

struct GlobalConfig {
    uint32_t canSpeed;
    bool isBigEndian;
    bool logHB;
    uint8_t bpMode; 
} gCfg;

const uint8_t buttonPins[] = {13, 14, 18, 19, 21, 22, 23, 25, 26, 27, 32, 33, 34, 35, 36, 39};
TxBtnConfig myButtons[16];
bool lastBtnStates[16], activeStates[16] = {false};
bool longPressTriggered[16] = {false};
unsigned long offTimestamps[16] = {0}, lastSentTimes[32] = {0}, pressStartTime[16] = {0};
int currentIdx[16] = {0}; 

uint32_t uniqueIDs[32];
uint32_t idHeartbeatRates[32];
uint8_t lastPayloads[32][8];
int uniqueCount = 0;

void refreshUniqueIDs() {
    uniqueCount = 0;
    for (int i = 0; i < 16; i++) {
        if (myButtons[i].mode == OFF) continue;
        for (int s = 0; s < 10; s++) {
            uint32_t id = myButtons[i].states[s].canID;
            if (id == 0) continue;
            int found = -1;
            for (int j = 0; j < uniqueCount; j++) if (uniqueIDs[j] == id) { found = j; break; }
            if (found == -1 && uniqueCount < 32) {
                uniqueIDs[uniqueCount] = id;
                idHeartbeatRates[uniqueCount] = myButtons[i].heartbeatMs;
                memset(lastPayloads[uniqueCount], 0, 8);
                uniqueCount++;
            }
        }
    }
}

void printConfig(int i) {
    if (i < 0 || i >= 16) return;
    const char* m = (myButtons[i].mode == OFF) ? "OFF" : (myButtons[i].mode == PUL ? "PUL" : (myButtons[i].mode == TOG ? "TOG" : (myButtons[i].mode == STA ? "STA" : "MOM")));
    Serial.printf("SET %d NAM=%s MDE=%s HB=%u PUL=%u\n", i, myButtons[i].name, m, myButtons[i].heartbeatMs, myButtons[i].pulseMs);
    for(int s = 0; s < 10; s++) {
        if (s > 0 && myButtons[i].mode != STA) continue;
        if (s > 0 && myButtons[i].states[s].canID == 0) continue;
        char ns[4], nl[4];
        if (myButtons[i].states[s].nextShort == 255) strcpy(ns, "X"); else sprintf(ns, "%d", myButtons[i].states[s].nextShort);
        if (myButtons[i].states[s].nextLong == 255) strcpy(nl, "X"); else sprintf(nl, "%d", myButtons[i].states[s].nextLong);
        Serial.printf("SET %d %d ID=%X BYT=%d MSK=%02X INV=%d LAT=%d NXT_S=%s NXT_L=%s %s\n", 
            i, s, myButtons[i].states[s].canID, myButtons[i].states[s].byteIdx, 
            myButtons[i].states[s].bitMask, myButtons[i].states[s].isInverted ? 1 : 0,
            myButtons[i].states[s].isLatching ? 1 : 0, ns, nl, (currentIdx[i] == s) ? "*" : "");
    }
}

void writeDefaults() {
    Serial.println("Restoring Defaults (500k, BIGENDIAN, HBOFF, BPON)...");
    gCfg = {500, true, false, 1}; 
    for(int i = 0; i < 16; i++) {
        memset(myButtons[i].name, 0, 16);
        sprintf(myButtons[i].name, "BTN_%d", i);
        myButtons[i].mode = OFF;
        myButtons[i].heartbeatMs = 0;
        myButtons[i].pulseMs = 0;
        for(int s = 0; s < 10; s++) {
            myButtons[i].states[s] = {0, 0, 0, false, true, (uint8_t)((s+1)%10), 255}; 
        }
    }
    strcpy(myButtons[3].name, "STA_DEMO"); myButtons[3].mode = STA; myButtons[3].pulseMs = 50;
    myButtons[3].states[0] = {0x200, 2, 0x00, false, true, 1, 0};
    myButtons[3].states[1] = {0x200, 2, 0x01, false, false, 2, 0}; 
    myButtons[3].states[2] = {0x200, 2, 0x02, false, true, 1, 0};

    prefs.begin("tx-btns", false);
    prefs.putBytes("config", myButtons, sizeof(myButtons));
    prefs.putBytes("glob", &gCfg, sizeof(gCfg));
    prefs.end();
    ESP.restart();
}

void handleConfigCommand(String cmd) {
    if (cmd.startsWith("BAUD=")) {
        gCfg.canSpeed = cmd.substring(5).toInt();
        prefs.begin("tx-btns", false); prefs.putBytes("glob", &gCfg, sizeof(gCfg)); prefs.end();
        Serial.printf("Baud set to %u. Restarting...\n", gCfg.canSpeed); delay(100); ESP.restart();
        return;
    }
    if (cmd.startsWith("ENDIN=")) {
        gCfg.isBigEndian = (cmd.substring(6).toInt() == 1);
        prefs.begin("tx-btns", false); prefs.putBytes("glob", &gCfg, sizeof(gCfg)); prefs.end();
        Serial.printf("Endian: %s\n", gCfg.isBigEndian ? "BIGENDIAN" : "LITTLEENDIAN"); 
        return;
    }

    char buf[128]; cmd.toCharArray(buf, sizeof(buf));
    strtok(buf, " ");
    char* btnToken = strtok(NULL, " ");
    if (!btnToken) { Serial.println("ERR: Missing Button Index"); return; }
    
    int i = atoi(btnToken);
    if (i < 0 || i >= 16) { Serial.printf("ERR: Invalid Button %d\n", i); return; }

    int stp = 0;
    char* nextToken = strtok(NULL, " ");
    if (nextToken && isDigit(nextToken[0]) && strchr(nextToken, '=') == NULL) {
        stp = constrain(atoi(nextToken), 0, 9);
        nextToken = strtok(NULL, " ");
    }

    bool updated = false;
    char* token = nextToken;
    while (token != NULL) {
        String t = String(token); int sep = t.indexOf('=');
        if (sep != -1) {
            String key = t.substring(0, sep); key.toUpperCase();
            String val = t.substring(sep + 1);
            updated = true;

            if (key == "NAM") { memset(myButtons[i].name, 0, 16); strncpy(myButtons[i].name, val.c_str(), 15); }
            else if (key == "MDE") { 
                val.toUpperCase(); myButtons[i].mode = (val=="STA")?STA:(val=="OFF")?OFF:(val=="PUL")?PUL:(val=="TOG")?TOG:MOM; 
                currentIdx[i] = 0; 
            }
            else if (key == "HB") myButtons[i].heartbeatMs = strtoul(val.c_str(), NULL, 10);
            else if (key == "PUL") myButtons[i].pulseMs = strtoul(val.c_str(), NULL, 10);
            else if (key == "ID") myButtons[i].states[stp].canID = strtoul(val.c_str(), NULL, 16);
            else if (key == "BYT") myButtons[i].states[stp].byteIdx = constrain(atoi(val.c_str()), 0, 7);
            else if (key == "MSK") myButtons[i].states[stp].bitMask = strtoul(val.c_str(), NULL, 16);
            else if (key == "INV") myButtons[i].states[stp].isInverted = (val == "1");
            else if (key == "LAT") myButtons[i].states[stp].isLatching = (val == "1");
            else if (key == "NXT_S" || key == "NXT_L") {
                uint8_t nextVal = (val.equalsIgnoreCase("X")) ? 255 : (uint8_t)atoi(val.c_str());
                if (key == "NXT_S") myButtons[i].states[stp].nextShort = nextVal;
                else myButtons[i].states[stp].nextLong = nextVal;
            } else { Serial.printf("ERR: Unknown Flag '%s'\n", key.c_str()); updated = false; }
        }
        token = strtok(NULL, " ");
    }

    if (updated) {
        prefs.begin("tx-btns", false); prefs.putBytes("config", myButtons, sizeof(myButtons)); prefs.end(); 
        refreshUniqueIDs(); Serial.printf("OK: Updated Button %d (Step %d)\n", i, stp); printConfig(i);
    }
}

void sendFrame(uint32_t id, uint8_t* data, int idx, bool hb) {
    memcpy(lastPayloads[idx], data, 8); lastSentTimes[idx] = millis();
    if (!hb || gCfg.logHB) { 
        Serial.printf("[%s] ID:0x%X Data:", hb ? "HB" : "CHG", id); 
        for(int d=0; d<8; d++) Serial.printf(" %02X", data[d]); Serial.println(); 
    }
#ifdef USE_CAN_BUS
    CanFrame f = {0}; f.identifier = id; f.extd = 0; f.data_length_code = 8; memcpy(f.data, data, 8);
    digitalWrite(STATUS_LED, HIGH); ESP32Can.writeFrame(f); digitalWrite(STATUS_LED, LOW);
#endif
}

void setup() {
    Serial.begin(115200); pinMode(STATUS_LED, OUTPUT);
    prefs.begin("tx-btns", false);
    if (!prefs.isKey("config") || !prefs.isKey("glob")) { prefs.end(); writeDefaults(); } 
    else { prefs.getBytes("config", myButtons, sizeof(myButtons)); prefs.getBytes("glob", &gCfg, sizeof(gCfg)); prefs.end(); }
    
    for(int i=0; i<16; i++) pinMode(buttonPins[i], (buttonPins[i] >= 34) ? INPUT : INPUT_PULLUP); 
    
    delay(500); 

    for(int i=0; i<16; i++) { 
        bool physicalLow = (digitalRead(buttonPins[i]) == LOW);
        lastBtnStates[i] = (buttonPins[i] >= 34) ? !physicalLow : physicalLow;
        currentIdx[i] = 0;
        pressStartTime[i] = 0; 
    }
    refreshUniqueIDs();
#ifdef USE_CAN_BUS
    ESP32Can.setPins(CAN_TX, CAN_RX); ESP32Can.begin(ESP32Can.convertSpeed(gCfg.canSpeed));
#endif
    Serial.printf("--- READY (BP=%d %s %dkbps) ---\n", gCfg.bpMode, gCfg.isBigEndian ? "BIGENDIAN" : "LITTLEENDIAN", gCfg.canSpeed);
}

void loop() {
    if (Serial.available()) {
        String input = Serial.readStringUntil('\n'); input.trim();
        if (input.equalsIgnoreCase("HELP")) {
            Serial.println("\n--- AVAILABLE COMMANDS ---");
            Serial.println("BAUD=[val]    : Set CAN speed (125, 250, 500, 1000). Reboots device.");
            Serial.println("ENDIN=[0/1]   : 0=LITTLEENDIAN, 1=BIGENDIAN.");
            Serial.println("HBON / HBOFF  : Toggle Heartbeat/Payload transmission logs.");
            Serial.println("BPOFF         : Disable all button press debug logs.");
            Serial.println("BPON          : Log button events (Short/Long) and State changes.");
            Serial.println("BPLONG        : Log detailed press/release durations and timestamps.");
            Serial.println("GET           : Show config for all 16 buttons.");
            Serial.println("GET [index]   : Show config for a specific button (0-15).");
            Serial.println("RESET         : Wipe all settings and restore factory defaults.");
            Serial.println("SET [btn] [stp] FLAG=VAL : Configure button logic.");
            Serial.println("   Flags: NAM=Name, MDE=MOM|TOG|PUL|STA|OFF, HB=HeartbeatMS, PUL=PulseMS,");
            Serial.println("          ID=HexID, BYT=0-7, MSK=HexMask, INV=0|1, LAT=0|1 (Latching),");
            Serial.println("          NXT_S=Next Short Step, NXT_L=Next Long Step (Use 'X' to Disable).");
        }
        else if (input.startsWith("SET ") || input.startsWith("BAUD=") || input.startsWith("ENDIN=")) handleConfigCommand(input);
        else if (input.equalsIgnoreCase("GET")) { for(int i=0; i<16; i++) printConfig(i); }
        else if (input.startsWith("GET ")) printConfig(input.substring(4).toInt());
        else if (input.equalsIgnoreCase("BPOFF")) gCfg.bpMode = 0;
        else if (input.equalsIgnoreCase("BPON")) gCfg.bpMode = 1;
        else if (input.equalsIgnoreCase("BPLONG")) gCfg.bpMode = 2;
        else if (input.equalsIgnoreCase("HBON")) gCfg.logHB = true;
        else if (input.equalsIgnoreCase("HBOFF")) gCfg.logHB = false;
        else if (input.equalsIgnoreCase("RESET")) writeDefaults();
        
        if (input.equalsIgnoreCase("BPOFF") || input.equalsIgnoreCase("BPON") || input.equalsIgnoreCase("BPLONG") || input.equalsIgnoreCase("HBON") || input.equalsIgnoreCase("HBOFF")) {
            prefs.begin("tx-btns", false); prefs.putBytes("glob", &gCfg, sizeof(gCfg)); prefs.end();
            Serial.printf("Mode Updated: BP=%d HB=%d\n", gCfg.bpMode, gCfg.logHB);
        }
    }

    unsigned long now = millis();
    for (int i = 0; i < 16; i++) {
        bool pressed = (digitalRead(buttonPins[i]) == LOW);
        if (buttonPins[i] >= 34) pressed = !pressed;

        if (pressed != lastBtnStates[i]) {
            if (pressed) { 
                if(gCfg.bpMode == 2) Serial.printf("BP: Button %d Pressed\n", i);
                pressStartTime[i] = now; 
                offTimestamps[i] = now + myButtons[i].pulseMs;
                longPressTriggered[i] = false;
                if (myButtons[i].mode != OFF && myButtons[i].mode != STA) activeStates[i] = true;
                if (myButtons[i].mode == TOG) activeStates[i] = !activeStates[i];
            } else { 
                unsigned long dur = now - pressStartTime[i];
                if(gCfg.bpMode == 2) Serial.printf("BP: Button %d Released (%lu ms)\n", i, dur);
                if (pressStartTime[i] != 0) {
                    if (dur < LONG_PRESS_MS) {
                        if(gCfg.bpMode >= 1) Serial.printf("BP: Button %d SHORT\n", i);
                        if (myButtons[i].mode == STA) {
                            uint8_t next = myButtons[i].states[currentIdx[i]].nextShort;
                            if (next != 255) {
                                currentIdx[i] = next;
                                if (!myButtons[i].states[currentIdx[i]].isLatching) {
                                    activeStates[i] = true;
                                    offTimestamps[i] = now + myButtons[i].pulseMs;
                                }
                            }
                        }
                    }
                    if (myButtons[i].mode == MOM && now >= offTimestamps[i]) activeStates[i] = false;
                    pressStartTime[i] = 0;
                }
            }
            lastBtnStates[i] = pressed; delay(20);
        }

        if (pressed && !longPressTriggered[i] && (now - pressStartTime[i] >= LONG_PRESS_MS)) {
            if(gCfg.bpMode >= 1) Serial.printf("BP: Button %d LONG\n", i);
            if (myButtons[i].mode == STA) {
                uint8_t next = myButtons[i].states[currentIdx[i]].nextLong;
                if (next != 255) {
                    currentIdx[i] = next;
                    if (!myButtons[i].states[currentIdx[i]].isLatching) {
                        activeStates[i] = true;
                        offTimestamps[i] = now + myButtons[i].pulseMs;
                    }
                }
            }
            longPressTriggered[i] = true; 
        }

        if (activeStates[i] && (myButtons[i].mode == PUL || myButtons[i].mode == STA || (myButtons[i].mode == MOM && !pressed))) {
            if (now >= offTimestamps[i]) activeStates[i] = false;
        }
    }

    for (int j = 0; j < uniqueCount; j++) {
        uint32_t id = uniqueIDs[j]; uint8_t currentData[8] = {0};
        for (int b = 0; b < 16; b++) {
            StateStep &s = myButtons[b].states[currentIdx[b]];
            if (s.canID == id) {
                bool val = (myButtons[b].mode == STA) ? (s.isLatching ? true : activeStates[b]) : (s.isInverted ? !activeStates[b] : activeStates[b]);
                if (val) {
                    int targetByte = gCfg.isBigEndian ? (7 - s.byteIdx) : s.byteIdx;
                    currentData[targetByte] |= s.bitMask;
                }
            }
        }
        bool hb = (idHeartbeatRates[j] > 0 && now - lastSentTimes[j] >= idHeartbeatRates[j]);
        bool dataChanged = memcmp(lastPayloads[j], currentData, 8) != 0;
        if (dataChanged || hb) sendFrame(id, currentData, j, hb && !dataChanged);
    }
}