/**
 * @file fori_bldc_control.ino
 * @brief Arduino Mega Firmware for FORI AGV BLDC Hub Motor Control
 * 
 * Communicates with two ZLAC8015D Dual-Channel Drivers via Modbus RTU (Serial1)
 * and receives speed commands from the PC (Serial) via USB.
 * 
 * Hardware Connections:
 * - Serial (USB): Communication with PC (115200 baud)
 * - Serial1: TX1 (Pin 18), RX1 (Pin 19) connected to TTL-to-RS485 Transceiver
 * - Pin 2: RS485 RE/DE (Direction control pin. High = Transmit, Low = Receive)
 * 
 * Slave Configuration:
 * - Driver 1 (Front Wheels): Slave ID 1. Ch1 = Left Front, Ch2 = Right Front
 * - Driver 2 (Rear Wheels): Slave ID 2. Ch1 = Left Rear, Ch2 = Right Rear
 */

#define RS485_DE_RE_PIN 2

// Modbus registers
#define REG_MODE 0x200D        // Operating mode: 3 = Velocity Control
#define REG_CONTROL_WORD 0x200E // Control word: 0x0008 = Enable, 0x0007 = Stop/Disable, 0x0006 = Clear Fault
#define REG_BRAKE_CH1 0x201A   // Brake Control Ch1: 0x0000 = Release, 0x0001 = Lock
#define REG_BRAKE_CH2 0x201B   // Brake Control Ch2: 0x0000 = Release, 0x0001 = Lock
#define REG_TARGET_RPM 0x2088  // Target RPM starting address (Ch1 at 0x2088, Ch2 at 0x2089)
#define REG_FEEDBACK_RPM 0x20AD // Actual feedback RPM starting address (Ch1 at 0x20AD, Ch2 at 0x20AE)

// Global state
bool drivers_enabled = false;
unsigned long last_cmd_time = 0;
const unsigned long CMD_TIMEOUT_MS = 500; // Safe timeout: stop if no commands received from PC

// Packet structures
struct PCCommand {
  int16_t left_rpm;
  int16_t right_rpm;
  uint8_t control_state; // 0 = Idle/Brake, 1 = Enable/Release Brake
};

struct DriverFeedback {
  int16_t front_left_rpm;
  int16_t front_right_rpm;
  int16_t rear_left_rpm;
  int16_t rear_right_rpm;
};

DriverFeedback feedback = {0, 0, 0, 0};

// Modbus CRC-16 Calculation
uint16_t calculateCRC(uint8_t* buf, int len) {
  uint16_t crc = 0xFFFF;
  for (int pos = 0; pos < len; pos++) {
    crc ^= (uint16_t)buf[pos];
    for (int i = 8; i != 0; i--) {
      if ((crc & 0x0001) != 0) {
        crc >>= 1;
        crc ^= 0xA001;
      } else {
        crc >>= 1;
      }
    }
  }
  return crc;
}

// Set RS485 to Transmit Mode
void setRS485Tx() {
  digitalWrite(RS485_DE_RE_PIN, HIGH);
  delayMicroseconds(50);
}

// Set RS485 to Receive Mode
void setRS485Rx() {
  delayMicroseconds(50);
  digitalWrite(RS485_DE_RE_PIN, LOW);
}

// Write Single Register via Modbus RTU
bool writeModbusSingle(uint8_t slave_id, uint16_t reg, uint16_t value) {
  uint8_t packet[8];
  packet[0] = slave_id;
  packet[1] = 0x06; // Write Single Register
  packet[2] = (reg >> 8) & 0xFF;
  packet[3] = reg & 0xFF;
  packet[4] = (value >> 8) & 0xFF;
  packet[5] = value & 0xFF;
  
  uint16_t crc = calculateCRC(packet, 6);
  packet[6] = crc & 0xFF;
  packet[7] = (crc >> 8) & 0xFF;

  setRS485Tx();
  Serial1.write(packet, 8);
  Serial1.flush();
  setRS485Rx();

  // Read response (expected 8 bytes)
  unsigned long start_time = millis();
  int idx = 0;
  uint8_t response[8];
  while (millis() - start_time < 50) {
    if (Serial1.available()) {
      response[idx++] = Serial1.read();
      if (idx >= 8) break;
    }
  }

  if (idx < 8) return false;
  uint16_t resp_crc = calculateCRC(response, 6);
  uint16_t calc_crc = (response[7] << 8) | response[6];
  return (resp_crc == calc_crc && response[0] == slave_id && response[1] == 0x06);
}

