#include <Keypad.h>
#include <EEPROM.h>   // internal EEPROM

// EEPROM wear-leveling configuration
const uint16_t EEPROM_START_ADDR = 0;   // starting offset in internal EEPROM
const uint8_t SLOT_COUNT = 32;          // number of rotating slots
const uint8_t SLOT_SIZE = 17;           // bytes per slot (magic + ver(2) + csum + payload)

// Slot layout (total SLOT_SIZE bytes):
// [0] magic(1) = 0xA5
// [1..2] version (uint16_t little-endian)
// [3] checksum (xor of payload bytes)
// [4] maxAntennas (1)
// [5] activeRelay (1)
// [6] rxRelay (1)
// [7] txRelay (1)
// [8] flags (bits: bit0=rxScanMode, bit1=txScanMode, bit2=swapMode)
// [9..16] blockedAntennas[8] (1 byte each, 0/1)

const uint8_t SLOT_MAGIC = 0xA5;

// Shift register pins
const int latchPin1 = 6;
const int clockPin1 = 7;
const int dataPin1 = 8;

const int latchPin2 = 9;
const int clockPin2 = 10;
const int dataPin2 = 11;

const int statusLedPin = 13;
const int swapLedPin = 12;
const int pttPin = 2;
const int txScanLedPin = 0;
const int MaxAntLedPin = 1;

// State variables
int activeRelay = 0;
int txRelay = 0;
int rxRelay = 0;
bool rxScanMode = false;
bool txScanMode = false;
bool swapMode = false;
int maxAntennas = 8;

// Array to track blocked antennas (true = blocked from TX)
bool blockedAntennas[8] = {false, false, false, false, false, false, false, false};

// Tracking changes
int lastRxRelay = -1;
int lastTxRelay = -1;
bool lastPttState = false;
bool lastRxScanMode = false;
bool lastTxScanMode = false;
bool lastSwapMode = false;

// Keypad configuration
const byte ROWS = 4;
const byte COLS = 3;
char keys[ROWS][COLS] = {
  {'1','2','3'},
  {'4','5','6'},
  {'7','8','9'},
  {'*','0','#'}
};
byte rowPins[ROWS] = {A0, A1, A2, A3};
byte colPins[COLS] = {3, 4, 5};
Keypad keypad = Keypad(makeKeymap(keys), rowPins, colPins, ROWS, COLS);

// ---------- Utility: pack/unpack and checksum ----------

uint8_t packFlags() {
  uint8_t f = 0;
  if (rxScanMode) f |= 0x01;
  if (txScanMode) f |= 0x02;
  if (swapMode)   f |= 0x04;
  return f;
}

void unpackFlags(uint8_t f) {
  rxScanMode = (f & 0x01);
  txScanMode = (f & 0x02);
  swapMode   = (f & 0x04);
}

uint8_t calcPayloadChecksum(const uint8_t *payload, uint8_t len) {
  uint8_t c = 0;
  for (uint8_t i = 0; i < len; i++) c ^= payload[i];
  return c;
}

// ---------- Defaults and helpers ----------

void loadDefaults() {
  maxAntennas = 8;
  activeRelay = 0;
  rxRelay = 0;
  txRelay = 0;
  rxScanMode = false;
  txScanMode = false;
  swapMode = false;
  
  for (int i = 0; i < 8; i++) blockedAntennas[i] = false;
}

// Read a single slot into buffer; returns false if out of EEPROM bounds
bool readSlot(uint8_t slotIndex, uint8_t *slotBuf) {
  if (slotIndex >= SLOT_COUNT) return false;
  uint16_t base = EEPROM_START_ADDR + (uint16_t)slotIndex * SLOT_SIZE;
  for (uint8_t i = 0; i < SLOT_SIZE; i++) {
    slotBuf[i] = EEPROM.read((int)(base + i));
  }
  return true;
}

// Write slotBuffer to EEPROM at slotIndex using EEPROM.update (byte-by-byte)
bool writeSlot(uint8_t slotIndex, const uint8_t *slotBuf) {
  if (slotIndex >= SLOT_COUNT) return false;
  uint16_t base = EEPROM_START_ADDR + (uint16_t)slotIndex * SLOT_SIZE;
  for (uint8_t i = 0; i < SLOT_SIZE; i++) {
    EEPROM.update((int)(base + i), slotBuf[i]);
  }
  return true;
}

// ---------- Wear-leveling load/save for entire state ----------

