#include "Config.h"
#include "Kinematics.h"
#include "GaitEngine.h"
#include "BodyPose.h"
#include "RobotIMU.h"
#include "StateMachine.h"
#include "InputMapper.h"
#include "Limits.h"
#include "Comms.h"
#include "ServoDriver.h"
#include "SmoothedValue.h"

// Uncomment to drive the robot from the Serial Monitor before a remote exists.
// See README.md for the key map.
#define TEST_SERIAL_INPUT

GaitEngine gait;
BodyPose pose;
RobotIMU imu;
StateMachine state_machine;
InputMapper mapper;
Comms comms;
ServoDriver servo_driver;

float current_heading_deg = 0;          // open-loop: accumulates commanded omega. See note in loop().
SmoothedValue stance_height(STANCE_HEIGHT_RATE_MM_S);   // remote height slider, rate-limited
float gyro_corr_roll_deg = 0, gyro_corr_pitch_deg = 0;  // robot-gyro self-leveling (integrated)
float last_joint_deg[NUM_LEGS][3] = {}; // for per-joint rate limiting
uint32_t last_loop_us = 0;
RobotState last_state = RobotState::DISARMED;

#ifdef TEST_SERIAL_INPUT
CmdPacket test_packet;
void pollSerialTestInput() {
  bool changed = false;
  while (Serial.available()) {
    char c = Serial.read();
    changed = true;
    switch (c) {
      case 'a': test_packet.flags ^= FLAG_ARMED; break;
      case 's': test_packet.flags ^= FLAG_STAND; break;
      // vx = forward (+X), vy = strafe left (+Y) - see HexapodProtocol.h axis convention.
      case 'w': test_packet.vx = min(test_packet.vx + 200, 1000);  test_packet.flags |= FLAG_WALK; break;
      case 'x': test_packet.vx = max(test_packet.vx - 200, -1000); test_packet.flags |= FLAG_WALK; break;
      case 'q': test_packet.vy = min(test_packet.vy + 200, 1000);  test_packet.flags |= FLAG_WALK; break;
      case 'e': test_packet.vy = max(test_packet.vy - 200, -1000); test_packet.flags |= FLAG_WALK; break;
      // yaw_rate + = turn left (CCW).
      case 'z': test_packet.yaw_rate = min(test_packet.yaw_rate + 200, 1000);  test_packet.flags |= FLAG_WALK; break;
      case 'c': test_packet.yaw_rate = max(test_packet.yaw_rate - 200, -1000); test_packet.flags |= FLAG_WALK; break;
      case 'g': test_packet.flags ^= FLAG_GAIT_WAVE; break;
      case 'm': test_packet.flags ^= FLAG_MODE_B; break;
      case ' ':
        test_packet.vx = 0; test_packet.vy = 0; test_packet.yaw_rate = 0;
        test_packet.flags &= ~FLAG_WALK;
        break;
      default: changed = false; break;
    }
  }
  if (changed) {
    test_packet.speed_scale = 180;
    comms.injectPacket(test_packet);
  }
}
#endif

void setup() {
  Serial.begin(115200);
  gait.begin();
  pose.begin();
  state_machine.begin();
  stance_height.reset(STANCE_HEIGHT_MM);

  if (!imu.begin())   Serial.println("robot IMU not responding - stabilization disabled");
  if (!comms.begin()) Serial.println("radio not responding - waiting for injected/test commands only");

  servo_driver.begin();
  last_loop_us = micros();
}

