#include <Wire.h>              // Library for I2C communication
#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
#include <SdFat.h>           // Enhanced SD card library

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

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

RS485 MODULE:
- DI -> Arduino MEGA TX1 (Pin 18)
- RO -> Arduino MEGA RX1 (Pin 19)
- DE & RE -> Arduino MEGA Pin 4
- VCC -> 5V
- GND -> GND

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

STATUS LEDs:
- LED A -> Arduino MEGA Pin 3
- LED B -> Arduino MEGA Pin 5
*/

// ==== CONFIGURATION SETTINGS ====
#define SD_CS_PIN 53          // SD card chip select pin (MEGA's SS pin)
#define DE_RE_PIN 4           // RS485 direction control
#define SLAVE_ID 101          // Modbus device address
#define SERIAL_BAUDRATE 115200    // Debug communication speed
#define MODBUS_BAUDRATE 9600  // Modbus communication speed
#define LED_A_PIN 3           // Activity LED
#define LED_B_PIN 5           // Error LED
#define MAX_RETRIES 3         // Maximum read attempts
#define ERROR_VALUE -999.99   // Error indicator value

// ==== SD CARD CONFIGURATION ====
#define SPI_CLOCK SD_SCK_MHZ(16)  // SD card speed (16MHz for stability)
#if HAS_SDIO_CLASS
#define SD_CONFIG SdioConfig(FIFO_SDIO)
#elif ENABLE_DEDICATED_SPI
#define SD_CONFIG SdSpiConfig(SD_CS_PIN, DEDICATED_SPI, SPI_CLOCK)
#else
#define SD_CONFIG SdSpiConfig(SD_CS_PIN, SHARED_SPI, SPI_CLOCK)
#endif

// ==== GLOBAL OBJECTS ====
RTC_DS3231 rtc;              // RTC object
SdFat32 sd;                  // SD card object
File dataFile;               // File object
ModbusMaster node;           // Modbus object

// ==== GLOBAL VARIABLES ====
unsigned long lastRefreshTime = 0;
bool headerWritten = false;
bool booted = false;

// ==== UTILITY FUNCTIONS ====
void flicker(uint8_t pin, uint8_t times, uint16_t speed) {
  for (uint8_t i = 0; i < times; i++) {
    digitalWrite(pin, HIGH);
    delay(speed);
    digitalWrite(pin, LOW);
    delay(speed);
  }
}

// ==== SETUP FUNCTION ====
void setup() {
  booted = false;
  
  // Initialize status LEDs
  pinMode(LED_A_PIN, OUTPUT);
  pinMode(LED_B_PIN, OUTPUT);
  digitalWrite(LED_A_PIN, LOW);
  digitalWrite(LED_B_PIN, HIGH);  // Error LED on until setup complete

  // Start debug serial communication
  Serial.begin(SERIAL_BAUDRATE);
  Serial.println(F("Startup"));

  // Start Modbus serial communication (Hardware Serial1)
  Serial1.begin(MODBUS_BAUDRATE);

  // Initialize RTC
  if (!rtc.begin()) {
    Serial.println(F("RTC initialization failed"));
    flicker(LED_B_PIN, 4, 1000);
    return;
  }

  if (rtc.lostPower()) {
    Serial.println(F("RTC lost power, setting time"));
    rtc.adjust(DateTime(F(__DATE__), F(__TIME__)));
    flicker(LED_B_PIN, 4, 500);
  }

  // Initialize SD card
  pinMode(SD_CS_PIN, OUTPUT);
  digitalWrite(SD_CS_PIN, HIGH);  // Ensure SS pin is high initially

  if (!sd.begin(SD_CONFIG)) {
    Serial.println(F("SD card initialization failed"));
    flicker(LED_B_PIN, 2, 1000);
    return;
  }

  // Initialize Modbus
  pinMode(DE_RE_PIN, OUTPUT);
  digitalWrite(DE_RE_PIN, LOW);  // Start in receive mode

  node.begin(SLAVE_ID, Serial1);
  node.preTransmission(preTransmission);
  node.postTransmission(postTransmission);

  flicker(LED_B_PIN, 10, 100);  // Setup success indicator
  digitalWrite(LED_B_PIN, LOW);
  booted = true;
}