void loadStateFromEEPROM() {
  // Scan slots and pick the record with the highest version that is valid (magic + checksum)
  uint16_t bestVersion = 0;
  bool found = false;
  uint8_t bestBuf[SLOT_SIZE];

  for (uint8_t s = 0; s < SLOT_COUNT; s++) {
    uint8_t buf[SLOT_SIZE];
    if (!readSlot(s, buf)) continue;

    if (buf[0] != SLOT_MAGIC) continue; // no record

    // read version (little-endian)
    uint16_t ver = (uint16_t)buf[1] | ((uint16_t)buf[2] << 8);

    // payload is bytes [4..(SLOT_SIZE-1)]
    uint8_t payloadLen = SLOT_SIZE - 4;
    uint8_t checksum = buf[3];
    uint8_t computed = calcPayloadChecksum(&buf[4], payloadLen);

    if (checksum != computed) continue; // corrupt

    // pick the highest version (uint16_t compare)
    if (!found || ver > bestVersion) {
      bestVersion = ver;
      memcpy(bestBuf, buf, SLOT_SIZE);
      found = true;
    }
  }

  if (!found) {
    // No valid stored state — use defaults
    loadDefaults();
    return;
  }

  // unpack payload
  uint8_t payloadIdx = 4;
  uint8_t ee_maxAntennas = bestBuf[payloadIdx++];
  uint8_t ee_activeRelay  = bestBuf[payloadIdx++];
  uint8_t ee_rxRelay      = bestBuf[payloadIdx++];
  uint8_t ee_txRelay      = bestBuf[payloadIdx++];
  uint8_t ee_flags        = bestBuf[payloadIdx++];
  bool ee_blockedAnt[8];
  for (uint8_t i = 0; i < 8; i++) ee_blockedAnt[i] = (bestBuf[payloadIdx++] == 1);

  // validate ranges
  bool dataValid = (ee_maxAntennas >= 1 && ee_maxAntennas <= 8 &&
                    ee_activeRelay < 8 && ee_rxRelay < 8 && ee_txRelay < 8);

  if (!dataValid) {
    loadDefaults();
    return;
  }

  // load into runtime
  maxAntennas = ee_maxAntennas;
  activeRelay = ee_activeRelay;
  rxRelay = ee_rxRelay;
  txRelay = ee_txRelay;
  unpackFlags(ee_flags);
  for (uint8_t i = 0; i < 8; i++) blockedAntennas[i] = ee_blockedAnt[i];
}

// Helper to compare current runtime state vs packed payload bytes
void packPayload(uint8_t *outPayload) {
  uint8_t idx = 0;
  outPayload[idx++] = (uint8_t)maxAntennas;
  outPayload[idx++] = (uint8_t)activeRelay;
  outPayload[idx++] = (uint8_t)rxRelay;
  outPayload[idx++] = (uint8_t)txRelay;
  outPayload[idx++] = packFlags();
  for (uint8_t i = 0; i < 8; i++) outPayload[idx++] = blockedAntennas[i] ? 1 : 0;
  // idx should equal payload length (SLOT_SIZE - 4)
}

// Save state to EEPROM using wear-leveling; writes only when state changed
void saveStateToEEPROM() {
  // Build current payload
  uint8_t payloadLen = SLOT_SIZE - 4;
  uint8_t payload[SLOT_SIZE - 4];
  packPayload(payload);

  // Compute checksum of payload
  uint8_t checksum = calcPayloadChecksum(payload, payloadLen);

  // Find latest slot/version
  uint16_t bestVersion = 0;
  bool found = false;
  uint8_t existingPayload[SLOT_SIZE - 4];

  for (uint8_t s = 0; s < SLOT_COUNT; s++) {
    uint8_t buf[SLOT_SIZE];
    if (!readSlot(s, buf)) continue;
    if (buf[0] != SLOT_MAGIC) continue;
    uint16_t ver = (uint16_t)buf[1] | ((uint16_t)buf[2] << 8);
    uint8_t csum = buf[3];
    // validate
    uint8_t comp = calcPayloadChecksum(&buf[4], payloadLen);
    if (csum != comp) continue;
    if (!found || ver > bestVersion) {
      bestVersion = ver;
      memcpy(existingPayload, &buf[4], payloadLen);
      found = true;
    }
  }

  // If found and payload identical, skip writing
  if (found) {
    bool same = true;
    for (uint8_t i = 0; i < payloadLen; i++) {
      if (existingPayload[i] != payload[i]) { same = false; break; }
    }
    if (same) return; // no change -> no write
  }

  // Determine next slot (rotate)
  uint8_t nextSlot = 0;
  if (!found) {
    nextSlot = 0;
    bestVersion = 0;
  } else {
    // pick slot that currently holds bestVersion, then advance one
    // We'll scan to find the slot index of bestVersion (first occurrence)
    bool foundSlotIdx = false;
    uint8_t bestSlotIdx = 0;
    for (uint8_t s = 0; s < SLOT_COUNT; s++) {
      uint8_t buf[SLOT_SIZE];
      if (!readSlot(s, buf)) continue;
      if (buf[0] != SLOT_MAGIC) continue;
      uint16_t ver = (uint16_t)buf[1] | ((uint16_t)buf[2] << 8);
      uint8_t csum = buf[3];
      uint8_t comp = calcPayloadChecksum(&buf[4], payloadLen);
      if (csum != comp) continue;
      if (ver == bestVersion) {
        bestSlotIdx = s;
        foundSlotIdx = true;
        break;
      }
    }
    if (foundSlotIdx) nextSlot = (bestSlotIdx + 1) % SLOT_COUNT;
    else nextSlot = 0;
  }

  // Prepare slot buffer
  uint8_t slotBuf[SLOT_SIZE];
  memset(slotBuf, 0xFF, SLOT_SIZE);
  slotBuf[0] = SLOT_MAGIC;
  uint16_t newVersion = (uint16_t)(bestVersion + 1);
  slotBuf[1] = (uint8_t)(newVersion & 0xFF);
  slotBuf[2] = (uint8_t)((newVersion >> 8) & 0xFF);
  slotBuf[3] = checksum;
  // copy payload
  for (uint8_t i = 0; i < payloadLen; i++) slotBuf[4 + i] = payload[i];

  // Write using update (only bytes that changed will be written)
  writeSlot(nextSlot, slotBuf);
}