void loop() {
  uint32_t now_us = micros();
  float dt = (now_us - last_loop_us) / 1e6f;
  if (dt < (1.0f / CONTROL_HZ)) return;   // hold to the fixed control rate
  last_loop_us = now_us;

#ifdef TEST_SERIAL_INPUT
  pollSerialTestInput();
#endif

  comms.update();
  CmdPacket cmd = comms.linkOk() ? comms.latest() : CmdPacket();   // link lost -> all-neutral, unarmed

  MapperOutput mapped = mapper.map(cmd, current_heading_deg, dt);

  // Standing height: every neutral foot position (gait, sit/stand, limits)
  // follows this, so it's set before anything below uses them.
  stance_height.setTarget(mapped.stance_height_mm);
  Kinematics::setStanceHeight(stance_height.update(dt));

  // Slow the walk command down to a stride the legs can really take (reach,
  // no neighbouring legs colliding) - the stride alone; the body pose then
  // makes room for it below.
  gait.setRequestedPattern(mapped.gait_wave ? &GAIT_WAVE : &GAIT_TRIPOD);
  mapped.twist = gait.limitTwist(mapped.twist, GaitEngine::NO_POSE);

  imu.enabled = mapped.robot_gyro_enabled;
  imu.update(dt);

  state_machine.update(mapped.armed, mapped.stand_request, mapped.walk_active, dt);
  RobotState state = state_machine.state();

  // Relax servos on the way into DISARMED, reattach on the way out.
  if (state == RobotState::DISARMED && last_state != RobotState::DISARMED) {
    servo_driver.relaxAll();
  } else if (state != RobotState::DISARMED && last_state == RobotState::DISARMED) {
    servo_driver.begin();
  }
  // The stand pose holds every foot at neutral, so walking starts from there.
  if (state == RobotState::WALKING && last_state != RobotState::WALKING) gait.resetFeet();
  last_state = state;

  // Stick pose + remote-tilt feedforward (gated on FLAG_REMOTE_GYRO) are
  // already combined in InputMapper.
  pose.setCommand(mapped.pose);

  // Robot-gyro self-leveling. The IMU measures the body's tilt in the world,
  // which already includes the correction applied so far - so the correction
  // integrates (a plain "correction = -measured tilt" only halves a slope).
  // The tilt the operator asked for is subtracted first, so leveling doesn't
  // cancel the stick. Switched off, the correction bleeds away smoothly.
  if (mapped.robot_gyro_enabled) {
    gyro_corr_roll_deg  -= ROBOT_GYRO_KI * (imu.rollDeg()  - pose.commandRollDeg())  * dt;
    gyro_corr_pitch_deg -= ROBOT_GYRO_KI * (imu.pitchDeg() - pose.commandPitchDeg()) * dt;
    gyro_corr_roll_deg  = constrain(gyro_corr_roll_deg,  -ROBOT_GYRO_MAX_CORR_DEG, ROBOT_GYRO_MAX_CORR_DEG);
    gyro_corr_pitch_deg = constrain(gyro_corr_pitch_deg, -ROBOT_GYRO_MAX_CORR_DEG, ROBOT_GYRO_MAX_CORR_DEG);
  } else {
    float step = POSE_TILT_MAX_RATE_DEG * dt;
    gyro_corr_roll_deg  = constrain(0.0f, gyro_corr_roll_deg - step,  gyro_corr_roll_deg + step);
    gyro_corr_pitch_deg = constrain(0.0f, gyro_corr_pitch_deg - step, gyro_corr_pitch_deg + step);
  }
  pose.setCorrection(gyro_corr_roll_deg, gyro_corr_pitch_deg);
  pose.update(dt);

  FootTarget gait_targets[NUM_LEGS];
  bool in_stance[NUM_LEGS];

  if (state == RobotState::WALKING) {
    // Walking has priority over the body pose: the pose shrinks until the
    // stride the walk command wants fits (and where the feet are right now,
    // since a stride in transition can be outside that envelope). The pose
    // shrinks at a bounded rate - no snapping - so meanwhile the walk slows
    // down against the pose as it is, and everything stays physical.
    FootTarget sets[GaitEngine::ENVELOPE_SETS + 1][NUM_LEGS];
    gait.strideEnvelope(mapped.twist, sets);
    for (int i = 0; i < NUM_LEGS; i++) sets[GaitEngine::ENVELOPE_SETS][i] = gait.feet()[i];
    pose.fitToReach(sets, GaitEngine::ENVELOPE_SETS + 1, dt);
    Twist twist = gait.limitTwist(mapped.twist, pose, true);

    gait.update(dt, twist, pose, gait_targets);
    for (int i = 0; i < NUM_LEGS; i++) in_stance[i] = gait.legInStance(i);

    // Simple open-loop heading estimate: accumulates the omega the planted
    // feet actually carried out (only while walking - the body only turns
    // then). Good enough for mode B in a normal session; fusing the robot
    // IMU's yaw rate in as well is the more robust upgrade if legs ever slip.
    current_heading_deg += degrees(gait.appliedTwist().omega) * dt;
  } else {
    // Sit/stand: blend between a crouched and a neutral stance height, no
    // stepping. This assumes "sit" keeps all six feet grounded - if your
    // sit pose rests the chassis on the ground instead, replace this with
    // a proper sequenced transition (feet positioned first, then lift).
    float blend = state_machine.sitStandBlend();
    for (int i = 0; i < NUM_LEGS; i++) {
      FootTarget neutral = Kinematics::neutralStance(i);
      FootTarget sit_target = { neutral.x, neutral.y, neutral.z * 0.4f };
      gait_targets[i].x = sit_target.x + (neutral.x - sit_target.x) * blend;
      gait_targets[i].y = sit_target.y + (neutral.y - sit_target.y) * blend;
      gait_targets[i].z = sit_target.z + (neutral.z - sit_target.z) * blend;
      in_stance[i] = true;
    }
    // Shrink the body pose as a whole until it's physically possible here.
    pose.fitToReach(&gait_targets, 1, dt);
  }

  bool stable = Limits::stancePolygonOk(in_stance);
  (void)stable;   // hook this up to your own response once you're ready - slow down, warn, etc.

  constexpr float MAX_JOINT_DEG_S = 300.0f;   // measure your actual servo speed and set this for real

  for (int i = 0; i < NUM_LEGS; i++) {
    FootTarget posed = pose.apply(gait_targets[i]);
    FootTarget clamped = Limits::clampWorkspace(i, posed);
    JointAngles angles = Kinematics::solve(i, clamped);

    angles.coxa_deg  = Limits::clampJointRate(last_joint_deg[i][0], angles.coxa_deg,  MAX_JOINT_DEG_S, dt);
    angles.femur_deg = Limits::clampJointRate(last_joint_deg[i][1], angles.femur_deg, MAX_JOINT_DEG_S, dt);
    angles.tibia_deg = Limits::clampJointRate(last_joint_deg[i][2], angles.tibia_deg, MAX_JOINT_DEG_S, dt);
    last_joint_deg[i][0] = angles.coxa_deg;
    last_joint_deg[i][1] = angles.femur_deg;
    last_joint_deg[i][2] = angles.tibia_deg;

    if (state != RobotState::DISARMED) {
      servo_driver.writeLeg(i, angles);
    }
  }
}