#include <Arduino.h>

// --- Pin Map ---
#define PIN_BAT_SENSE  PA0   // Battery (3.6V nominal)
#define PIN_5V_SENSE   PA4   // Powerboost 5V Rail
#define PIN_IGN_SENSE  PA5   // Extra Signal: 12V Ignition
#define PIN_PWR_CTRL   PA1   // MOSFET to Power SimpleRTK2B
#define PIN_LED        PC13  // Onboard LED (Active LOW)
#define PIN_EXT_LED    PB0   // External LED (via MOSFET)

// --- Constants & Thresholds ---
const float ADC_REF = 3.3;
const float DIVIDER_RATIO = 2.0; 
const float BAT_CRITICAL = 3.4;
const float BAT_RECOVERY = 3.7; 

// --- Timing (ms) ---
const unsigned long SOFT_RESET_DELAY = 300000;   // 5 mins
const unsigned long HARD_RESET_DELAY = 600000;   // 10 mins total
const unsigned long BOOT_STABILIZE_TIME = 15000; // Time for 4G/GNSS to boot
const unsigned long POWER_CYCLE_TIME = 5000;     // 5 seconds to drain caps on Hard Reset

// --- UBX Commands ---
// UBX-CFG-RATE: measRate = 1000ms, navRate = 1 (1Hz Output)
const byte UBX_1HZ[] = {0xB5, 0x62, 0x06, 0x08, 0x06, 0x00, 0xE8, 0x03, 0x01, 0x00, 0x01, 0x00, 0x01, 0x39}; 
// UBX-CFG-RATE: measRate = 1000ms, navRate = 30 (30s Output)
const byte UBX_30S[] = {0xB5, 0x62, 0x06, 0x08, 0x06, 0x00, 0xE8, 0x03, 0x1E, 0x00, 0x01, 0x00, 0x1E, 0xAD};
// UBX-CFG-RST: Hot Start / Software Reset
const byte UBX_SOFT_RESET[] = {0xB5, 0x62, 0x06, 0x04, 0x04, 0x00, 0x00, 0x00, 0x01, 0x00, 0x0F, 0x66};

// --- State Variables ---
enum SystemState { BROWNOUT, BOOTING, RUNNING, HARD_RESETTING };
SystemState currentState = BROWNOUT;

int fixQuality = 0; 
int satellites = 0; 
unsigned long rxCount = 0; 
float vBatF = 3.7, v5vF = 5.0; 
unsigned long lastDashUpdate = 0;
unsigned long bootTimer = 0;
unsigned long resetCooldownTimer = 0;
unsigned long brownoutTimer = 0; 

bool isSlowTxRate = false; 
bool ntripWasWorking = false;
unsigned long failureTimer = 0;
bool softResetDone = false;
bool newDataReceived = false; 

// Optimized: Using const char* prevents memory fragmentation on ARM MCUs
const char* sysMsg = "System Initialized."; 
String nmeaBuffer = ""; 

HardwareSerial GNSS(PA3, PA2); 

// --- UTILITIES ---

// Formats uptime safely into a character buffer
void getUptimeString(char* buffer, size_t maxLen) {
  unsigned long seconds = millis() / 1000;
  int d = seconds / 86400;
  int h = (seconds % 86400) / 3600;
  int m = (seconds % 3600) / 60;
  snprintf(buffer, maxLen, "%dd%dh%dm", d, h, m);
}

void updateSensors() {
  float rawBat = (analogRead(PIN_BAT_SENSE) / 4095.0) * ADC_REF * DIVIDER_RATIO;
  float raw5v  = (analogRead(PIN_5V_SENSE) / 4095.0) * ADC_REF * DIVIDER_RATIO;
  vBatF = (vBatF * 0.9) + (rawBat * 0.1); // Low-pass filter for stability
  v5vF  = (v5vF * 0.9) + (raw5v * 0.1);
}

void setPower(bool ON) {
  if (ON) {
    pinMode(PA2, OUTPUT); 
    pinMode(PA3, INPUT);  
    digitalWrite(PIN_PWR_CTRL, HIGH);
    GNSS.begin(115200);
    bootTimer = millis();
    currentState = BOOTING;
    fixQuality = 0;
    satellites = 0;
    sysMsg = "POWER ON: Booting...";
  } else {
    digitalWrite(PIN_PWR_CTRL, LOW);
    GNSS.end();
    pinMode(PA2, INPUT); 
    pinMode(PA3, INPUT); 
    ntripWasWorking = false;
    sysMsg = "POWER OFF: Standby";
  }
}

