#include <Arduino.h>
#include <unity.h>

#include "dd3_legacy_core.h"
#include "payload_codec.h"

static constexpr uint8_t kMaxSamples = 30;

static void fill_sparse_batch(BatchInput &in) {
  memset(&in, 0, sizeof(in));
  in.sender_id = 1;
  in.batch_id = 42;
  in.t_last = 1700000000;
  in.present_mask = (1UL << 0) | (1UL << 2) | (1UL << 3) | (1UL << 10) | (1UL << 29);
  in.n = 5;
  in.battery_mV = 3750;
  in.err_m = 2;
  in.err_d = 1;
  in.err_tx = 3;
  in.err_last = 2;
  in.err_rx_reject = 1;
  in.energy_wh[0] = 100000;
  in.energy_wh[1] = 100001;
  in.energy_wh[2] = 100050;
  in.energy_wh[3] = 100050;
  in.energy_wh[4] = 100200;
  in.p1_w[0] = -120;
  in.p1_w[1] = -90;
  in.p1_w[2] = 1910;
  in.p1_w[3] = -90;
  in.p1_w[4] = 500;
  in.p2_w[0] = 50;
  in.p2_w[1] = -1950;
  in.p2_w[2] = 60;
  in.p2_w[3] = 2060;
  in.p2_w[4] = -10;
  in.p3_w[0] = 0;
  in.p3_w[1] = 10;
  in.p3_w[2] = -1990;
  in.p3_w[3] = 10;
  in.p3_w[4] = 20;
}

static void fill_full_batch(BatchInput &in) {
  memset(&in, 0, sizeof(in));
  in.sender_id = 1;
  in.batch_id = 0xBEEF;
  in.t_last = 1769904999;
  in.present_mask = 0x3FFFFFFFUL;
  in.n = kMaxSamples;
  in.battery_mV = 4095;
  in.err_m = 10;
  in.err_d = 20;
  in.err_tx = 30;
  in.err_last = 3;
  in.err_rx_reject = 6;
  for (uint8_t i = 0; i < kMaxSamples; ++i) {
    in.energy_wh[i] = 500000UL + static_cast<uint32_t>(i) * static_cast<uint32_t>(i) * 3UL;
    in.p1_w[i] = static_cast<int16_t>(-1000 + static_cast<int16_t>(i) * 25);
    in.p2_w[i] = static_cast<int16_t>(500 - static_cast<int16_t>(i) * 30);
    in.p3_w[i] = static_cast<int16_t>(((i % 2) == 0 ? 100 : -100) + static_cast<int16_t>(i) * 5);
  }
}

static void assert_batch_equals(const BatchInput &expected, const BatchInput &actual) {
  TEST_ASSERT_EQUAL_UINT16(expected.sender_id, actual.sender_id);
  TEST_ASSERT_EQUAL_UINT16(expected.batch_id, actual.batch_id);
  TEST_ASSERT_EQUAL_UINT32(expected.t_last, actual.t_last);
  TEST_ASSERT_EQUAL_UINT32(expected.present_mask, actual.present_mask);
  TEST_ASSERT_EQUAL_UINT8(expected.n, actual.n);
  TEST_ASSERT_EQUAL_UINT16(expected.battery_mV, actual.battery_mV);
  TEST_ASSERT_EQUAL_UINT8(expected.err_m, actual.err_m);
  TEST_ASSERT_EQUAL_UINT8(expected.err_d, actual.err_d);
  TEST_ASSERT_EQUAL_UINT8(expected.err_tx, actual.err_tx);
  TEST_ASSERT_EQUAL_UINT8(expected.err_last, actual.err_last);
  TEST_ASSERT_EQUAL_UINT8(expected.err_rx_reject, actual.err_rx_reject);

  for (uint8_t i = 0; i < expected.n; ++i) {
    TEST_ASSERT_EQUAL_UINT32(expected.energy_wh[i], actual.energy_wh[i]);
    TEST_ASSERT_EQUAL_INT16(expected.p1_w[i], actual.p1_w[i]);
    TEST_ASSERT_EQUAL_INT16(expected.p2_w[i], actual.p2_w[i]);
    TEST_ASSERT_EQUAL_INT16(expected.p3_w[i], actual.p3_w[i]);
  }
  for (uint8_t i = expected.n; i < kMaxSamples; ++i) {
    TEST_ASSERT_EQUAL_UINT32(0, actual.energy_wh[i]);
    TEST_ASSERT_EQUAL_INT16(0, actual.p1_w[i]);
    TEST_ASSERT_EQUAL_INT16(0, actual.p2_w[i]);
    TEST_ASSERT_EQUAL_INT16(0, actual.p3_w[i]);
  }
}

