#include "motoraid.h"
#include "motorcontrol.h"
#include "robotconfig.h"

// Global instance
MotorAid motorAid;

extern RobotConfig config;

void MotorAid::addMotor(uint8_t motorID) {
  // Don't add duplicates
  if (findMotor(motorID) != nullptr) return;
  
  AidedMotor motor;
  motor.motorID = motorID;
  
  // Read actual position directly from servo (not config which may be stale)
  uint8_t model = config.getMotorModel(motorID);
  servoManager[0]->setFeetechMode(model);
  motor.lastPosition = servoManager[0]->getPosition(motorID);
  motor.lastUpdateTime = millis();
  
  // Initialize velocity samples to zero
  for (uint8_t i = 0; i < VELOCITY_SAMPLES; i++) {
    motor.velocitySamples[i] = 0;
  }
  
  aidedMotors.push_back(motor);
  
  // Disable torque on this motor so it can be moved freely
  servoManager[0]->disableTorque(motorID);
  
  Serial.print("[MotorAid] Added motor ");
  Serial.print(motorID);
  Serial.print(" at pos ");
  Serial.println(motor.lastPosition);
}

void MotorAid::removeMotor(uint8_t motorID) {
  for (auto it = aidedMotors.begin(); it != aidedMotors.end(); ++it) {
    if (it->motorID == motorID) {
      aidedMotors.erase(it);
      Serial.print("[MotorAid] Removed motor ");
      Serial.println(motorID);
      return;
    }
  }
}

bool MotorAid::isMotorAided(uint8_t motorID) const {
  return findMotor(motorID) != nullptr;
}

AidedMotor* MotorAid::findMotor(uint8_t motorID) {
  for (auto& motor : aidedMotors) {
    if (motor.motorID == motorID) return &motor;
  }
  return nullptr;
}

const AidedMotor* MotorAid::findMotor(uint8_t motorID) const {
  for (const auto& motor : aidedMotors) {
    if (motor.motorID == motorID) return &motor;
  }
  return nullptr;
}

int16_t MotorAid::calculateSmoothedVelocity(AidedMotor& motor) {
  // Calculate average of velocity samples
  int32_t sum = 0;
  for (uint8_t i = 0; i < VELOCITY_SAMPLES; i++) {
    sum += motor.velocitySamples[i];
  }
  return (int16_t)(sum / VELOCITY_SAMPLES);
}

void MotorAid::update() {
  unsigned long now = millis();
  
  for (AidedMotor& motor : aidedMotors) {
    // Calculate time delta
    unsigned long deltaTime = now - motor.lastUpdateTime;
    if (deltaTime < updateInterval) continue;
    
    // Get current position from config (already updated by updateMotorPositions)
    uint16_t currentPosition = config.getMotorPosition(motor.motorID);
    
    // Calculate instantaneous velocity (position units per second)
    int16_t positionDelta = (int16_t)currentPosition - (int16_t)motor.lastPosition;
    int16_t instantVelocity = (positionDelta * 1000) / (int16_t)deltaTime;
    
    // Add to velocity samples (circular buffer)
    motor.velocitySamples[motor.sampleIndex] = instantVelocity;
    motor.sampleIndex = (motor.sampleIndex + 1) % VELOCITY_SAMPLES;
    
    // Track samples collected for warmup
    if (motor.samplesCollected < VELOCITY_SAMPLES * 2) {
      motor.samplesCollected++;
    }
    
    // Don't calculate or trigger until we have enough samples (2x buffer = ~300ms warmup)
    if (motor.samplesCollected < VELOCITY_SAMPLES * 2) {
      continue;  // Still warming up
    }
    motor.warmedUp = true;
    
    // Calculate smoothed velocity
    motor.smoothedVelocity = calculateSmoothedVelocity(motor);
    
    // Update tracking
    motor.lastPosition = currentPosition;
    motor.lastUpdateTime = now;
    
    // Check if currently assisting
    if (motor.assisting) {
      // Calculate target position AHEAD of current position in direction of movement
      int16_t leadOffset = (motor.smoothedVelocity > 0) ? 150 : -150;
      motor.assistTargetPosition = constrain((int16_t)currentPosition + leadOffset, 0, 4095);
      
      // Send updated position command to lead the movement
      uint8_t model = config.getMotorModel(motor.motorID);
      servoManager[0]->setFeetechMode(model);
      
      uint8_t ids[1] = { motor.motorID };
      uint16_t positions[1] = { motor.assistTargetPosition };
      uint16_t speeds[1] = { assistSpeed };
      servoManager.syncWritePositions(ids, positions, speeds, 1, config, 0);
      
      // Extend assist if still moving fast
      if (abs(motor.smoothedVelocity) > resetThreshold) {
        motor.assistStartTime = now;  // Keep extending
      }
      
      // Check if assist duration has elapsed (only ends when movement stops)
      if (now - motor.assistStartTime >= assistDuration) {
        motor.assisting = false;
        // Disable torque again
        servoManager[0]->disableTorque(motor.motorID);
        
        Serial.print("[MotorAid] Assist ended, motor ");
        Serial.print(motor.motorID);
        Serial.print(" at pos ");
        Serial.println(currentPosition);
      }
      continue;
    }
    
    // Hysteresis: if we were triggered, wait for velocity to drop before re-triggering
    if (motor.wasTriggered) {
      if (abs(motor.smoothedVelocity) < resetThreshold) {
        motor.wasTriggered = false;  // Reset, can trigger again
      }
      continue;  // Don't trigger while in hysteresis zone
    }
    
    // Check if smoothed velocity exceeds threshold (in either direction)
    if (abs(motor.smoothedVelocity) > velocityThreshold) {
      // Someone is moving the motor - assist!
      motor.assisting = true;
      motor.assistStartTime = now;
      motor.wasTriggered = true;
      
      // Calculate target position AHEAD of current position in direction of movement
      int16_t leadOffset = (motor.smoothedVelocity > 0) ? 150 : -150;
      motor.assistTargetPosition = constrain((int16_t)currentPosition + leadOffset, 0, 4095);
      
      // Set servo mode based on motor model
      uint8_t model = config.getMotorModel(motor.motorID);
      servoManager[0]->setFeetechMode(model);
      
      // Enable torque and move AHEAD to assist
      servoManager[0]->enableTorque(motor.motorID);
      
      uint8_t ids[1] = { motor.motorID };
      uint16_t positions[1] = { motor.assistTargetPosition };
      uint16_t speeds[1] = { assistSpeed };
      servoManager.syncWritePositions(ids, positions, speeds, 1, config, 0);
      
      // Debug output
      Serial.print("[MotorAid] Assisting motor ");
      Serial.print(motor.motorID);
      Serial.print(" vel=");
      Serial.print(motor.smoothedVelocity);
      Serial.print(" pos=");
      Serial.print(currentPosition);
      Serial.print(" -> ");
      Serial.println(motor.assistTargetPosition);
    }
  }
}