#include <Wire.h>
#include <RTClib.h>
#include <NeoSWSerial.h>
#include <ModbusMaster.h>
#include "util.h"
#include "register_map_vsd.h"
#include <SPI.h>
#include <SdFat.h>

#define SD_CS_PIN 10
#define DE_RE_PIN 4
#define RX_PIN 8
#define TX_PIN 7
#define SLAVE_ID 1
#define SERIAL_BAUDRATE 9600
#define MODBUS_SERIAL_BAUDRATE 19200
#define LED_A_PID 3
#define LED_B_PID 5
#define MAX_RETRIES 1
#define ERROR_VALUE -999.99

#define SPI_CLOCK SD_SCK_MHZ(50)
#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

RTC_DS3231 rtc;
SdFat32 sd;
File dataFile;
NeoSWSerial modbusSerial(RX_PIN, TX_PIN);
ModbusMaster node;

unsigned long lastRefreshTime = 0;
bool headerWritten = false;
bool booted = false;

void flicker(uint8_t pin, uint8_t times, uint16_t speed) {
  while(times--) {
    delay(speed);
    digitalWrite(pin, HIGH);
    delay(speed);
    digitalWrite(pin, LOW);
  }
}

void setup() {
  booted = false;
  pinMode(LED_A_PID, OUTPUT);
  pinMode(LED_B_PID, OUTPUT);
  digitalWrite(LED_A_PID, LOW);
  digitalWrite(LED_B_PID, HIGH);

  Serial.begin(SERIAL_BAUDRATE);
  Serial.println(F("Startup \n"));

  if (!rtc.begin()) {
    Serial.println(F("Couldn't find RTC\n"));
    flicker(LED_B_PID, 4, 1000);
    digitalWrite(LED_B_PID, HIGH);
    digitalWrite(LED_A_PID, HIGH);
    return;
  }

  if (rtc.lostPower()) {
    Serial.println(F("RTC lost power, let's set the time!\n"));
    rtc.adjust(DateTime(F(__DATE__), F(__TIME__)));
    flicker(LED_B_PID, 4, 500);
  }

  pinMode(SD_CS_PIN, OUTPUT);
  if (!sd.begin(SD_CONFIG)) {
    flicker(LED_B_PID, 2, 1000);
    digitalWrite(LED_B_PID, HIGH);
    sd.initErrorHalt(&Serial);
    return;
  }

  pinMode(DE_RE_PIN, OUTPUT);
  digitalWrite(DE_RE_PIN, LOW);
  modbusSerial.begin(MODBUS_SERIAL_BAUDRATE);

  node.begin(SLAVE_ID, modbusSerial);
  node.preTransmission(preTransmission);
  node.postTransmission(postTransmission);
  flicker(LED_B_PID, 10, 100);
  digitalWrite(LED_B_PID, LOW);
  booted = true;
}

void preTransmission() {
  digitalWrite(DE_RE_PIN, HIGH);
  digitalWrite(LED_A_PID, HIGH);
}

void postTransmission() {
  digitalWrite(DE_RE_PIN, LOW);
  digitalWrite(LED_A_PID, LOW);
}

void writeDateTime(File &file) {
  DateTime now = rtc.now();
  file.print('\n');
  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(',');
}

void getFilename(char* buffer) {
  DateTime now = rtc.now();
  sprintf(buffer, "pm8k_%d%02d%02d.csv", now.year(), now.month(), now.day());
}

float readRegisterWithRetry(uint16_t addr, uint8_t regtype) {
  for(uint8_t retry = 0; retry < MAX_RETRIES; retry++) {
    delay(5);
    uint8_t result = node.readHoldingRegisters(addr - 1, 2);
    if(result == node.ku8MBSuccess) {
      switch(regtype) {
        case 1:
          return node.getResponseBuffer(0);
        case 2:
          return getRegisterFloat(node.getResponseBuffer(0), node.getResponseBuffer(1));
        case 3:
          return getRegisterInt64(node.getResponseBuffer(0), node.getResponseBuffer(1),
                                node.getResponseBuffer(2), node.getResponseBuffer(3));
      }
    }
    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));
    flicker(LED_B_PID, 1, 50);
  }
  return ERROR_VALUE;
}

void writeHeader() {
  if (!headerWritten) {
    dataFile.print("\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("@");
      dataFile.print(regaddr);
      dataFile.print(",");
    }
    headerWritten = true;
    flicker(LED_A_PID, 50, 10);
  }
}

void loop() {
  if (!booted) {
    delay(10000);
    digitalWrite(LED_A_PID, LOW);
    return;
  }

  if (millis() - lastRefreshTime >= 1000) {
    lastRefreshTime += 1000;
    char filename[20];
    getFilename(filename);

    if (!dataFile.open(filename, FILE_WRITE)) {
      flicker(LED_B_PID, 6, 500);
      return;
    }

    writeHeader();
    writeDateTime(dataFile);
    
    const uint16_t totalReg = sizeof(registers) / sizeof(registers[0]);
    float baseValues[4] = {ERROR_VALUE, ERROR_VALUE, ERROR_VALUE, ERROR_VALUE};
    uint8_t errorCount = 0;
    
    // Single pass for both reading and processing
    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);
      const float scale = pgm_read_float(&registers[i].scale);
      float value = ERROR_VALUE;

      if (regtype <= 3 && regaddr > 0) {
        value = readRegisterWithRetry(regaddr, regtype);
        if (value == ERROR_VALUE) {
          errorCount++;
          if (errorCount > 5) {
            dataFile.close();
            flicker(LED_B_PID, 10, 100);
            return;
          }
        }
        if (i < 4) baseValues[i] = value;
      } else {
        bool validBase = true;
        for(uint8_t j = 0; j < 4; j++) {
          if (baseValues[j] == ERROR_VALUE) {
            validBase = false;
            break;
          }
        }
        
        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;
          }
        }
      }
      
      if (value != ERROR_VALUE) {
        value *= scale;
      }
      dataFile.print(value);
      dataFile.print(',');
    }

    dataFile.close();
    
     
    if (errorCount > 0) {
      Serial.print(F("Cycle completed with "));
      Serial.print(errorCount);
      Serial.println(F(" errors"));
      flicker(LED_B_PID, errorCount, 200);
    } else {
      Serial.println(F("Cycle completed successfully"));
      flicker(LED_A_PID, 4, 100);
    }
    
    if (errorCount > 5) {
      Serial.println(F("Too many errors, aborting cycle"));
      dataFile.close();
      return;
    }
  }
}