// ------------------------
// Keypad event for long press
// ------------------------
void keypadEvent(KeypadEvent key) {
  switch (keypad.getState()) {
    case HOLD:
      if (key >= '1' && key <= '8') {
        int antenna = key - '1';
        if (antenna >= 0 && antenna < 8) {
          toggleAntennaBlock(antenna);
        }
      }
      break;
  }
}


// ------------------------
// Setup
// ------------------------
void setup() {
  
  pinMode(latchPin1, OUTPUT);
  pinMode(clockPin1, OUTPUT);
  pinMode(dataPin1, OUTPUT);
  pinMode(latchPin2, OUTPUT);
  pinMode(clockPin2, OUTPUT);
  pinMode(dataPin2, OUTPUT);
  pinMode(pttPin, INPUT);
  pinMode(statusLedPin, OUTPUT);
  pinMode(swapLedPin, OUTPUT);

  digitalWrite(latchPin1, LOW);
  digitalWrite(clockPin1, LOW);
  digitalWrite(dataPin1, LOW);
  digitalWrite(latchPin2, LOW);
  digitalWrite(clockPin2, LOW);
  digitalWrite(dataPin2, LOW);

  pinMode(txScanLedPin, OUTPUT);
  digitalWrite(txScanLedPin, LOW);
  pinMode(MaxAntLedPin, OUTPUT);
  digitalWrite(MaxAntLedPin, LOW);

  keypad.addEventListener(keypadEvent);
  keypad.setHoldTime(3000);

  // Load state from EEPROM (wear-leveled)
  delay(100);
  loadStateFromEEPROM();
  
  updateAntennas();
  updateMaxAntLed();

  delay(2000);
}

// ------------------------
// Loop and other functions
// ------------------------
void loop() {
  handleKeys();
  checkPTT();
  updateAntennas();

  bool currentPtt = digitalRead(pttPin);
  int currentRx = (rxScanMode || txScanMode) ? rxRelay : activeRelay;
  int currentTx = (rxScanMode || txScanMode) ? txRelay : activeRelay;

  if (currentRx != lastRxRelay || currentTx != lastTxRelay ||
      currentPtt != lastPttState || rxScanMode != lastRxScanMode ||
      txScanMode != lastTxScanMode || swapMode != lastSwapMode) {

    lastRxRelay = currentRx;
    lastTxRelay = currentTx;
    lastPttState = currentPtt;
    lastRxScanMode = rxScanMode;
    lastTxScanMode = txScanMode;
    lastSwapMode = swapMode;
  }
}

void blinkTxIndicator(int antenna) {
  for (int i = 0; i < 2; i++) {
    digitalWrite(latchPin2, LOW);
    shiftOut(dataPin2, clockPin2, MSBFIRST, 1 << antenna);
    digitalWrite(latchPin2, HIGH);
    delay(300);

    digitalWrite(latchPin2, LOW);
    shiftOut(dataPin2, clockPin2, MSBFIRST, 0);
    digitalWrite(latchPin2, HIGH);
    delay(300);
  }
  updateLatch2();
}

