#include "HexapodSimApi.h"
#include "SimHardware.h"
#include "SimRobotAccess.h"
#include "SimBody.h"
#include "../robot/Config.h"    // geometry and pins as the firmware sees them
#include <HexapodProtocol.h>

static_assert(NUM_LEGS == HXSIM_NUM_LEGS, "HxGeometry arrays are sized for 6 legs");

// The real Teensy spins loop() far faster than this; 250 us keeps the
// control-rate jitter under 2% at 60 Hz while staying cheap.
static constexpr uint32_t HXSIM_LOOP_PERIOD_US = 250;

static bool g_started = false;
static uint32_t g_control_ticks = 0;
static bool g_manual_imu = false;       // hxsim_set_imu() called: stop feeding the IMU from the body
static uint64_t g_last_body_us = 0;

// After each control tick: move the kinematic body to the new feet, and feed
// the simulated MPU6050 what it would measure on that body.
static void bodyStep() {
  float feet[18];
  hxsim_get_feet_mm(feet);
  uint64_t now = simhw::nowMicros();
  simbody::update(feet, (now - g_last_body_us) / 1e6f);
  g_last_body_us = now;
  if (!g_manual_imu) {
    float a[3], g[3];
    simbody::imu(a, g);
    simhw::imuSet(a[0], a[1], a[2], g[0], g[1], g[2]);
  }
}

HXSIM_API int32_t hxsim_api_version(void) { return HXSIM_API_VERSION; }
HXSIM_API int32_t hxsim_packet_size(void) { return (int32_t)sizeof(CmdPacket); }

HXSIM_API void hxsim_get_geometry(HxGeometry* g) {
  g->num_legs = NUM_LEGS;
  for (int i = 0; i < NUM_LEGS; i++) {
    g->mount_x_mm[i] = LEG_MOUNT_X_MM[i];
    g->mount_y_mm[i] = LEG_MOUNT_Y_MM[i];
    g->mount_angle_deg[i] = LEG_MOUNT_ANGLE_DEG[i];
  }
  g->coxa_mm = COXA_LEN_MM;
  g->femur_mm = FEMUR_LEN_MM;
  g->tibia_mm = TIBIA_LEN_MM;
  g->stance_radius_mm = STANCE_RADIUS_MM;
  g->stance_height_mm = STANCE_HEIGHT_MM;
  g->coxa_min_deg = COXA_MIN_DEG;   g->coxa_max_deg = COXA_MAX_DEG;
  g->femur_min_deg = FEMUR_MIN_DEG; g->femur_max_deg = FEMUR_MAX_DEG;
  g->tibia_min_deg = TIBIA_MIN_DEG; g->tibia_max_deg = TIBIA_MAX_DEG;
  g->control_hz = CONTROL_HZ;
  g->leg_clearance_mm = LEG_CLEARANCE_MM;
  g->stance_height_range_mm = STANCE_HEIGHT_RANGE_MM;
}

HXSIM_API void hxsim_set_radio_present(int32_t present) { simhw::radioSetPresent(present != 0); }
HXSIM_API void hxsim_set_imu_present(int32_t present)   { simhw::imuSetPresent(present != 0); }

HXSIM_API void hxsim_start(void) {
  if (g_started) return;
  g_started = true;
  simhw::imuSet(0, 0, 1, 0, 0, 0);   // level and still until the body model says otherwise
  simbody::reset();
  simrobot::setupOnce();
  bodyStep();
}

HXSIM_API void hxsim_step(uint32_t dt_us) {
  if (!g_started) hxsim_start();
  while (dt_us > 0) {
    uint32_t chunk = dt_us < HXSIM_LOOP_PERIOD_US ? dt_us : HXSIM_LOOP_PERIOD_US;
    simhw::advanceMicros(chunk);
    dt_us -= chunk;
    uint32_t before = simrobot::lastLoopMicros();
    simrobot::loopOnce();
    if (simrobot::lastLoopMicros() != before) {
      g_control_ticks++;
      bodyStep();
    }
  }
}

HXSIM_API void hxsim_radio_send(const uint8_t* bytes, int32_t len) { simhw::radioSend(bytes, len); }

HXSIM_API void hxsim_set_imu(float ax, float ay, float az, float gx, float gy, float gz) {
  g_manual_imu = true;
  simhw::imuSet(ax, ay, az, gx, gy, gz);
}