static void test_encode_decode_roundtrip_schema_v3() {
  BatchInput in = {};
  fill_sparse_batch(in);

  uint8_t encoded[256] = {};
  size_t encoded_len = 0;
  TEST_ASSERT_TRUE(encode_batch(in, encoded, sizeof(encoded), &encoded_len));
  TEST_ASSERT_TRUE(encoded_len > 24);

  BatchInput out = {};
  TEST_ASSERT_TRUE(decode_batch(encoded, encoded_len, &out));
  assert_batch_equals(in, out);
}

static void test_decode_rejects_bad_magic_schema_flags() {
  BatchInput in = {};
  fill_sparse_batch(in);

  uint8_t encoded[256] = {};
  size_t encoded_len = 0;
  TEST_ASSERT_TRUE(encode_batch(in, encoded, sizeof(encoded), &encoded_len));

  BatchInput out = {};

  uint8_t bad_magic[256] = {};
  memcpy(bad_magic, encoded, encoded_len);
  bad_magic[0] = 0x00;
  TEST_ASSERT_FALSE(decode_batch(bad_magic, encoded_len, &out));

  uint8_t bad_schema[256] = {};
  memcpy(bad_schema, encoded, encoded_len);
  bad_schema[2] = 0x02;
  TEST_ASSERT_FALSE(decode_batch(bad_schema, encoded_len, &out));

  uint8_t bad_flags[256] = {};
  memcpy(bad_flags, encoded, encoded_len);
  bad_flags[3] = 0x00;
  TEST_ASSERT_FALSE(decode_batch(bad_flags, encoded_len, &out));
}

static void test_decode_rejects_truncated_and_length_mismatch() {
  BatchInput in = {};
  fill_sparse_batch(in);

  uint8_t encoded[256] = {};
  size_t encoded_len = 0;
  TEST_ASSERT_TRUE(encode_batch(in, encoded, sizeof(encoded), &encoded_len));

  BatchInput out = {};
  TEST_ASSERT_FALSE(decode_batch(encoded, encoded_len - 1, &out));
  TEST_ASSERT_FALSE(decode_batch(encoded, 12, &out));

  uint8_t with_tail[257] = {};
  memcpy(with_tail, encoded, encoded_len);
  with_tail[encoded_len] = 0xAA;
  TEST_ASSERT_FALSE(decode_batch(with_tail, encoded_len + 1, &out));
}

static void test_encode_and_decode_reject_invalid_present_mask() {
  BatchInput in = {};
  fill_sparse_batch(in);

  uint8_t encoded[256] = {};
  size_t encoded_len = 0;

  in.present_mask = 0x40000000UL;
  TEST_ASSERT_FALSE(encode_batch(in, encoded, sizeof(encoded), &encoded_len));

  fill_sparse_batch(in);
  TEST_ASSERT_TRUE(encode_batch(in, encoded, sizeof(encoded), &encoded_len));

  BatchInput out = {};

  uint8_t invalid_bits[256] = {};
  memcpy(invalid_bits, encoded, encoded_len);
  invalid_bits[15] |= 0x40;
  TEST_ASSERT_FALSE(decode_batch(invalid_bits, encoded_len, &out));

  uint8_t bitcount_mismatch[256] = {};
  memcpy(bitcount_mismatch, encoded, encoded_len);
  bitcount_mismatch[16] = 0x01; // n=1 while mask has 5 bits set
  TEST_ASSERT_FALSE(decode_batch(bitcount_mismatch, encoded_len, &out));
}

