#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"

// 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().
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;
      case 'w': test_packet.vy = min(test_packet.vy + 200, 1000);  test_packet.flags |= FLAG_WALK; break;
      case 'x': test_packet.vy = max(test_packet.vy - 200, -1000); test_packet.flags |= FLAG_WALK; break;
      case 'q': test_packet.vx = max(test_packet.vx - 200, -1000); test_packet.flags |= FLAG_WALK; break;
      case 'e': test_packet.vx = min(test_packet.vx + 200, 1000);  test_packet.flags |= FLAG_WALK; break;
      case 'z': test_packet.yaw_rate = max(test_packet.yaw_rate - 200, -1000); test_packet.flags |= FLAG_WALK; break;
      case 'c': test_packet.yaw_rate = min(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();

  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);

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

  // Simple open-loop heading estimate: accumulates the omega we've been
  // commanding. 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 on your surface.
  current_heading_deg += degrees(mapped.twist.omega) * 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();
  }
  last_state = state;

  gait.setRequestedPattern(mapped.gait_wave ? &GAIT_WAVE : &GAIT_TRIPOD);

  // Remote-gyro toggle off -> drop commanded tilt but keep XYZ offset/height.
  PoseCommand pose_cmd = mapped.pose;
  if (!mapped.remote_gyro_enabled) {
    pose_cmd.roll_deg = 0;
    pose_cmd.pitch_deg = 0;
  }
  pose.setCommand(pose_cmd);
  pose.setCorrection(mapped.robot_gyro_enabled ? -imu.rollDeg() : 0,
                      mapped.robot_gyro_enabled ? -imu.pitchDeg() : 0);
  pose.update(dt);

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

  if (state == RobotState::WALKING) {
    gait.update(dt, mapped.twist, gait_targets);
    for (int i = 0; i < NUM_LEGS; i++) in_stance[i] = gait.legInStance(i);
  } 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;
    }
  }

  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);
    }
  }
}