// Write Multiple Registers via Modbus RTU (e.g. Target Speed Ch1 & Ch2)
bool writeModbusMultiple(uint8_t slave_id, uint16_t start_reg, uint16_t num_regs, int16_t* values) {
  uint8_t packet[32];
  packet[0] = slave_id;
  packet[1] = 0x10; // Write Multiple Registers
  packet[2] = (start_reg >> 8) & 0xFF;
  packet[3] = start_reg & 0xFF;
  packet[4] = (num_regs >> 8) & 0xFF;
  packet[5] = num_regs & 0xFF;
  packet[6] = num_regs * 2; // Byte count
  
  for (int i = 0; i < num_regs; i++) {
    packet[7 + i*2] = (values[i] >> 8) & 0xFF;
    packet[8 + i*2] = values[i] & 0xFF;
  }
  
  int len = 7 + num_regs * 2;
  uint16_t crc = calculateCRC(packet, len);
  packet[len] = crc & 0xFF;
  packet[len+1] = (crc >> 8) & 0xFF;

  setRS485Tx();
  Serial1.write(packet, len + 2);
  Serial1.flush();
  setRS485Rx();

  // Read response (expected 8 bytes)
  unsigned long start_time = millis();
  int idx = 0;
  uint8_t response[8];
  while (millis() - start_time < 50) {
    if (Serial1.available()) {
      response[idx++] = Serial1.read();
      if (idx >= 8) break;
    }
  }

  if (idx < 8) return false;
  uint16_t resp_crc = calculateCRC(response, 6);
  uint16_t calc_crc = (response[7] << 8) | response[6];
  return (resp_crc == calc_crc && response[0] == slave_id && response[1] == 0x10);
}

// Read Holding Registers via Modbus RTU
bool readModbus(uint8_t slave_id, uint16_t start_reg, uint16_t num_regs, int16_t* output) {
  uint8_t packet[8];
  packet[0] = slave_id;
  packet[1] = 0x03; // Read Holding Registers
  packet[2] = (start_reg >> 8) & 0xFF;
  packet[3] = start_reg & 0xFF;
  packet[4] = (num_regs >> 8) & 0xFF;
  packet[5] = num_regs & 0xFF;
  
  uint16_t crc = calculateCRC(packet, 6);
  packet[6] = crc & 0xFF;
  packet[7] = (crc >> 8) & 0xFF;

  // Clear serial read buffer before sending
  while (Serial1.available()) Serial1.read();

  setRS485Tx();
  Serial1.write(packet, 8);
  Serial1.flush();
  setRS485Rx();

  // Read response
  int expected_len = 5 + num_regs * 2;
  uint8_t response[32];
  int idx = 0;
  unsigned long start_time = millis();
  
  while (millis() - start_time < 50) {
    if (Serial1.available()) {
      response[idx++] = Serial1.read();
      if (idx >= expected_len) break;
    }
  }

  if (idx < expected_len) return false;
  
  uint16_t resp_crc = calculateCRC(response, expected_len - 2);
  uint16_t calc_crc = (response[expected_len-1] << 8) | response[expected_len-2];
  if (resp_crc != calc_crc) return false;

  if (response[0] != slave_id || response[1] != 0x03) return false;

  for (int i = 0; i < num_regs; i++) {
    output[i] = (response[3 + i*2] << 8) | response[4 + i*2];
  }
  return true;
}

// Initialize Driver Mode and Settings
void initDrivers() {
  // Set Operating Mode to Velocity (3) for both drivers
  writeModbusSingle(1, REG_MODE, 3);
  delay(10);
  writeModbusSingle(2, REG_MODE, 3);
  delay(10);

  // Stop/Disable drivers on boot to keep it safe
  disableDrivers();
}

void enableDrivers() {
  if (drivers_enabled) return;

  // 1. Release electromagnetic brakes (write 0 to 0x201A & 0x201B)
  writeModbusSingle(1, REG_BRAKE_CH1, 0);
  writeModbusSingle(1, REG_BRAKE_CH2, 0);
  writeModbusSingle(2, REG_BRAKE_CH1, 0);
  writeModbusSingle(2, REG_BRAKE_CH2, 0);
  delay(20);

  // 2. Enable drivers (write 8 to 0x200E)
  writeModbusSingle(1, REG_CONTROL_WORD, 8);
  writeModbusSingle(2, REG_CONTROL_WORD, 8);
  delay(20);

  drivers_enabled = true;
}

