#include "commands.h"
#include "nodegraph.h"
#include "sensors.h"
#include "esp_system.h"
#include "soc/rtc_cntl_reg.h"
#include <vector>

// External references
extern RobotConfig config;
extern ServoManager servoManager;

// Global state instances
AnimationState animState;
MotorStreamState motorStream;

// ============================================================================
// AnimationState
// ============================================================================

void AnimationState::play(PlayMode mode, uint8_t repeats) {
  current = &animation;
  current->setActive(true);
  playMode = mode;
  repeatsRemaining = repeats;
}

void AnimationState::stop() {
  if (current) {
    current->setActive(false);
  }
  playMode = PLAY_IDLE;
}

// ============================================================================
// MotorStreamState
// ============================================================================

void MotorStreamState::start() {
  active = true;
  lastStreamTime = millis();
}

void MotorStreamState::stop() {
  active = false;
}

bool MotorStreamState::shouldStream() {
  if (!active) return false;
  if (millis() - lastStreamTime >= STREAM_INTERVAL_MS) {
    lastStreamTime = millis();
    return true;
  }
  return false;
}

// ============================================================================
// Command Dispatcher
// ============================================================================

void dispatchCommand() {
  const char* tag = getReceivedTag();
  const uint8_t* payload = getReceivedPayload();
  uint16_t len = getReceivedPayloadLen();
  
  // Identity & Config
  if (tagMatches(tag, Tag::IDENT)) {
    handleIdent(payload, len);
  }
  else if (tagMatches(tag, Tag::CONFG)) {
    handleConfig(payload, len);
  }
  // File Operations
  else if (tagMatches(tag, Tag::FLIST)) {
    handleFileList(payload, len);
  }
  else if (tagMatches(tag, Tag::FLOAD)) {
    handleFileLoad(payload, len);
  }
  else if (tagMatches(tag, Tag::FSAVE)) {
    handleFileSave(payload, len);
  }
  else if (tagMatches(tag, Tag::FDELE)) {
    handleFileDelete(payload, len);
  }
  else if (tagMatches(tag, Tag::FPLAY)) {
    handleFilePlay(payload, len);
  }
  // Motor Control
  else if (tagMatches(tag, Tag::MSET)) {
    handleMotorSet(payload, len);
  }
  else if (tagMatches(tag, Tag::MSCAN)) {
    handleMotorScan(payload, len);
  }
  else if (tagMatches(tag, Tag::MWRIT)) {
    handleMotorWrite(payload, len);
  }
  else if (tagMatches(tag, Tag::MSTRM)) {
    handleMotorStream(payload, len);
  }
  // System
  else if (tagMatches(tag, Tag::BOOT)) {
    handleBootloader(payload, len);
  }
  else {
    char tagStr[5] = {tag[0], tag[1], tag[2], tag[3], 0};
    sendMessage("Unknown tag: " + String(tagStr));
  }
}

// ============================================================================
// Identity & Config Handlers
// ============================================================================

void handleIdent(const uint8_t* payload, uint16_t len) {
  std::vector<uint8_t> response = config.serializeToBytes();
  sendPacket(Tag::IDENT, response.data(), response.size());
}

void handleConfig(const uint8_t* payload, uint16_t len) {
  RobotConfig newConfig;
  uint16_t offset = 0;

  // Decode robot name
  if (offset >= len) { sendNack(Tag::CONFG, "Missing name"); return; }
  uint8_t nameLength = payload[offset++];
  for (uint8_t i = 0; i < nameLength && offset < len; ++i) {
    newConfig.deviceName += (char)payload[offset++];
  }

  // Decode firmware version
  if (offset + 2 > len) { sendNack(Tag::CONFG, "Missing version"); return; }
  newConfig.firmwareVersion.major = payload[offset++];
  newConfig.firmwareVersion.minor = payload[offset++];

  // Decode motor count
  if (offset >= len) { sendNack(Tag::CONFG, "Missing motors"); return; }
  uint8_t motorCount = payload[offset++];

  // Decode motors
  for (uint8_t i = 0; i < motorCount && offset < len; ++i) {
    if (offset + 5 > len) break;

    ServoModel model;
    model.major = payload[offset++];
    model.minor = payload[offset++];

    uint16_t motorID = payload[offset++] | (payload[offset++] << 8);

    uint8_t motorNameLength = payload[offset++];
    String motorName = "";
    for (uint8_t j = 0; j < motorNameLength && offset < len; ++j) {
      motorName += (char)payload[offset++];
    }

    Motor motor;
    motor.motorID = motorID;
    motor.servoModel = model;
    motor.name = motorName;
    motor.position = 0;
    motor.isEnabled = true;

    newConfig.motors.push_back(motor);
  }

  config = newConfig;
  config.saveToFFat();
  
  sendAck(Tag::CONFG);
  sendMessage("Config updated: " + config.deviceName);
}

