#include <Wire.h>              // Library for I2C communication (used by RTC)
#include <RTClib.h>           // Library for Real Time Clock
#include <ModbusMaster.h>     // Library for Modbus communication
#include "util.h"             // Custom utility functions
#include "register_map.h"     // Map of Modbus registers to read
#include <SPI.h>              // Library for SPI communication (used by SD card)
#include <SdFat.h>           // Enhanced SD card library

// ==== PIN CONNECTIONS AND SETTINGS ====
/*
Physical Connections Guide:

SD CARD MODULE:
- CS (Chip Select) -> Arduino MEGA pin 53
- MOSI -> Arduino MEGA pin 51
- MISO -> Arduino MEGA pin 50
- SCK -> Arduino MEGA pin 52
- VCC -> 5V
- GND -> GND

RS485 MODULE:
- DI (Data In) -> Arduino MEGA TX1 (Pin 18)
- RO (Receive Out) -> Arduino MEGA RX1 (Pin 19)
- DE & RE (Data/Receive Enable) -> Arduino MEGA Pin 4
- VCC -> 5V
- GND -> GND
- A & B -> To Modbus device (polarity sensitive)

RTC MODULE (DS3231):
- SDA -> Arduino MEGA Pin 20
- SCL -> Arduino MEGA Pin 21
- VCC -> 5V
- GND -> GND

STATUS LEDs:
- LED A (Activity) -> Arduino MEGA Pin 3 (blinks during normal operation)
- LED B (Error) -> Arduino MEGA Pin 5 (blinks during errors)
*/

// ==== CONFIGURATION SETTINGS ====
#define SD_CS_PIN 53          // SD card chip select pin (uses MEGA's default SS pin)
#define DE_RE_PIN 4           // Controls RS485 direction (transmit/receive switching)
#define SLAVE_ID 1            // Modbus device address (change to match your device)
#define SERIAL_BAUDRATE 115200     // Speed for debug messages via USB

#define MODBUS_SERIAL_BAUDRATE 9600  // Speed for Modbus communication
#define LED_A_PIN 3           // Activity LED (blinks during normal operation)
#define LED_B_PIN 5           // Error LED (blinks when problems occur)
#define MAX_RETRIES 1         // Number of times to retry failed readings
#define ERROR_VALUE -999.99   // Value used to indicate reading errors

// ==== SD CARD CONFIGURATION ====
// Sets up the SD card for optimal performance with MEGA
#define SPI_CLOCK SD_SCK_MHZ(16)  // //50 (fast) //16 (half) //4 (slow) SD card speed (50MHz)
// Choose the best SD card mode based on hardware capabilities
#if HAS_SDIO_CLASS
#define SD_CONFIG SdioConfig(FIFO_SDIO)  // Use SDIO if available (faster)
#elif ENABLE_DEDICATED_SPI
#define SD_CONFIG SdSpiConfig(SD_CS_PIN, DEDICATED_SPI, SPI_CLOCK)  // Dedicated SPI bus
#else
#define SD_CONFIG SdSpiConfig(SD_CS_PIN, SHARED_SPI, SPI_CLOCK)     // Shared SPI bus
#endif

// ==== GLOBAL OBJECTS ====
RTC_DS3231 rtc;              // Real Time Clock object for timekeeping
SdFat32 sd;                  // SD card object for file operations
File dataFile;               // File object for data logging
ModbusMaster node;           // Modbus communication object

// ==== GLOBAL VARIABLES ====
unsigned long lastRefreshTime = 0;  // Tracks time between readings
bool headerWritten = false;         // Tracks if CSV header has been written
bool booted = false;                // Tracks if setup completed successfully

// ==== UTILITY FUNCTIONS ====
// Makes an LED blink a specified number of times
// pin: which LED to blink
// times: how many blinks
// speed: how fast to blink (in milliseconds)
void flicker(uint8_t pin, uint8_t times, uint16_t speed) {
  while(times--) {
    delay(speed);
    digitalWrite(pin, HIGH);  // Turn LED on
    delay(speed);
    digitalWrite(pin, LOW);   // Turn LED off
  }
}