static void test_encode_rejects_invalid_n_and_regression_cases() {
  BatchInput in = {};
  fill_sparse_batch(in);

  uint8_t encoded[256] = {};
  size_t encoded_len = 0;

  in.n = 31;
  TEST_ASSERT_FALSE(encode_batch(in, encoded, sizeof(encoded), &encoded_len));

  fill_sparse_batch(in);
  in.n = 0;
  in.present_mask = 1;
  TEST_ASSERT_FALSE(encode_batch(in, encoded, sizeof(encoded), &encoded_len));

  fill_sparse_batch(in);
  in.n = 2;
  in.present_mask = 0x00000003UL;
  in.energy_wh[1] = in.energy_wh[0] - 1;
  TEST_ASSERT_FALSE(encode_batch(in, encoded, sizeof(encoded), &encoded_len));

  fill_sparse_batch(in);
  TEST_ASSERT_FALSE(encode_batch(in, encoded, 10, &encoded_len));
}

static const uint8_t VECTOR_SYNC_EMPTY[] = {
    0xB3, 0xDD, 0x03, 0x01, 0x01, 0x00, 0x34, 0x12, 0xE4, 0x97, 0x7E, 0x69, 0x00, 0x00, 0x00, 0x00, 0x00, 0xA6, 0x0E,
    0x00, 0x00, 0x00, 0x00, 0x00};

static const uint8_t VECTOR_SPARSE_5[] = {
    0xB3, 0xDD, 0x03, 0x01, 0x01, 0x00, 0x2A, 0x00, 0x00, 0xF1, 0x53, 0x65, 0x0D, 0x04, 0x00, 0x20, 0x05, 0xA6, 0x0E,
    0x02, 0x01, 0x03, 0x02, 0x01, 0xA0, 0x86, 0x01, 0x00, 0x01, 0x31, 0x00, 0x96, 0x01, 0x88, 0xFF, 0x3C, 0xA0, 0x1F,
    0x9F, 0x1F, 0x9C, 0x09, 0x32, 0x00, 0x9F, 0x1F, 0xB4, 0x1F, 0xA0, 0x1F, 0xAB, 0x20, 0x00, 0x00, 0x14, 0x9F, 0x1F,
    0xA0, 0x1F, 0x14};

static const uint8_t VECTOR_FULL_30[] = {
    0xB3, 0xDD, 0x03, 0x01, 0x01, 0x00, 0xEF, 0xBE, 0x67, 0x9B, 0x7E, 0x69, 0xFF, 0xFF, 0xFF, 0x3F, 0x1E, 0xFF, 0x0F,
    0x0A, 0x14, 0x1E, 0x03, 0x06, 0x20, 0xA1, 0x07, 0x00, 0x03, 0x09, 0x0F, 0x15, 0x1B, 0x21, 0x27, 0x2D, 0x33, 0x39,
    0x3F, 0x45, 0x4B, 0x51, 0x57, 0x5D, 0x63, 0x69, 0x6F, 0x75, 0x7B, 0x81, 0x01, 0x87, 0x01, 0x8D, 0x01, 0x93, 0x01,
    0x99, 0x01, 0x9F, 0x01, 0xA5, 0x01, 0xAB, 0x01, 0x18, 0xFC, 0x32, 0x32, 0x32, 0x32, 0x32, 0x32, 0x32, 0x32, 0x32,
    0x32, 0x32, 0x32, 0x32, 0x32, 0x32, 0x32, 0x32, 0x32, 0x32, 0x32, 0x32, 0x32, 0x32, 0x32, 0x32, 0x32, 0x32, 0x32,
    0x32, 0xF4, 0x01, 0x3B, 0x3B, 0x3B, 0x3B, 0x3B, 0x3B, 0x3B, 0x3B, 0x3B, 0x3B, 0x3B, 0x3B, 0x3B, 0x3B, 0x3B, 0x3B,
    0x3B, 0x3B, 0x3B, 0x3B, 0x3B, 0x3B, 0x3B, 0x3B, 0x3B, 0x3B, 0x3B, 0x3B, 0x3B, 0x64, 0x00, 0x85, 0x03, 0x9A, 0x03,
    0x85, 0x03, 0x9A, 0x03, 0x85, 0x03, 0x9A, 0x03, 0x85, 0x03, 0x9A, 0x03, 0x85, 0x03, 0x9A, 0x03, 0x85, 0x03, 0x9A,
    0x03, 0x85, 0x03, 0x9A, 0x03, 0x85, 0x03, 0x9A, 0x03, 0x85, 0x03, 0x9A, 0x03, 0x85, 0x03, 0x9A, 0x03, 0x85, 0x03,
    0x9A, 0x03, 0x85, 0x03, 0x9A, 0x03, 0x85, 0x03, 0x9A, 0x03, 0x85, 0x03};