// ============================================================================
// File Operation Handlers
// ============================================================================

void handleFileList(const uint8_t* payload, uint16_t len) {
  File root = FFat.open("/");
  if (!root || !root.isDirectory()) {
    sendPacket(Tag::FLIST, nullptr, 0);
    return;
  }

  String fileList = "";
  File file = root.openNextFile();
  while (file) {
    if (!file.isDirectory()) {
      fileList += String(file.name()) + "\n";
    }
    file = root.openNextFile();
  }

  sendPacket(Tag::FLIST, (const uint8_t*)fileList.c_str(), fileList.length());
}

void handleFileLoad(const uint8_t* payload, uint16_t len) {
  if (len == 0 || len >= 128) {
    sendNack(Tag::FLOAD, "Invalid filename");
    return;
  }

  char filename[128];
  memcpy(filename, payload, len);
  filename[len] = '\0';

  File file = FFat.open(filename, "r");
  if (!file || !file.available()) {
    sendNack(Tag::FLOAD, "File not found");
    return;
  }

  size_t fileSize = file.size();
  if (fileSize > MAX_PAYLOAD_SIZE) {
    sendNack(Tag::FLOAD, "File too large");
    file.close();
    return;
  }

  uint8_t* buffer = (uint8_t*)malloc(fileSize);
  if (!buffer) {
    sendNack(Tag::FLOAD, "Out of memory");
    file.close();
    return;
  }

  file.read(buffer, fileSize);
  file.close();

  sendPacket(Tag::FLOAD, buffer, fileSize);
  free(buffer);
}

void handleFileSave(const uint8_t* payload, uint16_t len) {
  if (parseAndSaveAnimation(payload, len, animState.animation)) {
    sendAck(Tag::FSAVE);
  } else {
    sendNack(Tag::FSAVE, "Parse failed");
  }
}

void handleFileDelete(const uint8_t* payload, uint16_t len) {
  if (len < 2) {
    sendNack(Tag::FDELE, "Invalid request");
    return;
  }
  
  uint16_t filenameLen = payload[0] | (payload[1] << 8);
  if (len < 2 + filenameLen) {
    sendNack(Tag::FDELE, "Invalid filename");
    return;
  }

  char filename[128];
  memcpy(filename, payload + 2, min((uint16_t)127, filenameLen));
  filename[min((uint16_t)127, filenameLen)] = '\0';

  String fullPath = "/" + String(filename);
  deleteFile(FFat, fullPath.c_str());
  
  sendAck(Tag::FDELE);
  sendMessage("Deleted: " + String(filename));
}

void handleFilePlay(const uint8_t* payload, uint16_t len) {
  if (len < 4) {
    sendNack(Tag::FPLAY, "Invalid request");
    return;
  }
  
  uint16_t filenameLen = payload[0] | (payload[1] << 8);
  if (len < 2 + filenameLen + 2) {
    sendNack(Tag::FPLAY, "Invalid filename");
    return;
  }

  char filename[128];
  memcpy(filename, payload + 2, min((uint16_t)127, filenameLen));
  filename[min((uint16_t)127, filenameLen)] = '\0';

  uint16_t offset = 2 + filenameLen;
  PlayMode mode = static_cast<PlayMode>(payload[offset]);
  uint8_t repeats = payload[offset + 1];

  animState.animation.clear();
  String fullPath = "/" + String(filename);
  
  if (animState.animation.loadFromFile(fullPath.c_str())) {
    animState.play(mode, repeats);
    sendAck(Tag::FPLAY);
    sendMessage("Playing: " + String(filename));
  } else {
    sendNack(Tag::FPLAY, "Load failed");
  }
}

// ============================================================================
// Motor Control Handlers
// ============================================================================