// ==== SETUP FUNCTION ====
// Runs once when the Arduino starts or resets
// Initializes all hardware and prepares for operation
void setup() {
  booted = false;  // Mark as not ready
  
  // Setup status LEDs
  pinMode(LED_A_PIN, OUTPUT);
  pinMode(LED_B_PIN, OUTPUT);
  digitalWrite(LED_A_PIN, LOW);     // Activity LED off
  digitalWrite(LED_B_PIN, HIGH);    // Error LED on until setup complete

  // Start serial communication for debugging via USB
  Serial.begin(SERIAL_BAUDRATE);
  Serial.println(F("Startup \n"));

  // Start serial communication for Modbus (using hardware Serial1)
  Serial1.begin(MODBUS_SERIAL_BAUDRATE);

  // Initialize Real Time Clock
  if (!rtc.begin()) {
    Serial.println(F("Couldn't find RTC\n"));
    flicker(LED_B_PIN, 4, 1000);  // Error pattern: 4 slow blinks
    digitalWrite(LED_B_PIN, HIGH);
    digitalWrite(LED_A_PIN, HIGH);
    return;  // Stop if RTC fails
  }

  // Check if RTC lost power and reset time if needed
  if (rtc.lostPower()) {
    Serial.println(F("RTC lost power, let's set the time!\n"));
    rtc.adjust(DateTime(F(__DATE__), F(__TIME__)));  // Set to compile time
    flicker(LED_B_PIN, 4, 500);  // Warning pattern: 4 medium blinks
  }

  // Initialize SD card
  pinMode(SD_CS_PIN, OUTPUT);
  // Set pin to high, hopefully fix Mega issue
  digitalWrite(SD_CS_PIN, HIGH);

  if (!sd.begin(SD_CONFIG)) {
    flicker(LED_B_PIN, 2, 1000);  // Error pattern: 2 slow blinks
    digitalWrite(LED_B_PIN, HIGH);
    sd.initErrorHalt(&Serial);
    return;  // Stop if SD card fails
  }

  // Setup RS485 communication direction control
  pinMode(DE_RE_PIN, OUTPUT);
  digitalWrite(DE_RE_PIN, LOW);  // Start in receive mode

  // Initialize Modbus communication
  node.begin(SLAVE_ID, Serial1);  // Using Hardware Serial1 for Modbus
  node.preTransmission(preTransmission);   // Set callbacks for RS485 direction control
  node.postTransmission(postTransmission);
  
  flicker(LED_B_PIN, 10, 100);  // Success pattern: 10 quick blinks
  digitalWrite(LED_B_PIN, LOW);  // Turn off error LED
  booted = true;                 // Mark setup as complete
}

// ==== RS485 CONTROL FUNCTIONS ====
// Called before Modbus transmission begins
void preTransmission() {
  digitalWrite(DE_RE_PIN, HIGH);    // Enable transmitter
  digitalWrite(LED_A_PIN, HIGH);    // Turn on activity LED
}

// Called after Modbus transmission completes
void postTransmission() {
  digitalWrite(DE_RE_PIN, LOW);     // Enable receiver
  digitalWrite(LED_A_PIN, LOW);     // Turn off activity LED
}

// ==== FILE OPERATIONS ====
// Writes the current date and time to the CSV file
void writeDateTime(File &file) {
  DateTime now = rtc.now();  // Get current time from RTC
  file.print('\n');         // Start new line
  // Write date and time in format: YYYY-MM-DD HH:MM:SS,
  file.print(now.year(), DEC);
  file.print('-');
  file.print(now.month(), DEC);
  file.print('-');
  file.print(now.day(), DEC);
  file.print(' ');
  file.print(now.hour(), DEC);
  file.print(':');
  file.print(now.minute(), DEC);
  file.print(':');
  file.print(now.second(), DEC);
  file.print(',');
}

// Generates filename based on current date (format: pm8k_YYYYMMDD.csv)
void getFilename(char* buffer) {
  DateTime now = rtc.now();
  sprintf(buffer, "log_%d%02d%02d.csv", now.year(), now.month(), now.day());
}

// ==== MODBUS OPERATIONS ====
// Reads a Modbus register with retry capability
// addr: register address to read
// regtype: type of register (1=int, 2=float, 3=long)
// Returns: register value or ERROR_VALUE if failed
float readRegisterWithRetry(uint16_t addr, uint8_t regtype) {
  for(uint8_t retry = 0; retry < MAX_RETRIES; retry++) {
    delay(5);  // Short delay between attempts
    uint8_t result = node.readHoldingRegisters(addr - 1, 2);  // Read register
    
    if(result == node.ku8MBSuccess) {
      // Convert register value based on type
      switch(regtype) {
        case 1:  // Integer
          return node.getResponseBuffer(0);
        case 2:  // Float
          return getRegisterFloat(node.getResponseBuffer(0), node.getResponseBuffer(1));
        case 3:  // Long
          return getRegisterInt64(node.getResponseBuffer(0), node.getResponseBuffer(1),
                                node.getResponseBuffer(2), node.getResponseBuffer(3));
      }
    }
    
    // Log error if read failed
    Serial.print(F("Read error at register "));
    Serial.print(addr);
    Serial.print(F(", attempt "));
    Serial.print(retry + 1);
    Serial.print(F(" of "));
    Serial.print(MAX_RETRIES);
    Serial.print(F(", error code: "));
    Serial.println(result);
    
    delay(5 * (retry + 1));  // Increasing delay between retries
    flicker(LED_B_PIN, 1, 50);  // Quick error blink
  }
  return ERROR_VALUE;  // Return error value if all retries failed
}

