#include "pro_codec.h"
#include "unitree_legged_sdk/a1_const.h"
#include "unitree_legged_sdk/aliengo_const.h"
#include "unitree_legged_sdk/b1_const.h"
#include "unitree_legged_sdk/go1_const.h"
#include "unitree_legged_sdk/unitree_legged_sdk.h"

#include <array>
#include <atomic>
#include <chrono>
#include <cmath>
#include <cstdint>
#include <fstream>
#include <iterator>
#include <string>
#include <thread>
#include <type_traits>
#include <vector>

using namespace UNITREE_LEGGED_SDK;

namespace {

#define CHECK(expression) do { if (!(expression)) return __LINE__; } while (false)

uint64_t Fnv1a(const uint8_t* data, std::size_t size) {
  uint64_t result = 1469598103934665603ull;
  for (std::size_t i = 0; i < size; ++i) {
    result ^= data[i];
    result *= 1099511628211ull;
  }
  return result;
}

std::vector<uint8_t> ReadFile(const std::string& path) {
  std::ifstream stream(path, std::ios::binary);
  return std::vector<uint8_t>(
      (std::istreambuf_iterator<char>(stream)), std::istreambuf_iterator<char>());
}

}  // namespace

int main() {
  static_assert(std::is_same<decltype(LowCmd::motorCmd), std::array<MotorCmd, 20>>::value,
                "official LowCmd field shape changed");
  static_assert(sizeof(IMU) == 53, "IMU must match official packed declaration");
  static_assert(sizeof(MotorCmd) == 33, "MotorCmd must match official packed declaration");
  static_assert(sizeof(MotorState) == 38, "MotorState must match official packed declaration");
  static_assert(sizeof(BmsState) == 34, "BmsState must match official packed declaration");
  static_assert(sizeof(LowCmd) == 734, "LowCmd must match official packed declaration");
  static_assert(sizeof(LowState) == 937, "LowState must match official packed declaration");
  static_assert(sizeof(HighCmd) == 129, "HighCmd must match official packed declaration");
  static_assert(sizeof(HighState) == 1087, "HighState must match official packed declaration");
  static_assert(go1_Hip_max > 1.0 && a1_Hip_max < 1.0 &&
                    aliengo_Hip_max > 1.0 && b1_Hip_max < 1.0,
                "official robot constant headers must be available");

  go1_pro_internal::ProCodec codec(go1_pro_internal::ProCodec::FindStateFile());
  LowCmd initialized{};
  go1_pro_internal::InitLowCmd(initialized);
  CHECK(initialized.head[0] == 0xfe && initialized.head[1] == 0xef);
  CHECK(initialized.levelFlag == LOWLEVEL && initialized.bandWidth == 0x3ac0);
  for (const auto& motor : initialized.motorCmd) {
    CHECK(motor.mode == 0x0a);
    CHECK(motor.q == static_cast<float>(PosStopF));
    CHECK(motor.dq == static_cast<float>(VelStopF));
    CHECK(motor.Kp == 0.0f && motor.Kd == 0.0f && motor.tau == 0.0f);
  }
  std::array<uint8_t, 8> zero{};
  std::array<uint8_t, 8> encrypted{};
  codec.Encrypt(zero.data(), encrypted.data(), encrypted.size());
  const std::array<uint8_t, 8> known = {{0x12, 0xaa, 0x03, 0x54, 0x23, 0x6b, 0x66, 0xe3}};
  CHECK(encrypted == known);

  LowCmd cmd{};
  cmd.head = {{0xfe, 0xef}};
  cmd.levelFlag = LOWLEVEL;
  cmd.bandWidth = 0x3ac0;
  const auto packet = codec.EncodeLowCmd(cmd);
  CHECK(Fnv1a(packet.data(), packet.size()) == 0x63c5ef1a61977e18ull);

#if GO1_PRO_HAVE_CAPTURE_FIXTURES
  const auto real_lowcmd = ReadFile(GO1_PRO_TEST_LOWCMD);
  const auto real_lowcmd_plain = ReadFile(GO1_PRO_TEST_LOWCMD_PLAIN);
  CHECK(real_lowcmd.size() == go1_pro_internal::kLowCmdWireSize);
  CHECK(real_lowcmd_plain.size() == go1_pro_internal::kLowCmdWireSize);
  std::vector<uint8_t> decrypted_lowcmd(real_lowcmd.size());
  codec.Decrypt(real_lowcmd.data(), decrypted_lowcmd.data(), real_lowcmd.size());
  CHECK(decrypted_lowcmd == real_lowcmd_plain);
  std::vector<uint8_t> encrypted_lowcmd(real_lowcmd_plain.size());
  codec.Encrypt(real_lowcmd_plain.data(), encrypted_lowcmd.data(), real_lowcmd_plain.size());
  CHECK(encrypted_lowcmd == real_lowcmd);

  const auto state_packet = ReadFile(GO1_PRO_TEST_CAPTURE);
  CHECK(state_packet.size() == go1_pro_internal::kLowStateDatagramSize);
  LowState state{};
  CHECK(codec.DecodeLowState(state_packet.data(), state_packet.size(), state));
  CHECK(state.head[0] == 0xfe && state.head[1] == 0xef);
  CHECK(state.levelFlag == LOWLEVEL);
  CHECK(state.motorState[FR_1].temperature == 50);
  CHECK(std::isfinite(state.motorState[FR_1].q));
  CHECK(state.wirelessRemote[0] == 0x55 && state.wirelessRemote[1] == 0x51);

  const auto old_state_packet = ReadFile(GO1_PRO_TEST_CAPTURE_OLD);
  CHECK(old_state_packet.size() == go1_pro_internal::kLowStateDatagramSize);
  LowState old_state{};
  CHECK(codec.DecodeLowState(old_state_packet.data(), old_state_packet.size(), old_state));
  CHECK(old_state.bms.SOC == 76);
  CHECK(old_state.motorState[FR_0].temperature == 52);
#endif

  Safety safety(LeggedType::Go1);
  cmd.motorCmd[FR_0].q = 4.0f;
  safety.PositionLimit(cmd);
  CHECK(cmd.motorCmd[FR_0].q == 0.78f);
  cmd.motorCmd[FR_1].q = static_cast<float>(PosStopF);
  safety.PositionLimit(cmd);
  CHECK(cmd.motorCmd[FR_1].q == static_cast<float>(PosStopF));

  std::atomic<int> callbacks{0};
  LoopFunc loop("compat-test", 0.001f, Callback([&callbacks] { ++callbacks; }));
  loop.start();
  std::this_thread::sleep_for(std::chrono::milliseconds(10));
  loop.shutdown();
  CHECK(callbacks.load() > 0);
  return 0;
}