void parseNMEA() {
  while (GNSS.available()) {
    char c = GNSS.read();
    if (c == '$') { nmeaBuffer = "$"; continue; }
    if (c == '\n') {
      if (nmeaBuffer.startsWith("$GNGGA") || nmeaBuffer.startsWith("$GPGGA")) {
        int commaIndex = 0;
        int satStart = -1;
        for (unsigned int i = 0; i < nmeaBuffer.length(); i++) {
          if (nmeaBuffer[i] == ',') {
            commaIndex++;
            if (commaIndex == 6 && (i + 1) < nmeaBuffer.length()) { 
              fixQuality = nmeaBuffer[i + 1] - '0';
            } 
            else if (commaIndex == 7 && (i + 1) < nmeaBuffer.length()) {
              satStart = i + 1;
            } 
            else if (commaIndex == 8 && satStart != -1) {
              satellites = nmeaBuffer.substring(satStart, i).toInt();
              break; 
            }
          }
        }
        rxCount++;
        newDataReceived = true; 
      }
      nmeaBuffer = ""; 
    } else if (c != '\r' && nmeaBuffer.length() < 100) {
      nmeaBuffer += c;
    }
  }
}

// Local Dashboard (ANSI Colors for PuTTY/TeraTerm)
void drawDashboard(bool pa4_on, bool pa5_on) {
  char uptimeBuf[16];
  getUptimeString(uptimeBuf, sizeof(uptimeBuf));

  Serial.print("\033[H"); // Reset cursor to top-left
  Serial.println("===========================================");
  Serial.println("    RTK BLACKPILL TELEMETRY SYSTEM       ");
  Serial.println("===========================================");
  
  // 1. System State
  Serial.print("SYSTEM STATE: ");
  if (currentState == RUNNING)        Serial.println("\033[1;32mRUNNING       \033[0m");
  else if (currentState == BOOTING)   Serial.println("\033[1;33mBOOTING       \033[0m");
  else if (currentState == HARD_RESETTING) Serial.println("\033[1;31mPWR CYCLING   \033[0m");
  else                                Serial.println("\033[1;35mBROWNOUT      \033[0m");

  // 2. Hardware Monitors
  Serial.print("BATTERY:      "); Serial.print(vBatF, 2); Serial.println(" V  ");
  Serial.print("UPTIME:       "); Serial.println(uptimeBuf);
  
  Serial.print("PWR IN (PA4): ");
  if (pa4_on) Serial.println("\033[1;32mACTIVE (>4.0V)\033[0m  ");
  else        Serial.println("\033[1;37mOFF (BATTERY)\033[0m   ");

  Serial.print("IGN IN (PA5): ");
  if (pa5_on) Serial.println("\033[1;32mACTIVE (HIGH)\033[0m   ");
  else        Serial.println("\033[1;37mOFF (LOW)\033[0m        ");

  // 3. Update Rate (Restored)
  Serial.print("GNSS TX RATE: ");
  if (currentState == BOOTING || currentState == BROWNOUT) {
    Serial.println("\033[1;37mOFFLINE         \033[0m");
  } else if (isSlowTxRate) {
    Serial.println("\033[1;33mEVERY 30 SECONDS\033[0m");
  } else {
    Serial.println("\033[1;36m1Hz (1 SECOND)  \033[0m");
  }

  // 4. GNSS Status
  Serial.print("RTK FIX:      ");
  if (fixQuality == 4)      Serial.println("\033[1;32m[4] FIXED \033[0m    ");
  else if (fixQuality == 5) Serial.println("\033[1;36m[5] FLOAT \033[0m    ");
  else if (fixQuality > 0)  Serial.println("\033[1;33m[1] 3D FIX\033[0m    ");
  else                      Serial.println("\033[1;31m[0] NO FIX\033[0m    ");

  Serial.print("PACKETS TX'd: "); Serial.print(rxCount); Serial.println("        ");
  
  // 5. System Logs
  Serial.println("-------------------------------------------");
  Serial.print("LAST EVENT:   "); 
  
  // Padding the const char* message to cleanly overwrite previous longer messages
  String paddedMsg = String(sysMsg);
  while(paddedMsg.length() < 30) paddedMsg += " ";
  Serial.println(paddedMsg); 
  Serial.println("===========================================");
}

// Visual State Indicator
void updateLEDs() {
  unsigned long m = millis();
  bool ext = false; 
  switch (currentState) {
    case BROWNOUT: 
      ext = (m % 2000 < 50); // Heartbeat pulse (Sleep)
      break;
    
    case BOOTING:
      ext = (m % 500 < 250); // Steady Medium Blink (Booting GNSS)
      break;

    case HARD_RESETTING: 
      ext = (m % 100 < 50);  // Aggressive Strobe (Fatal Error / Cycle)
      break;
      
    case RUNNING:
      if (fixQuality == 4) {
        ext = true; // Solid ON (RTK Fixed - Perfect)
      } 
      else if (fixQuality > 0) {
        unsigned long cycle = m % 1000;
        ext = (cycle < 100) || (cycle > 200 && cycle < 300); // Double blink (Float/3D)
      } 
      else {
        ext = (m % 1000 < 500); // Steady Slow Blink (Searching)
      }
      break;
  }
  
  digitalWrite(PIN_EXT_LED, ext);
  digitalWrite(PIN_LED, !ext); // Inverted for onboard Active-Low LED
}