void handleKeys() {
  static String inputBuffer = "";
  static unsigned long zeroPressTime = 0;
  static unsigned long lastStarPressTime = 0;

  char key = keypad.getKey();
  if (!key) {
    if (inputBuffer == "00" && (millis() - zeroPressTime > 2000)) {
      inputBuffer = "";
    }
    return;
  }

  if (inputBuffer == "00" && isDigit(key)) {
    processKey(key, inputBuffer, zeroPressTime, lastStarPressTime);
    return;
  }

  if (digitalRead(pttPin) == HIGH) {
    if (key) {
      delay(1000);
    }
    return;
  }

  if (key >= '1' && key <= '8') {
    int antenna = key - '1';
    if (antenna >= 0 && antenna < maxAntennas) {
      if (rxScanMode || txScanMode) {
        setScanAntenna(antenna);
      } else {
        setActiveAntenna(antenna);
      }
    }
    if (inputBuffer == "0") inputBuffer = "";
    return;
  } else {
    processKey(key, inputBuffer, zeroPressTime, lastStarPressTime);
  }
}

void processKey(char key, String &inputBuffer, unsigned long &zeroPressTime, unsigned long &lastStarPressTime) {
  if (key == '0') {
    if (inputBuffer == "") {
      inputBuffer = "0";
    } else if (inputBuffer == "0") {
      inputBuffer = "00";
      zeroPressTime = millis();
    }
    return;
  }

  if (inputBuffer == "00" && isDigit(key)) {
    int newMax = key - '0';
    if (newMax >= 1 && newMax <= 8) {
      maxAntennas = newMax;
      updateMaxAntLed();
      delay(1000);
      if (activeRelay >= maxAntennas) activeRelay = maxAntennas - 1;
      if (rxRelay >= maxAntennas) rxRelay = maxAntennas - 1;
      if (txRelay >= maxAntennas) txRelay = maxAntennas - 1;
      updateAntennas();
      saveStateToEEPROM();
    }
    inputBuffer = "";
    return;
  }

  if (inputBuffer == "00" && !isDigit(key)) {
    inputBuffer = "";
  }

  if (key == '9') {
    toggleSwapMode();
  } else if (key == '*') {
    unsigned long now = millis();
    if (now - lastStarPressTime < 500) {
      toggleTxScanMode();
    } else {
      toggleRxScanMode();
    }
    lastStarPressTime = now;
  } else if (key == '#') {
    toggleTxScanMode();
  }
}

void toggleAntennaBlock(int antennaNumber) {
  if (antennaNumber < 0 || antennaNumber >= 8) return;

  blockedAntennas[antennaNumber] = !blockedAntennas[antennaNumber];

  blinkTxIndicator(antennaNumber);

  if (blockedAntennas[antennaNumber] && txRelay == antennaNumber) {
    findNewTxAntenna();
  }
  
  saveStateToEEPROM();
}

void findNewTxAntenna() {
  for (int i = 0; i < maxAntennas; i++) {
    if (!blockedAntennas[i]) {
      txRelay = i;
      updateLatch2();
      saveStateToEEPROM();
      return;
    }
  }
  txRelay = 0;
  updateLatch2();
  saveStateToEEPROM();
}

bool canUseAsTx(int antennaNumber) {
  if (antennaNumber < 0 || antennaNumber >= 8) return false;
  return !blockedAntennas[antennaNumber];
}

void toggleRxScanMode() {
  bool cameFromTxScan = false;
  if (txScanMode) { 
    rxRelay = txRelay; 
    txScanMode = false; 
    cameFromTxScan = true; 
  }
  rxScanMode = !rxScanMode;
  digitalWrite(statusLedPin, rxScanMode ? HIGH : LOW);

  if (rxScanMode) {
    if (!cameFromTxScan) {
      txRelay = activeRelay;
      rxRelay = activeRelay;
    }
    swapMode = false;
    digitalWrite(swapLedPin, LOW);
    updateLatch2();
  } else {
    activeRelay = rxRelay;
    swapMode = false;
    digitalWrite(swapLedPin, LOW);
    digitalWrite(latchPin2, LOW);
    shiftOut(dataPin2, clockPin2, MSBFIRST, 0);
    digitalWrite(latchPin2, HIGH);
  }
  saveStateToEEPROM();
}