void handleMotorSet(const uint8_t* payload, uint16_t len) {
  if (len % 3 != 0) {
    sendNack(Tag::MSET, "Invalid length");
    return;
  }

  uint8_t motorCount = len / 3;
  uint8_t ids[motorCount];
  uint16_t positions[motorCount];
  uint16_t speeds[motorCount];

  for (uint8_t i = 0; i < motorCount; ++i) {
    ids[i] = payload[i * 3];
    positions[i] = (payload[i * 3 + 2] << 8) | payload[i * 3 + 1];
    speeds[i] = 0;
  }

  servoManager.syncWritePositions(ids, positions, speeds, motorCount, config, 0);
  sendAck(Tag::MSET);
}

void handleMotorScan(const uint8_t* payload, uint16_t len) {
  if (len < 1) {
    sendNack(Tag::MSCAN, "Invalid request");
    return;
  }

  uint8_t channel = payload[0];
  Feetech* servo = servoManager[channel];

  for (int i = 0; i < 254; i++) {
    servo->sendPing(i);
    uint8_t val = servo->waitOnData1Byte(10);

    if (val != 0) {
      uint8_t response[33];
      response[0] = channel;
      response[1] = i;

      uint16_t model = servo->getModel(i);
      servo->setFeetechMode(model);
      
      uint16_t minAngle = servo->getMinAngleLimit(i);
      uint16_t maxAngle = servo->getMaxAngleLimit(i);
      uint16_t position = servo->getPosition(i);
      uint8_t cwDead = servo->getCWDeadZone(i);
      uint8_t ccwDead = servo->getCCWDeadZone(i);
      uint16_t offset = servo->getOffset(i);
      uint8_t mode = servo->getMode(i);
      uint8_t torque = servo->getTorqueEnable(i);
      uint8_t accel = servo->getAcceleration(i);
      uint16_t goalPos = servo->getGoalPosition(i);
      uint16_t goalTime = servo->getGoalTime(i);
      uint16_t goalSpeed = servo->getGoalSpeed(i);
      uint8_t lock = servo->getLock(i);
      int16_t speed = servo->getSpeed(i);
      uint16_t load = servo->getLoad(i);
      uint8_t temp = servo->getTemperature(i);
      uint8_t moving = servo->getMoving(i);
      uint8_t current = servo->getCurrent(i);
      uint8_t voltage = servo->getVoltage(i);

      // Pack response
      response[2] = (model >> 8) & 0xFF;
      response[3] = model & 0xFF;
      response[4] = (minAngle >> 8) & 0xFF;
      response[5] = minAngle & 0xFF;
      response[6] = (maxAngle >> 8) & 0xFF;
      response[7] = maxAngle & 0xFF;
      response[8] = (position >> 8) & 0xFF;
      response[9] = position & 0xFF;
      response[10] = cwDead;
      response[11] = ccwDead;
      response[12] = (offset >> 8) & 0xFF;
      response[13] = offset & 0xFF;
      response[14] = mode;
      response[15] = torque;
      response[16] = accel;
      response[17] = (goalPos >> 8) & 0xFF;
      response[18] = goalPos & 0xFF;
      response[19] = (goalTime >> 8) & 0xFF;
      response[20] = goalTime & 0xFF;
      response[21] = (goalSpeed >> 8) & 0xFF;
      response[22] = goalSpeed & 0xFF;
      response[23] = lock;
      response[24] = (speed >> 8) & 0xFF;
      response[25] = speed & 0xFF;
      response[26] = (load >> 8) & 0xFF;
      response[27] = load & 0xFF;
      response[28] = temp;
      response[29] = moving;
      response[30] = (current >> 8) & 0xFF;
      response[31] = current & 0xFF;
      response[32] = voltage;

      sendPacket(Tag::MSCAN, response, 33);
    }
  }

  // Signal end of scan
  uint8_t endResponse[2] = { channel, 255 };
  sendPacket(Tag::MSCAN, endResponse, 2);
}

