#include <base64.h>
#include "feetech.h"
#include "animation.h"

#define DEVICE_NAME "Little Sophia"
#define FIRMWARE_VERSION "0.0.1"

#define HEADER1 0xAA
#define HEADER2 0x55

#define CMD_ID_REQUEST 0x01
#define CMD_FILE_LIST 0x02
#define CMD_LOAD_FILE 0x03
#define CMD_DELETE_FILE 0x04
#define CMD_LOAD_FILE_CHUNK 0x05


// ESP32 S2 PINOUT
#define RX_PIN 17  // DI
#define TX_PIN 18  // RO
#define DE_PIN 33  // Driver Enable
#define RE_PIN 3   // Receiver Enable

Animation sweep;
Animation stagger;

Feetech servos = Feetech(Serial1, DE_PIN, RE_PIN, TX_PIN, RX_PIN);

uint16_t flipBytes(uint16_t value) {
  return (value >> 8) | (value << 8);
}


uint8_t ids[NUM_CHANNELS] = { 10, 11, 12, 13, 14 };
uint16_t pos1[] = { 0, 0, 0, 0, 0 };
uint16_t pos2[] = { 1023, 1023, 1023, 1023, 1023 };

void setup() {
  Serial.begin(115200);
  for (int i = 0; i < 5; i++) {
    Serial.println(i);
    delay(500);
  }


  pos2[3] = flipBytes(pos2[3]);

  servos.begin();

  if (!FFat.begin()) {
    Serial.println("FFat mount failed");
    return;
  }

  // sweep.clear();
  // sweep.createSampleSweep(4);
  // sweep.saveToFile("/sweep.anim");
  // sweep.clear();
  // sweep.createStaggeredSweep(4);
  // sweep.saveToFile("/stagger.anim");
  // delay(9999);

  // anim.clear();
  // if (!sweep.loadFromFile("/sweep.anim")) {
  //   Serial.println("Failed to load animation");
  //   return;
  // } else {
  //   Serial.println("Loaded sweep anim");
  // }

  // if (!stagger.loadFromFile("/stagger.anim")) {
  //   Serial.println("Failed to load animation");
  //   return;
  // } else {
  //   Serial.println("Loaded stagger anim");
  // }


  // Serial.println(anim.getFrame(0, 0));  // Should show saved value

  //SetID(11, 14);
}

void SetID(uint8_t oldID, uint8_t newID) {
  Serial.println("Setting Lock to 0");
  Serial.println(servos.setLock(oldID, 0));
  delay(1000);
  Serial.print("Changing ID ");
  Serial.print(oldID);
  Serial.print(" to ");
  Serial.println(newID);
  Serial.println(servos.setID(oldID, newID));
  delay(1000);
  Serial.println("Setting Lock to 1");
  Serial.println(servos.setLock(newID, 1));
  delay(1000);
}

unsigned long lastSend = 0;
void loop() {
  HandleSerialRequests();

  // put your main code here, to run repeatedly:
  //PingAll();


  // playAnimation(sweep);
  // playAnimation(stagger);
  // playLayeredAnimation(sweep, stagger);


  // servos.syncWritePos(ids, pos1, 5);
  // delay(1000);

  // servos.syncWritePos(ids, pos2, 5);
  // delay(1000);

  if (millis() - lastSend > 1000) {
    //sendMessageFromESP32(String(millis()));
    //PrintFileList();
    lastSend = millis();
  }
}

void HandleSerialRequests() {
  if (Serial.available() >= 4) {
    if (Serial.read() == HEADER1 && Serial.read() == HEADER2) {
      uint8_t command = Serial.read();
      uint8_t length = Serial.read();

      String payload = "";
      for (int i = 0; i < length; i++) {
        while (!Serial.available())
          ;
        payload += (char)Serial.read();
      }

      handleCommand(command, payload);
    }
  }
}

void handleCommand(uint8_t command, const String& payload) {
  switch (command) {
    case CMD_ID_REQUEST:
      sendIdPacket();
      break;
    case CMD_FILE_LIST:
      sendFileList();
      break;
    case CMD_LOAD_FILE:
      sendFileContent(payload);
      break;
    case CMD_DELETE_FILE:
      deleteFile(payload);
      break;
      //default:
      //Serial.println("{\"error\":\"Unknown command\"}");
  }
}

void sendReply(uint8_t command, const String& payload) {
  uint16_t length = payload.length();  // Supports up to 65,535 bytes

  // Calculate checksum using both length bytes
  uint8_t checksum = command ^ (length >> 8) ^ (length & 0xFF);
  for (int i = 0; i < length; i++) {
    checksum ^= payload[i];
  }

  // Send packet with 2-byte length
  Serial.write(HEADER1);
  Serial.write(HEADER2);
  Serial.write(command);
  Serial.write((length >> 8) & 0xFF);  // High byte
  Serial.write(length & 0xFF);         // Low byte
  Serial.write((const uint8_t*)payload.c_str(), length);
  Serial.write(checksum);
}


void sendChunkReply(uint8_t command, const String& filename, const uint8_t* chunkData, size_t chunkSize, size_t offset, size_t totalSize) {
  // Build JSON metadata
  String payload = "{\"file\":\"" + filename + "\",\"offset\":" + offset + ",\"totalSize\":" + totalSize + ",\"chunk\":[";

  for (size_t i = 0; i < chunkSize; i++) {
    payload += String(chunkData[i]);
    if (i < chunkSize - 1) payload += ",";
  }

  payload += "]}";

  // Calculate payload length
  uint16_t length = payload.length();

  // Calculate checksum
  uint8_t checksum = command ^ (length >> 8) ^ (length & 0xFF);
  for (int i = 0; i < length; i++) {
    checksum ^= payload[i];
  }

  // Send packet
  Serial.write(HEADER1);
  Serial.write(HEADER2);
  Serial.write(command);
  Serial.write((length >> 8) & 0xFF);  // High byte
  Serial.write(length & 0xFF);         // Low byte
  Serial.write((const uint8_t*)payload.c_str(), length);
  Serial.write(checksum);
}