HXSIM_API void hxsim_set_ground_tilt(float roll_deg, float pitch_deg) { simbody::setGroundTilt(roll_deg, pitch_deg); }

HXSIM_API void hxsim_get_body(float* pos_mm, float* quat_wxyz, float* yaw_deg, int32_t* contact_mask) {
  if (pos_mm) simbody::position(pos_mm);
  if (quat_wxyz) simbody::orientation(quat_wxyz);
  if (yaw_deg) *yaw_deg = simbody::yawDeg();
  if (contact_mask) *contact_mask = simbody::contactMask();
}

HXSIM_API void hxsim_serial_input(const char* bytes, int32_t len) { simhw::serialFeed(bytes, len); }
HXSIM_API int32_t hxsim_serial_output(char* out, int32_t cap) { return simhw::serialDrain(out, cap); }

HXSIM_API void hxsim_get_servos(int32_t* us, int32_t* attached, uint32_t* writes) {
  for (int i = 0; i < NUM_LEGS; i++)
    for (int j = 0; j < 3; j++) {
      const simhw::ServoPin& p = simhw::servoPin(SERVOS[i][j].pin);
      if (us)       us[i * 3 + j] = p.us;
      if (attached) attached[i * 3 + j] = p.attached ? 1 : 0;
      if (writes)   writes[i * 3 + j] = p.writes;
    }
}

// Inverse of ServoDriver::degToMicros through the same per-servo table:
// what angle the firmware meant by this pulse.
static float pulseToDeg(int leg, int joint, int us) {
  const ServoCal& c = SERVOS[leg][joint];
  return (us - c.us_at_0deg) / c.us_per_deg;
}

HXSIM_API void hxsim_get_joints_deg(float* deg) {
  for (int i = 0; i < NUM_LEGS; i++)
    for (int j = 0; j < 3; j++)
      deg[i * 3 + j] = pulseToDeg(i, j, simhw::servoPin(SERVOS[i][j].pin).us);
}

// Forward kinematics, written independently of robot/Kinematics.cpp so it
// can catch IK mistakes. Joint convention (the one Config.h's limits imply):
//   coxa  - CCW from the leg's mount angle, viewed from above
//   femur - elevation above horizontal, + = up
//   tibia - relative to the femur line, + = up, so a knee-down leg is negative
HXSIM_API void hxsim_get_feet_mm(float* xyz) {
  float deg[NUM_LEGS * 3];
  hxsim_get_joints_deg(deg);
  for (int i = 0; i < NUM_LEGS; i++) {
    float a = radians(LEG_MOUNT_ANGLE_DEG[i] + deg[i * 3 + 0]);
    float f = radians(deg[i * 3 + 1]);
    float t = radians(deg[i * 3 + 1] + deg[i * 3 + 2]);
    float horiz = COXA_LEN_MM + FEMUR_LEN_MM * cosf(f) + TIBIA_LEN_MM * cosf(t);
    xyz[i * 3 + 0] = LEG_MOUNT_X_MM[i] + horiz * cosf(a);
    xyz[i * 3 + 1] = LEG_MOUNT_Y_MM[i] + horiz * sinf(a);
    xyz[i * 3 + 2] = FEMUR_LEN_MM * sinf(f) + TIBIA_LEN_MM * sinf(t);
  }
}

HXSIM_API void hxsim_get_status(HxStatus* s) {
  s->robot_state = simrobot::state();
  s->sit_stand_blend = simrobot::sitStandBlend();
  s->link_ok = simrobot::linkOk() ? 1 : 0;
  s->heading_deg = simrobot::headingDeg();
  s->imu_roll_deg = simrobot::imuRollDeg();
  s->imu_pitch_deg = simrobot::imuPitchDeg();
  s->imu_enabled = simrobot::imuEnabled() ? 1 : 0;
  s->sim_time_ms = (uint32_t)(simhw::nowMicros() / 1000);
  s->control_ticks = g_control_ticks;
  s->pose_reach_scale = simrobot::poseReachScale();
  s->twist_scale = simrobot::twistScale();
  s->stance_height_mm = simrobot::stanceHeightMm();
  s->stance_mask = 0;
  for (int i = 0; i < NUM_LEGS; i++)
    if (simrobot::legInStance(i)) s->stance_mask |= 1 << i;
}