// Writes the CSV header row if it hasn't been written yet
void writeHeader() {
  if (!headerWritten) {
    dataFile.print("\nDate Time,");
    // Write register addresses as column headers
    const uint16_t totalReg = sizeof(registers) / sizeof(registers[0]);
    for (uint16_t i = 0; i < totalReg; i++) {
      const uint16_t regaddr = pgm_read_word(&registers[i].regaddr);
      dataFile.print("@");
      dataFile.print(regaddr);
      dataFile.print(",");
    }
    headerWritten = true;
    flicker(LED_A_PIN, 50, 10);  // Success pattern: 50 quick blinks
  }
}

// ==== MAIN PROGRAM LOOP ====
void loop() {
  // If setup failed, wait and try again
  if (!booted) {
    delay(10000);  // Wait 10 seconds
    digitalWrite(LED_A_PIN, LOW);
    return;
  }

  // Check if it's time for next reading (every 1000ms)
  if (millis() - lastRefreshTime >= 1000) {
    lastRefreshTime += 1000;
    
    // Create new file for today's date
    char filename[20];
    getFilename(filename);

    // Try to open the data file
    if (!dataFile.open(filename, FILE_WRITE)) {
      flicker(LED_B_PIN, 6, 500);  // Error pattern: 6 medium blinks
      return;
    }

    // Write header if needed and timestamp
    writeHeader();
    writeDateTime(dataFile);
    
    // Initialize variables for reading registers
    const uint16_t totalReg = sizeof(registers) / sizeof(registers[0]);
    float baseValues[4] = {ERROR_VALUE, ERROR_VALUE, ERROR_VALUE, ERROR_VALUE};
    uint8_t errorCount = 0;
    
    // Read and process all registers
    for (uint16_t i = 0; i < totalReg; i++) {
      // Get register information
      const uint16_t regaddr = pgm_read_word(&registers[i].regaddr);
      const uint8_t regtype = pgm_read_word(&registers[i].regtype);
      const float scale = pgm_read_float(&registers[i].scale);
      float value = ERROR_VALUE;

      // Read basic register types
      if (regtype <= 3 && regaddr > 0) {
        value = readRegisterWithRetry(regaddr, regtype);
        if (value == ERROR_VALUE) {
          errorCount++;
          if (errorCount > 5) {  // Too many errors, abort
            dataFile.close();
            flicker(LED_B_PIN, 10, 100);  // Error pattern: 10 quick blinks
            return;
          }
        }
        if (i < 4) baseValues[i] = value;  // Store first 4 values for calculations
      } else {
        // Check if we have valid base values for calculations
        bool validBase = true;
        for(uint8_t j = 0; j < 4; j++) {
          if (baseValues[j] == ERROR_VALUE) {
            validBase = false;
            break;
          }
        }
        
        // Calculate derived values if base values are valid
        if (validBase) {
          switch(regtype) {
            case 4:
              value = calculateStatusWord(baseValues);
              break;
            case 5:
              value = calculateThermal(baseValues);
              break;
            case 6:
              value = calculatePower(baseValues);
              break;
            case 7:
              value = calculateRPM(baseValues);
              break;
          }
        }
      }
      
      // Apply scaling factor if value is valid
      if (value != ERROR_VALUE) {
        value *= scale;
      }
      // Write value to file
      dataFile.print(value);
      dataFile.print(',');
    }

    // Close file after writing
    dataFile.close();
    
    // Report status
    if (errorCount > 0) {
      Serial.print(F("Cycle completed with "));
      Serial.print(errorCount);
      Serial.println(F(" errors"));
      flicker(LED_B_PIN, errorCount, 200);  // Error pattern: blink count = error count
    } else {
      Serial.println(F("Cycle completed successfully"));
      flicker(LED_A_PIN, 4, 100);  // Success pattern: 4 quick blinks
    }
    
    // Abort if too many errors
    if (errorCount > 5) {
      Serial.println(F("Too many errors, aborting cycle"));
      dataFile.close();
      return;
    }
  }
}