// --- SETUP ---
void setup() {
  pinMode(PIN_PWR_CTRL, OUTPUT);
  pinMode(PIN_LED, OUTPUT);
  pinMode(PIN_EXT_LED, OUTPUT);
  pinMode(PIN_IGN_SENSE, INPUT_PULLDOWN); 
  
  Serial.begin(115200);
  analogReadResolution(12);
  
  // Reserving memory prevents heap fragmentation
  nmeaBuffer.reserve(100);

  Serial.print("\033[2J\033[H"); // Clear entire screen in PuTTY on boot
  setPower(false);
  currentState = BROWNOUT; 
}

// --- MAIN LOOP ---
void loop() {
  updateSensors();
  parseNMEA();
  updateLEDs();

  bool pa4_on = (v5vF > 4.0); 
  bool pa5_on = (digitalRead(PIN_IGN_SENSE) == HIGH); 
  
  // Rate Truth Table: 1Hz only when both 5V and Ignition are active
  bool fastTxMode = (pa4_on && pa5_on); 

  // Dashboard Update
  if (newDataReceived || millis() - lastDashUpdate > 2000) {
    drawDashboard(pa4_on, pa5_on);
    lastDashUpdate = millis();
    newDataReceived = false;
  }

  // Brownout Check Failsafe
  if (vBatF < BAT_CRITICAL && !pa4_on && !pa5_on) {
    if (brownoutTimer == 0) brownoutTimer = millis();
    if (millis() - brownoutTimer > 3000) {
      if (currentState != BROWNOUT && currentState != HARD_RESETTING) {
        sysMsg = "ERR: Bat Critical. Sleeping.";
        setPower(false);
        currentState = BROWNOUT;
      }
    }
  } else { brownoutTimer = 0; }

  // State Machine
  switch (currentState) {
    case BROWNOUT:
      if (pa4_on || pa5_on || vBatF > BAT_RECOVERY) setPower(true);
      break;

    case BOOTING:
      if (millis() - bootTimer > BOOT_STABILIZE_TIME) {
        sysMsg = "Applying initial UBX.";
        if (fastTxMode) { 
          GNSS.write(UBX_1HZ, sizeof(UBX_1HZ)); 
          isSlowTxRate = false; 
        } else { 
          GNSS.write(UBX_30S, sizeof(UBX_30S)); 
          isSlowTxRate = true; 
        }
        currentState = RUNNING;
        failureTimer = 0; 
      }
      break;
      
    case HARD_RESETTING:
      if (millis() - resetCooldownTimer > POWER_CYCLE_TIME) currentState = BROWNOUT; 
      break;

    case RUNNING:
      // A. Dynamic Rate Switching
      if (fastTxMode && isSlowTxRate) {
        sysMsg = "Ignition ON: 1Hz Mode";
        GNSS.write(UBX_1HZ, sizeof(UBX_1HZ));
        isSlowTxRate = false;
      } else if (!fastTxMode && !isSlowTxRate) {
        sysMsg = "Ignition OFF: 30s Mode";
        GNSS.write(UBX_30S, sizeof(UBX_30S));
        isSlowTxRate = true;
      }

      // B. Watchdog Logic (Reset triggers if RTK drops)
      if (fixQuality >= 4) { 
        if (!ntripWasWorking) { 
          sysMsg = "Fix OK. Watchdog Armed."; 
          ntripWasWorking = true; 
        }
        failureTimer = 0; 
        softResetDone = false;
      } 
      else if (ntripWasWorking && fixQuality < 4) { 
        if (failureTimer == 0) failureTimer = millis();
        unsigned long elapsed = millis() - failureTimer;
        
        // Soft Reset Trigger
        if (elapsed > SOFT_RESET_DELAY && !softResetDone) {
          sysMsg = "WARN: Fix Lost. Soft Reset...";
          GNSS.write(UBX_SOFT_RESET, sizeof(UBX_SOFT_RESET));
          softResetDone = true;
        }
        
        // Hard Reset Trigger
        if (elapsed > HARD_RESET_DELAY) {
          sysMsg = "FATAL: Hard Cycle...";
          setPower(false); 
          currentState = HARD_RESETTING;
          resetCooldownTimer = millis();
        }
      }
      break;
  }
}