void sendIdPacket() {
  String payload = "{\"name\":\"" + String(DEVICE_NAME) + "\",\"version\":\"" + String(FIRMWARE_VERSION) + "\"}";
  sendReply(CMD_ID_REQUEST, payload);
}


void PrintFileList() {
  File root = FFat.open("/");
  if (!root || !root.isDirectory()) {
    Serial.println("Failed to open FFat root directory");
    return;
  }

  Serial.println("Files in FFat:");

  File file = root.openNextFile();
  while (file) {
    Serial.print("  ");
    Serial.print(file.name());
    Serial.print("  |  Size: ");
    Serial.println(file.size());
    file = root.openNextFile();
  }

  Serial.println("End of file list.");
}

void sendFileList() {
  File root = FFat.open("/");
  if (!root || !root.isDirectory()) {
    sendReply(CMD_FILE_LIST, "[]");
    return;
  }

  File file = root.openNextFile();
  String payload = "[";

  bool first = true;
  while (file) {
    if (!file.isDirectory()) {
      if (!first) payload += ",";
      payload += "\"" + String(file.name()) + "\"";
      first = false;
    }
    file = root.openNextFile();
  }

  payload += "]";
  sendReply(CMD_FILE_LIST, payload);
}

void sendFileContent(const String& filename) {
  sendFileInChunks(filename);
  return;
  File file = FFat.open(filename, FILE_READ);
  if (!file) {
    sendReply(CMD_LOAD_FILE, "{\"error\":\"File not found\"}");
    return;
  }

  String raw;
  while (file.available()) {
    raw += (char)file.read();
  }
  file.close();

  String encoded = base64::encode(raw);  // This must be base64

  String payload = "{\"file\":\"" + filename + "\",\"content\":\"" + encoded + "\"}";
  sendReply(CMD_LOAD_FILE, payload);
}

void sendFileInChunks(const String& filename) {
  File file = FFat.open(filename, FILE_READ);
  if (!file) {
    sendReply(CMD_LOAD_FILE, "{\"error\":\"File not found\"}");
    return;
  }

  const size_t chunkSize = 512;
  size_t totalSize = file.size();
  size_t offset = 0;
  size_t chunksSent = 0;

  unsigned long startTime = millis();

  while (offset < totalSize) {
    size_t remaining = totalSize - offset;
    size_t thisChunk = remaining < chunkSize ? remaining : chunkSize;

    uint8_t buffer[thisChunk];
    file.read(buffer, thisChunk);

    sendChunkReply(CMD_LOAD_FILE_CHUNK, filename, buffer, thisChunk, offset, totalSize);

    offset += thisChunk;
    chunksSent++;
    delay(10);  // Optional pacing
  }

  file.close();

  unsigned long duration = millis() - startTime;

  String finalReply = "{\"status\":\"complete\",\"file\":\"" + filename +
                      "\",\"chunks\":" + chunksSent +
                      ",\"bytesSent\":" + totalSize +
                      ",\"durationMs\":" + duration + "}";

  sendReply(CMD_LOAD_FILE, finalReply);
}



void deleteFile(const String& filename) {
  String payload = "{\"deleted\":\"" + filename + "\"}";
  sendReply(CMD_DELETE_FILE, payload);
}

void playAnimation(Animation& anim) {
  uint16_t positions[NUM_CHANNELS];
  const uint16_t frameCount = anim.getFrameCount();
  const uint32_t frameDelay = 1000 / FRAMES_PER_SECOND;  // 20 ms

  uint32_t nextFrameTime = millis();

  for (uint16_t frame = 0; frame < frameCount; frame++) {
    // Wait until it's time for the next frame
    while (millis() < nextFrameTime) {
      // Optional: yield or do background tasks here
      delay(1);
    }

    // Send frame to servos
    if (anim.getFramePositions(frame, positions)) {
      servos.syncWritePos(ids, positions, NUM_CHANNELS);
    }

    // Schedule next frame
    nextFrameTime += frameDelay;
  }
}

void playLayeredAnimation(Animation& base, Animation& overlay) {
  uint16_t basePositions[NUM_CHANNELS];
  uint16_t overlayPositions[NUM_CHANNELS];
  uint16_t finalPositions[NUM_CHANNELS];

  const uint16_t frameCount = base.getFrameCount();
  const uint32_t frameDelay = 1000 / FRAMES_PER_SECOND;
  uint32_t nextFrameTime = millis();

  for (uint16_t frame = 0; frame < frameCount; frame++) {
    while (millis() < nextFrameTime) delay(1);

    base.getFramePositions(frame, basePositions);
    overlay.getFramePositions(frame, overlayPositions);

    for (uint8_t ch = 0; ch < NUM_CHANNELS; ch++) {
      finalPositions[ch] = (basePositions[ch] + overlayPositions[ch]) / 2;
    }

    servos.syncWritePos(ids, finalPositions, NUM_CHANNELS);
    nextFrameTime += frameDelay;
  }
}



void PingAll() {
  std::vector<uint8_t> successfulAddresses;
  servos.pingAll(successfulAddresses);

  // Now successfulAddresses contains all successful pings
  Serial.println("Successful Addresses:");
  for (uint8_t address : successfulAddresses) {
    Serial.print(address);
    Serial.print(" ");
  }
  Serial.println();
}