void toggleTxScanMode() {
  bool cameFromRxScan = false;
  if (rxScanMode) { 
    txRelay = rxRelay; 
    rxScanMode = false; 
    cameFromRxScan = true; 
    digitalWrite(statusLedPin, LOW);
  }
  txScanMode = !txScanMode;

  if (txScanMode) {
    if (!cameFromRxScan) {
      rxRelay = activeRelay;
      txRelay = activeRelay;
    }
    swapMode = false;
    digitalWrite(swapLedPin, LOW);
    updateLatch2();
  } else {
    activeRelay = txRelay;
    swapMode = false;
    digitalWrite(swapLedPin, LOW);
    digitalWrite(latchPin2, LOW);
    shiftOut(dataPin2, clockPin2, MSBFIRST, 0);
    digitalWrite(latchPin2, HIGH);
  }
  saveStateToEEPROM();
}

void toggleSwapMode() {
  if (!(rxScanMode || txScanMode)) return;
  if (!swapMode) {
    if (canUseAsTx(rxRelay) && canUseAsTx(txRelay)) {
      swapMode = true;
      int temp = rxRelay;
      rxRelay = txRelay;
      txRelay = temp;
      digitalWrite(swapLedPin, HIGH);
    } else {
      blinkTxIndicator(txRelay);
      delay(1200);
      return;
    }
  } else {
    swapMode = false;
    int temp = rxRelay;
    rxRelay = txRelay;
    txRelay = temp;
    digitalWrite(swapLedPin, LOW);
  }
  saveStateToEEPROM();
  updateLatch2();
}

void setScanAntenna(int antennaNumber) {
  if (antennaNumber < 0 || antennaNumber >= maxAntennas) return;
  if (swapMode) return;

  bool stateChanged = false;
  
  if (rxScanMode) {
    if (rxRelay != antennaNumber) {
      rxRelay = antennaNumber;
      stateChanged = true;
    }
  } else if (txScanMode) {
    if (canUseAsTx(antennaNumber)) {
      if (txRelay != antennaNumber) {
        txRelay = antennaNumber;
        updateLatch2();
        stateChanged = true;
      }
    } else {
      blinkTxIndicator(antennaNumber);
      delay(1000);
      return;
    }
  }
  
  if (stateChanged) {
    saveStateToEEPROM();
  }
}

void setActiveAntenna(int antennaNumber) {
  if (antennaNumber < 0 || antennaNumber >= maxAntennas) return;
  
  if (activeRelay != antennaNumber) {
    activeRelay = antennaNumber;
    if (!rxScanMode && !txScanMode) {
      txRelay = activeRelay;
      rxRelay = activeRelay;
    }
    saveStateToEEPROM();
  }
}

void checkPTT() {
  static bool lastPttState = LOW;
  bool currentPttState = digitalRead(pttPin);
  if (currentPttState != lastPttState) {
    lastPttState = currentPttState;
    updateAntennas();
  }
}

void updateAntennas() {
  bool pttActive = digitalRead(pttPin);

  int currentRx = (rxScanMode || txScanMode) ? rxRelay : activeRelay;
  int currentTx;

  if (rxScanMode || txScanMode) {
    if (blockedAntennas[txRelay]) {
      currentTx = currentRx;
    } else {
      currentTx = txRelay;
    }
  } else {
    currentTx = txRelay;
  }

  int displayAntenna = (rxScanMode || txScanMode) ? (pttActive ? currentTx : currentRx) : activeRelay;

  digitalWrite(latchPin1, LOW);
  shiftOut(dataPin1, clockPin1, MSBFIRST, 1 << displayAntenna);
  digitalWrite(latchPin1, HIGH);

  updateLatch2();
  digitalWrite(txScanLedPin, txScanMode ? HIGH : LOW);
}

void updateLatch2() {
  byte latch2Pattern = 0;
  bool pttActive = digitalRead(pttPin);

  int currentTx = -1;

  if (rxScanMode || txScanMode) {
    currentTx = blockedAntennas[txRelay] ? rxRelay : txRelay;
  } else {
    if (blockedAntennas[activeRelay]) {
      if (pttActive) {
        for (int i = activeRelay + 1; i < maxAntennas; i++) {
          if (!blockedAntennas[i]) {
            currentTx = i;
            break;
          }
        }
        if (currentTx == -1) {
          for (int i = 0; i < activeRelay; i++) {
            if (!blockedAntennas[i]) {
              currentTx = i;
              break;
            }
          }
        }
      }
    } else {
      if (pttActive) currentTx = activeRelay;
    }
  }

  if (currentTx >= 0) {
    latch2Pattern = 1 << currentTx;
  }

  digitalWrite(latchPin2, LOW);
  shiftOut(dataPin2, clockPin2, MSBFIRST, latch2Pattern);
  digitalWrite(latchPin2, HIGH);
}

void updateMaxAntLed() {
  digitalWrite(MaxAntLedPin, maxAntennas < 8 ? HIGH : LOW);
}