void handleMotorWrite(const uint8_t* payload, uint16_t len) {
  if (len < 5) {
    sendNack(Tag::MWRIT, "Invalid request");
    return;
  }

  uint8_t channel = payload[0];
  uint8_t id = payload[1];
  uint8_t reg = payload[2];
  uint8_t dataLen = payload[3];

  Feetech* servo = servoManager[channel];
  uint8_t model = servo->getMajorModel(id);
  servo->setFeetechMode(model);

  // Special case for ID changes (reg 5)
  if (reg == 5 && dataLen == 1) {
    servo->changeID(id, payload[4]);
    sendAck(Tag::MWRIT);
    sendMessage("ID changed");
    return;
  }

  if (dataLen == 1) {
    servo->write1Byte(id, reg, payload[4]);
    uint8_t response = servo->waitOnData1Byte(10);
    sendPacket(Tag::MWRIT, &response, 1);
  } else {
    servo->write2Bytes(id, reg, payload[4] | (payload[5] << 8));
    uint16_t response = servo->waitOnData2Bytes(10);
    uint8_t respBuf[2] = { (uint8_t)(response & 0xFF), (uint8_t)(response >> 8) };
    sendPacket(Tag::MWRIT, respBuf, 2);
  }
}

void handleMotorStream(const uint8_t* payload, uint16_t len) {
  if (len < 1) {
    sendNack(Tag::MSTRM, "Invalid request");
    return;
  }

  bool enable = payload[0] != 0;

  if (enable) {
    // Disable torque for position reading
    for (const Motor& motor : config.motors) {
      servoManager[0]->setFeetechMode(motor.servoModel.major);
      servoManager[0]->disableTorque(motor.motorID);
    }
    motorStream.start();
  } else {
    motorStream.stop();
  }

  sendAck(Tag::MSTRM);
}

// ============================================================================
// System Handlers
// ============================================================================

void handleBootloader(const uint8_t* payload, uint16_t len) {
  sendMessage("Entering bootloader...");
  Serial.flush();
  delay(100);
  
  REG_WRITE(RTC_CNTL_OPTION1_REG, RTC_CNTL_FORCE_DOWNLOAD_BOOT);
  esp_restart();
}

// ============================================================================
// Helper Functions
// ============================================================================

void sendMotorPositions() {
  std::vector<uint8_t> payload;

  for (const Motor& motor : config.motors) {
    payload.push_back(motor.motorID);
    payload.push_back(motor.position & 0xFF);
    payload.push_back((motor.position >> 8) & 0xFF);
  }

  sendPacket(Tag::MPOS, payload.data(), payload.size());
}

bool parseAndSaveAnimation(const uint8_t* payload, uint16_t len, Animation& animation) {
  if (len < 2) return false;
  
  uint16_t filenameLen = payload[0] | (payload[1] << 8);
  if (len < 2 + filenameLen + 18) return false;

  char filename[128];
  memcpy(filename, payload + 2, min((uint16_t)127, filenameLen));
  filename[min((uint16_t)127, filenameLen)] = '\0';

  const uint8_t* ptr = payload + 2 + filenameLen;
  uint16_t remaining = len - 2 - filenameLen;

  sendMessage("Saving: " + String(filename));

  // Parse header
  if (remaining < 18) return false;
  memcpy(animation.header.magic, ptr, 4);
  if (strncmp(animation.header.magic, "ANIM", 4) != 0) {
    sendMessage("Invalid magic header");
    return false;
  }

  animation.header.frameCount = ptr[4] | (ptr[5] << 8);
  animation.header.version = ptr[6];
  animation.header.frameRate = ptr[7];
  memcpy(animation.header.reserved, ptr + 8, 8);

  uint16_t curveCount = ptr[16] | (ptr[17] << 8);
  ptr += 18;
  remaining -= 18;

  // Parse curves
  uint16_t curveDataSize = curveCount * sizeof(CurveSegment);
  if (remaining < curveDataSize) return false;

  animation.clearAllCurves();
  for (uint16_t i = 0; i < curveCount; i++) {
    CurveSegment seg;
    memcpy(&seg, ptr, sizeof(CurveSegment));
    animation.addCurveSegment(seg);
    ptr += sizeof(CurveSegment);
  }
  remaining -= curveDataSize;

  // Parse node graph
  if (remaining > 0) {
    loadNodeGraph(ptr, remaining, animation.nodeGraph);
    animation.nodeGraph.bindAnimationContext(&animation);
  }

  // Save to file
  String fullPath = "/" + String(filename);
  animation.saveToFile(fullPath.c_str());
  
  sendMessage("Saved: " + String(filename));
  return true;
}

void deleteFile(fs::FS& fs, const char* path) {
  if (fs.remove(path)) {
    Serial.printf("Deleted: %s\n", path);
  } else {
    Serial.printf("Delete failed: %s\n", path);
  }
}