// ==== RS485 CONTROL FUNCTIONS ====
void preTransmission() {
  digitalWrite(DE_RE_PIN, HIGH);
  digitalWrite(LED_A_PIN, HIGH);
  delayMicroseconds(50);  // Give RS485 time to switch
}

void postTransmission() {
  delayMicroseconds(50);  // Give RS485 time to switch
  digitalWrite(DE_RE_PIN, LOW);
  digitalWrite(LED_A_PIN, LOW);
}

// ==== FILE OPERATIONS ====
String getFilename() {
  DateTime now = rtc.now();
  char buffer[20];
  sprintf(buffer, "pm8k_%d%02d%02d.csv", now.year(), now.month(), now.day());
  return String(buffer);
}

void writeDateTime() {
  DateTime now = rtc.now();
  dataFile.print('\n');
  dataFile.print(now.year(), DEC);
  dataFile.print('-');
  dataFile.print(now.month(), DEC);
  dataFile.print('-');
  dataFile.print(now.day(), DEC);
  dataFile.print(' ');
  dataFile.print(now.hour(), DEC);
  dataFile.print(':');
  dataFile.print(now.minute(), DEC);
  dataFile.print(':');
  dataFile.print(now.second(), DEC);
  dataFile.print(',');
}

// ==== MODBUS OPERATIONS ====
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);
    
    if (result == node.ku8MBSuccess) {
      switch(regtype) {
        case 1:  // Integer
          return node.getResponseBuffer(0);
        case 2:  // Float
          return getRegisterFloat(node.getResponseBuffer(0), node.getResponseBuffer(1));
        case 0:  // 32-bit Integer
          return getRegisterInt32(node.getResponseBuffer(0), node.getResponseBuffer(1));
        case 5:  // String
          String str;
          for (uint8_t j = 0; j < 20; j++) {
            uint8_t v = node.getResponseBuffer(j);
            if (v == 0) break;
            str += (char)v;
          }
          return str.toFloat();
      }
    }
    
    Serial.print(F("Read error at register "));
    Serial.print(addr);
    Serial.print(F(", attempt "));
    Serial.println(retry + 1);
    
    delay(50 * (retry + 1));  // Increasing delay between retries
    flicker(LED_B_PIN, 1, 50);
  }
  return ERROR_VALUE;
}

// ==== MAIN LOOP ====
void loop() {
  if (!booted) {
    Serial.println(F("Boot failed, retrying in 10 seconds"));
    delay(10000);
    return;
  }

  if (millis() - lastRefreshTime >= 1000) {
    lastRefreshTime += 1000;
    String filename = getFilename();
    uint8_t errorCount = 0;

    if (!dataFile.open(filename.c_str(), FILE_WRITE)) {
      Serial.println(F("Failed to open file"));
      flicker(LED_B_PIN, 6, 500);
      return;
    }

    if (!headerWritten) {
      dataFile.print(F("\nDate Time,"));
      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(F("@"));
        dataFile.print(regaddr);
        dataFile.print(F(","));
      }
      headerWritten = true;
      flicker(LED_A_PIN, 3, 100);
    }

    writeDateTime();

    // Read all registers
    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);
      const uint8_t regtype = pgm_read_word(&registers[i].regtype);

      if (regaddr > 0) {
        float value = readRegisterWithRetry(regaddr, regtype);
        if (value == ERROR_VALUE) {
          errorCount++;
          if (errorCount > 5) {
            Serial.println(F("Too many errors, aborting cycle"));
            dataFile.close();
            return;
          }
        }
        dataFile.print(value);
        dataFile.print(F(","));
      }
    }

    dataFile.close();
    
    if (errorCount > 0) {
      flicker(LED_B_PIN, errorCount, 200);
    } else {
      flicker(LED_A_PIN, 1, 100);
    }
  }
}