static void test_payload_golden_vectors() {
  BatchInput expected_sync = {};
  expected_sync.sender_id = 1;
  expected_sync.batch_id = 0x1234;
  expected_sync.t_last = 1769904100;
  expected_sync.present_mask = 0;
  expected_sync.n = 0;
  expected_sync.battery_mV = 3750;
  expected_sync.err_m = 0;
  expected_sync.err_d = 0;
  expected_sync.err_tx = 0;
  expected_sync.err_last = 0;
  expected_sync.err_rx_reject = 0;

  BatchInput expected_sparse = {};
  fill_sparse_batch(expected_sparse);

  BatchInput expected_full = {};
  fill_full_batch(expected_full);

  struct VectorCase {
    const char *name;
    const uint8_t *bytes;
    size_t len;
    const BatchInput *expected;
  } cases[] = {
      {"sync_empty", VECTOR_SYNC_EMPTY, sizeof(VECTOR_SYNC_EMPTY), &expected_sync},
      {"sparse_5", VECTOR_SPARSE_5, sizeof(VECTOR_SPARSE_5), &expected_sparse},
      {"full_30", VECTOR_FULL_30, sizeof(VECTOR_FULL_30), &expected_full},
  };

  for (size_t i = 0; i < (sizeof(cases) / sizeof(cases[0])); ++i) {
    BatchInput decoded = {};
    TEST_ASSERT_TRUE_MESSAGE(decode_batch(cases[i].bytes, cases[i].len, &decoded), cases[i].name);
    assert_batch_equals(*cases[i].expected, decoded);

    uint8_t reencoded[512] = {};
    size_t reencoded_len = 0;
    TEST_ASSERT_TRUE_MESSAGE(encode_batch(*cases[i].expected, reencoded, sizeof(reencoded), &reencoded_len), cases[i].name);
    TEST_ASSERT_EQUAL_UINT_MESSAGE(cases[i].len, reencoded_len, cases[i].name);
    TEST_ASSERT_EQUAL_UINT8_ARRAY_MESSAGE(cases[i].bytes, reencoded, cases[i].len, cases[i].name);
  }
}

void setup() {
  dd3_legacy_core_force_link();
  UNITY_BEGIN();
  RUN_TEST(test_encode_decode_roundtrip_schema_v3);
  RUN_TEST(test_decode_rejects_bad_magic_schema_flags);
  RUN_TEST(test_decode_rejects_truncated_and_length_mismatch);
  RUN_TEST(test_encode_and_decode_reject_invalid_present_mask);
  RUN_TEST(test_encode_rejects_invalid_n_and_regression_cases);
  RUN_TEST(test_payload_golden_vectors);
  UNITY_END();
}

void loop() {}