void disableDrivers() {
  // 1. Send zero speeds
  int16_t stop_speeds[2] = {0, 0};
  writeModbusMultiple(1, REG_TARGET_RPM, 2, stop_speeds);
  writeModbusMultiple(2, REG_TARGET_RPM, 2, stop_speeds);

  // 2. Lock brakes (write 1 to 0x201A & 0x201B)
  writeModbusSingle(1, REG_BRAKE_CH1, 1);
  writeModbusSingle(1, REG_BRAKE_CH2, 1);
  writeModbusSingle(2, REG_BRAKE_CH1, 1);
  writeModbusSingle(2, REG_BRAKE_CH2, 1);
  delay(20);

  // 3. Disable drivers (write 7 to 0x200E)
  writeModbusSingle(1, REG_CONTROL_WORD, 7);
  writeModbusSingle(2, REG_CONTROL_WORD, 7);

  drivers_enabled = false;
}

void setup() {
  pinMode(RS485_DE_RE_PIN, OUTPUT);
  setRS485Rx();

  // USB serial communication to PC
  Serial.begin(115200);
  
  // RS485 Serial communication to ZLAC8015D drivers
  Serial1.begin(115200);

  // Wait for drivers to boot up
  delay(1000);
  initDrivers();
}

void loop() {
  // Check Serial command packet from PC
  // Format: [0xFE, left_rpm_H, left_rpm_L, right_rpm_H, right_rpm_L, control_state, checksum] -> 7 bytes
  if (Serial.available() >= 7) {
    if (Serial.read() == 0xFE) {
      uint8_t buffer[6];
      Serial.readBytes(buffer, 6);
      
      uint8_t checksum = 0;
      for (int i = 0; i < 5; i++) {
        checksum += buffer[i];
      }
      
      if (checksum == buffer[5]) {
        last_cmd_time = millis();
        PCCommand cmd;
        cmd.left_rpm = (buffer[0] << 8) | buffer[1];
        cmd.right_rpm = (buffer[2] << 8) | buffer[3];
        cmd.control_state = buffer[4];

        if (cmd.control_state == 1) {
          enableDrivers();
          
          // Write speeds (invert right motor target RPM due to mirrored physical mounting)
          int16_t front_speeds[2] = {cmd.left_rpm, -cmd.right_rpm};
          int16_t rear_speeds[2] = {cmd.left_rpm, -cmd.right_rpm};
          
          writeModbusMultiple(1, REG_TARGET_RPM, 2, front_speeds);
          writeModbusMultiple(2, REG_TARGET_RPM, 2, rear_speeds);
        } else {
          disableDrivers();
        }
      }
    }
  }

  // Safety Timeout: stop the robot if connection to PC is lost
  if (drivers_enabled && (millis() - last_cmd_time > CMD_TIMEOUT_MS)) {
    disableDrivers();
  }

  // Periodic Telemetry Feedback to PC (approx. 20Hz)
  static unsigned long last_telemetry_time = 0;
  if (millis() - last_telemetry_time >= 50) {
    last_telemetry_time = millis();
    
    int16_t front_data[2] = {0, 0};
    int16_t rear_data[2] = {0, 0};
    
    // Read actual RPMs (invert right motor feedback RPM due to mirrored physical mounting)
    bool read_front = readModbus(1, REG_FEEDBACK_RPM, 2, front_data);
    bool read_rear = readModbus(2, REG_FEEDBACK_RPM, 2, rear_data);
    
    if (read_front) {
      feedback.front_left_rpm = front_data[0];
      feedback.front_right_rpm = -front_data[1]; // Invert right
    }
    if (read_rear) {
      feedback.rear_left_rpm = rear_data[0];
      feedback.rear_right_rpm = -rear_data[1]; // Invert right
    }
    
    // Send feedback packet to PC
    // Format: [0xFD, f_lh, f_ll, f_rh, f_rl, r_lh, r_ll, r_rh, r_rl, checksum] -> 10 bytes
    uint8_t resp[10];
    resp[0] = 0xFD;
    resp[1] = (feedback.front_left_rpm >> 8) & 0xFF;
    resp[2] = feedback.front_left_rpm & 0xFF;
    resp[3] = (feedback.front_right_rpm >> 8) & 0xFF;
    resp[4] = feedback.front_right_rpm & 0xFF;
    resp[5] = (feedback.rear_left_rpm >> 8) & 0xFF;
    resp[6] = feedback.rear_left_rpm & 0xFF;
    resp[7] = (feedback.rear_right_rpm >> 8) & 0xFF;
    resp[8] = feedback.rear_right_rpm & 0xFF;
    
    uint8_t checksum = 0;
    for (int i = 1; i < 9; i++) {
      checksum += resp[i];
    }
    resp[9] = checksum;
    
    Serial.write(resp, 10);
  }
}
