// Copyright 2022 Dolphin Emulator Project
// SPDX-License-Identifier: GPL-2.0-or-later

#include "Core/HW/WiimoteEmu/DesiredWiimoteState.h"

#include <cstring>
#include <optional>
#include <type_traits>
#include <variant>

#include "Common/BitUtils.h"
#include "Common/CommonTypes.h"
#include "Common/VariantUtil.h"

#include "Core/HW/WiimoteEmu/MotionPlus.h"

namespace WiimoteEmu
{
SerializedWiimoteState SerializeDesiredState(const DesiredWiimoteState& state)
{
  const u8 has_buttons = (state.buttons.hex & WiimoteCommon::ButtonData::BUTTON_MASK) != 0 ? 1 : 0;
  const u8 has_accel = state.acceleration != DesiredWiimoteState::DEFAULT_ACCELERATION ? 1 : 0;
  const u8 has_camera = state.camera_points != DesiredWiimoteState::DEFAULT_CAMERA ? 1 : 0;
  const u8 has_motion_plus = state.motion_plus.has_value() ? 1 : 0;

  // Right now we support < 16 extensions so the info which extension is in use fits into 4 bits.
  // This allows 'empty' packets to be a single byte, which is very nice for reducing bandwidth.
  // If we ever support 16 or more we have to redesign this a bit; ideally use a variable-length
  // encoding so that typical extensions (None, Nunchuk, Classic Controller) still fit into the
  // initial 4 bits.
  static_assert(std::variant_size_v<DesiredExtensionState::ExtensionData> <= (1 << 4));
  const u8 extension = u8(state.extension.data.index());

  SerializedWiimoteState s;
  s.length = 0;
  s.data[s.length++] = u8(has_buttons | (has_accel << 1) | (has_camera << 2) |
                          (has_motion_plus << 3) | (extension << 4));

  if (has_buttons)
  {
    const u8 buttons = u8((state.buttons.a) | (state.buttons.b << 1) | (state.buttons.plus << 2) |
                          (state.buttons.minus << 3) | (state.buttons.one << 4) |
                          (state.buttons.two << 5) | (state.buttons.home << 6));
    const u8 dpad = u8((state.buttons.up) | (state.buttons.down << 1) | (state.buttons.left << 2) |
                       (state.buttons.right << 3));
    s.data[s.length++] = buttons;
    s.data[s.length++] = dpad;
  }

  if (has_accel)
  {
    const u16 accel_x = state.acceleration.value.x;  // 10 bits
    const u16 accel_y = state.acceleration.value.y;  // 9 bits (ignore lowest bit)
    const u16 accel_z = state.acceleration.value.z;  // 9 bits (ignore lowest bit)
    const u8 accel_x_high = u8(accel_x >> 2);
    const u8 accel_y_high = u8(accel_y >> 2);
    const u8 accel_z_high = u8(accel_z >> 2);
    const u8 accel_low = u8((accel_x & 0b11) | (Common::ExtractBit<1>(accel_y) << 2) |
                            (Common::ExtractBit<1>(accel_z) << 3));

    if (has_buttons)
    {
      // can use the high bits of the dpad field from buttons
      s.data[s.length - 1] |= u8(accel_low << 4);
    }
    else
    {
      s.data[s.length++] = u8(accel_low << 4);
    }

    s.data[s.length++] = accel_x_high;
    s.data[s.length++] = accel_y_high;
    s.data[s.length++] = accel_z_high;
  }

  if (has_camera)
  {
    for (auto& camera_point : state.camera_points)
    {
      IRExtended irext;
      irext.SetPosition(camera_point.position);
      irext.size = camera_point.size;
      Common::BitCastPtr<IRExtended>(&s.data[s.length]) = irext;
      s.length += sizeof(irext);
    }
  }

  if (has_motion_plus)
  {
    const u16 pitch_slow = state.motion_plus->is_slow.x ? 1 : 0;
    const u16 roll_slow = state.motion_plus->is_slow.y ? 1 : 0;
    const u16 yaw_slow = state.motion_plus->is_slow.z ? 1 : 0;
    const u16 pitch_value = state.motion_plus->gyro.value.x;  // 14 bits
    const u16 roll_value = state.motion_plus->gyro.value.y;   // 14 bits
    const u16 yaw_value = state.motion_plus->gyro.value.z;    // 14 bits
    s.data[s.length++] = u8(pitch_value);
    s.data[s.length++] = u8(((pitch_value >> 8) & 0x3f) | (pitch_slow << 7));
    s.data[s.length++] = u8(roll_value);
    s.data[s.length++] = u8(((roll_value >> 8) & 0x3f) | (roll_slow << 7));
    s.data[s.length++] = u8(yaw_value);
    s.data[s.length++] = u8(((yaw_value >> 8) & 0x3f) | (yaw_slow << 7));
  }

  if (extension)
  {
    std::visit(
        [&s](const auto& arg) {
          using T = std::decay_t<decltype(arg)>;
          static_assert(sizeof(arg) <= 6);
          Common::BitCastPtr<T>(&s.data[s.length]) = arg;
          s.length += sizeof(arg);
        },
        state.extension.data);
  }

  return s;
}

bool DeserializeDesiredState(DesiredWiimoteState* state, const SerializedWiimoteState& serialized)
{
  // clear state
  state->buttons.hex = 0;
  state->acceleration = DesiredWiimoteState::DEFAULT_ACCELERATION;
  state->camera_points = DesiredWiimoteState::DEFAULT_CAMERA;
  state->motion_plus = std::nullopt;
  state->extension.data = std::monostate();

  if (serialized.length < 1)
  {
    // can't be valid
    return false;
  }

  const auto& d = serialized.data;
  const u8 has_buttons = d[0] & 1;
  const u8 has_accel = (d[0] >> 1) & 1;
  const u8 has_camera = (d[0] >> 2) & 1;
  const u8 has_motion_plus = (d[0] >> 3) & 1;
  const u8 extension = (d[0] >> 4);

  if (extension >= ExtensionNumber::MAX)
  {
    // invalid extension
    return false;
  }

  const size_t expected_size = [&] {
    size_t s = 1;
    if (has_buttons && has_accel)
      s += 5;
    else if (has_buttons)
      s += 2;
    else if (has_accel)
      s += 4;
    if (has_camera)
      s += 12;
    if (has_motion_plus)
      s += 6;
    if (extension)
    {
      WithVariantAlternative<DesiredExtensionState::ExtensionData>(
          extension, [&]<typename T>() { s += sizeof(T); });
    }
    return s;
  }();

  if (serialized.length != expected_size)
  {
    // invalid length
    return false;
  }

  size_t pos = 1;

  if (has_buttons)
  {
    state->buttons.a = d[pos] & 1;
    state->buttons.b = (d[pos] >> 1) & 1;
    state->buttons.plus = (d[pos] >> 2) & 1;
    state->buttons.minus = (d[pos] >> 3) & 1;
    state->buttons.one = (d[pos] >> 4) & 1;
    state->buttons.two = (d[pos] >> 5) & 1;
    state->buttons.home = (d[pos] >> 6) & 1;
    state->buttons.up = d[pos + 1] & 1;
    state->buttons.down = (d[pos + 1] >> 1) & 1;
    state->buttons.left = (d[pos + 1] >> 2) & 1;
    state->buttons.right = (d[pos + 1] >> 3) & 1;
    pos += 2;
  }

  if (has_accel)
  {
    if (has_buttons)
      pos -= 1;
    const u8 accel_low = d[pos] >> 4;
    const u8 accel_x_high = d[pos + 1];
    const u8 accel_y_high = d[pos + 2];
    const u8 accel_z_high = d[pos + 3];
    state->acceleration.value.x = (accel_x_high << 2) | (accel_low & 0b11);
    state->acceleration.value.y =
        Common::ExpandValue<u16>((accel_y_high << 1) | Common::ExtractBit<2>(accel_low), 1);
    state->acceleration.value.z =
        Common::ExpandValue<u16>((accel_z_high << 1) | Common::ExtractBit<3>(accel_low), 1);
    pos += 4;
  }

  if (has_camera)
  {
    for (auto& camera_point : state->camera_points)
    {
      const IRExtended irext = Common::BitCastPtr<IRExtended>(&d[pos]);
      camera_point = CameraPoint(irext.GetPosition(), irext.size);
      pos += sizeof(irext);
    }
  }

  if (has_motion_plus)
  {
    const u16 pitch_value = d[pos] | ((d[pos + 1] & 0x3f) << 8);
    const u16 roll_value = d[pos + 2] | ((d[pos + 3] & 0x3f) << 8);
    const u16 yaw_value = d[pos + 4] | ((d[pos + 5] & 0x3f) << 8);
    const bool pitch_slow = (d[pos + 1] & 0x80) != 0;
    const bool roll_slow = (d[pos + 3] & 0x80) != 0;
    const bool yaw_slow = (d[pos + 5] & 0x80) != 0;
    state->motion_plus = MotionPlus::DataFormat::Data{
        MotionPlus::DataFormat::GyroRawValue{
            MotionPlus::DataFormat::GyroType(pitch_value, roll_value, yaw_value)},
        MotionPlus::DataFormat::SlowType(pitch_slow, roll_slow, yaw_slow)};
    pos += 6;
  }

  if (extension)
  {
    WithVariantAlternative<DesiredExtensionState::ExtensionData>(extension, [&]<typename T>() {
      state->extension.data.emplace<T>(Common::BitCastPtr<T>(&d[pos]));
    });
  }

  return true;
}
}  // namespace WiimoteEmu
