#include <algorithm>
#include <array>
#include <cstdint>
#include <vector>

#include "catch2/catch_test_macros.hpp"
#include "fsfw/globalfunctions/CobsEncoder.h"
#include "fsfw/returnvalues/returnvalue.h"

// no zero bytes
constexpr std::array<uint8_t, 3> ti0 = {1, 2, 43};
constexpr std::array<uint8_t, 5> to0 = {0x04, 0x01, 0x02, 0x2b, 0x00};

// single zero byte
constexpr std::array<uint8_t, 4> ti1 = {1, 2, 0, 3};
constexpr std::array<uint8_t, 6> to1 = {3, 1, 2, 2, 3, 0};

// multiple zero bytes in sequence
constexpr std::array<uint8_t, 6> ti2 = {1, 2, 0, 0, 0, 3};
constexpr std::array<uint8_t, 8> to2 = {3, 1, 2, 1, 1, 2, 3, 0};

// stuffing bytes required because of the length without a single zero byte
// clang-format off
constexpr std::array<uint8_t, 257> ti3 = {
  1, 2, 3, 4,5 , 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16,
  1, 2, 3, 4,5 , 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16,
  1, 2, 3, 4,5 , 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16,
  1, 2, 3, 4,5 , 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16,
  1, 2, 3, 4,5 , 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16,
  1, 2, 3, 4,5 , 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16,
  1, 2, 3, 4,5 , 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16,
  1, 2, 3, 4,5 , 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16,
  1, 2, 3, 4,5 , 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16,
  1, 2, 3, 4,5 , 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16,
  1, 2, 3, 4,5 , 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16,
  1, 2, 3, 4,5 , 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16,
  1, 2, 3, 4,5 , 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16,
  1, 2, 3, 4,5 , 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16,
  1, 2, 3, 4,5 , 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16,
  1, 2, 3, 4,5 , 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16,
  // #257
  0xaa,
};
constexpr std::array<uint8_t, 260> to3 = {
  0xff,
  1, 2, 3, 4,5 , 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16,
  1, 2, 3, 4,5 , 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16,
  1, 2, 3, 4,5 , 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16,
  1, 2, 3, 4,5 , 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16,
  1, 2, 3, 4,5 , 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16,
  1, 2, 3, 4,5 , 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16,
  1, 2, 3, 4,5 , 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16,
  1, 2, 3, 4,5 , 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16,
  1, 2, 3, 4,5 , 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16,
  1, 2, 3, 4,5 , 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16,
  1, 2, 3, 4,5 , 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16,
  1, 2, 3, 4,5 , 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16,
  1, 2, 3, 4,5 , 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16,
  1, 2, 3, 4,5 , 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16,
  1, 2, 3, 4,5 , 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16,
  1, 2, 3, 4,5 , 6, 7, 8, 9, 10, 11, 12, 13, 14,
  // #255 -> stuffing
  0x04,
  15, 16,
  0xaa, 0
};
// clang-format on

// real packet data

// PLOC startup packets (usually in one or two messages, but 3 separate COBS encoded packets)
// clang-format off
constexpr std::array<uint8_t, 36> ploc_startup_1_i = {
  0x08,0x42,0xc0,0x00,0x00,0x1d,0x20,0x05,0x02,0x00,0x00,0x00,0x00,0x40,0x61,0x0d,0x02,0x8e,
  0xb4,0xef,0x2e,0xe1,0x44,0x00,0x1f,0x03,0x00,0x00,0x52,0x18,0x00,0x00,0x00,0x00,0x1b,0xe2,
};
constexpr std::array<uint8_t, 36> ploc_startup_2_i = {
  0x08,0x42,0xc0,0x01,0x00,0x1d,0x20,0x05,0x02,0x00,0x01,0x00,0x00,0x40,0x61,0x0d,0x02,0x8e,
  0xb4,0xef,0x30,0x74,0x44,0x00,0x1f,0x07,0x00,0x00,0x50,0x03,0x00,0x00,0x00,0x00,0x98,0xe3,
};
constexpr std::array<uint8_t, 36> ploc_startup_3_i = {
  0x08,0x42,0xc0,0x02,0x00,0x1d,0x20,0x05,0x01,0x00,0x02,0x00,0x00,0x40,0x61,0x0d,0x02,0x8e,
  0xb5,0xb8,0x30,0x73,0x44,0x00,0x1f,0x07,0x00,0x00,0x3b,0x04,0x00,0x00,0x00,0x00,0x22,0x68,
};
constexpr std::array<uint8_t, 38> ploc_startup_1_o = {
  0x04,0x08,0x42,0xc0,0x01,0x05,0x1d,0x20,0x05,0x02,
  0x01,0x01,0x01,0x0b,0x40,0x61,0x0d,0x02,0x8e,0xb4,
  0xef,0x2e,0xe1,0x44,0x03,0x1f,0x03,0x01,0x03,0x52,
  0x18,0x01,0x01,0x01,0x03,0x1b,0xe2, 0x00,
};
constexpr std::array<uint8_t, 38> ploc_startup_2_o = {
  0x05, 0x08,
  0x42,0xc0,0x01,0x05,0x1d,0x20,0x05,0x02,0x02,0x01,
  0x01,0x0b,0x40,0x61,0x0d,0x02,0x8e,0xb4,0xef,0x30,
  0x74,0x44,0x03,0x1f,0x07,0x01,0x03,0x50,0x03,0x01,
  0x01,0x01,0x03,0x98,0xe3,0x00
};
constexpr std::array<uint8_t, 38> ploc_startup_3_o = {
  0x05,0x08,0x42,0xc0,0x02,0x05,0x1d,0x20,0x05,0x01,
  0x02,0x02,0x01,0x0b,0x40,0x61,0x0d,0x02,0x8e,0xb5,
  0xb8,0x30,0x73,0x44,0x03,0x1f,0x07,0x01,0x03,0x3b,
  0x04,0x01,0x01,0x01,0x03,0x22,0x68,0x00,
};
// clang-format on

// PLOC Ping (sent from OBC)
// clang-format off
constexpr std::array<uint8_t, 13> ploc_ping_i = {
  // as received by PLOC
  0x18,0x42,0xc0,0x00,0x00,0x06,0x2f,0x11,0x01,0x00,
  0x00,0xd5,0xc8,
};
constexpr std::array<uint8_t, 15> ploc_ping_o = {
  // as sent from OBC
  0x04,0x18,0x42,0xc0,0x01,0x05,0x06,0x2f,0x11,0x01,0x01,0x03,0xd5,0xc8,0x00,
};
// clang-format on
// Corresponding PLOC Pong (sent from PLOC)
// clang-format off
constexpr std::array<uint8_t, 26> ploc_pong_i = {
  // as received by OBC
  0x08,0x42,0xc0,0x05,0x00,0x13,0x20,0x01,0x07,0x00,
  0x01,0x00,0x00,0x40,0x61,0x0d,0x02,0x98,0xbd,0xab,
  0x18,0x42,0xc0,0x00,0x13,0xd5,
};
constexpr std::array<uint8_t, 28> ploc_pong_o = {
  // as sent by PLOC
  0x05,0x08,0x42,0xc0,0x05,0x05,0x13,0x20,0x01,0x07,
  0x02,0x01,0x01,0x0b,0x40,0x61,0x0d,0x02,0x98,0xbd,
  0xab,0x18,0x42,0xc0,0x03,0x13,0xd5,0x00,
};
// clang-format on

#define FOR_EACH_TEST_ARRAY(macro)           \
  macro(ti0, to0);                           \
  macro(ti1, to1);                           \
  macro(ti2, to2);                           \
  macro(ti3, to3);                           \
  macro(ploc_startup_1_i, ploc_startup_1_o); \
  macro(ploc_startup_2_i, ploc_startup_2_o); \
  macro(ploc_startup_3_i, ploc_startup_3_o); \
  macro(ploc_ping_i, ploc_ping_o);           \
  macro(ploc_pong_i, ploc_pong_o);

template <typename FI, typename FE, typename SI>
void check_iter_equal(FI first_begin, FE first_end, SI second_begin) {
  auto first_iter = first_begin;
  auto second_iter = second_begin;
  auto i = 0;
  while (first_iter < first_end) {
    CHECK(*first_iter == *second_iter);
    ++first_iter;
    ++second_iter;
    ++i;
  }

  // safety check
  CHECK(std::equal(first_begin, first_end, second_begin));
}

#define ENCODE_TEST(input, output)                                                  \
  TEST_CASE("COBS encode " #input, "[cobs]") {                                      \
    std::vector<uint8_t> outstream;                                                 \
    outstream.resize(output.size());                                                \
    size_t encodedLen = 0;                                                          \
    auto result = CobsEncoder::encode(input.data(), input.size(), outstream.data(), \
                                      outstream.size(), &encodedLen);               \
    if (result == returnvalue::OK) {                                                \
      CHECK(encodedLen == output.size());                                           \
      check_iter_equal(outstream.begin(), outstream.end(), output.begin());         \
    } else {                                                                        \
      CHECK(result == returnvalue::OK);                                             \
    }                                                                               \
  }

FOR_EACH_TEST_ARRAY(ENCODE_TEST)
#undef ENCODE_TEST

#define DECODE_TEST(input, output)                                                             \
  TEST_CASE("COBS decode " #output, "[cobs]") {                                                \
    std::vector<uint8_t> instream;                                                             \
    instream.resize(input.size());                                                             \
    size_t decodedLen = 0;                                                                     \
    size_t readLen = 0;                                                                        \
    auto result = CobsEncoder::decode(output.data(), output.size(), &readLen, instream.data(), \
                                      instream.size(), &decodedLen);                           \
    if (result == returnvalue::OK) {                                                           \
      CHECK(readLen == output.size());                                                         \
      CHECK(decodedLen == input.size());                                                       \
      check_iter_equal(instream.begin(), instream.end(), input.begin());                       \
    } else {                                                                                   \
      CHECK(result == returnvalue::OK);                                                        \
    }                                                                                          \
  }

FOR_EACH_TEST_ARRAY(DECODE_TEST)
#undef DECODE_TEST

// payloads far past any block count round-trip and the worst case bound is tight enough to hold
// them
TEST_CASE("Large payload", "[cobs]") {
  // no zero bytes at all, so this needs the maximum number of blocks the input length allows
  std::vector<uint8_t> input(64 * 1024, 1);
  std::vector<uint8_t> encoded(CobsEncoder::worstCaseEncodedLen(input.size()));
  size_t encodedLen = 0;

  auto result =
      CobsEncoder::encode(input.data(), input.size(), encoded.data(), encoded.size(), &encodedLen);
  REQUIRE(result == returnvalue::OK);
  CHECK(encodedLen == encoded.size());

  std::vector<uint8_t> decoded(input.size());
  size_t decodedLen = 0;
  size_t readLen = 0;
  result = CobsEncoder::decode(encoded.data(), encodedLen, &readLen, decoded.data(), decoded.size(),
                               &decodedLen);
  REQUIRE(result == returnvalue::OK);
  CHECK(readLen == encodedLen);
  REQUIRE(decodedLen == input.size());
  CHECK(std::equal(decoded.begin(), decoded.end(), input.begin()));
}

// check that too small target buffers are not blindly used
TEST_CASE("Insufficient space", "[cobs]") {
  std::vector<uint8_t> output{0, 2};
  size_t encodedLen = 0;

  auto result =
      CobsEncoder::encode(ti1.data(), ti1.size(), output.data(), output.size(), &encodedLen);
  CHECK(result == CobsEncoder::INSUFFICIENT_SPACE);

  // even equal size should not work
  output.resize(ti1.size());
  result = CobsEncoder::encode(ti1.data(), ti1.size(), output.data(), output.size(), &encodedLen);
  CHECK(result == CobsEncoder::INSUFFICIENT_SPACE);

  // check reverse direction for decoding
  output.resize(2);
  size_t readLen;
  size_t decodedLen;
  result = CobsEncoder::decode(to1.data(), to1.size(), &readLen, output.data(), output.size(),
                               &decodedLen);
  CHECK(result == CobsEncoder::INSUFFICIENT_SPACE);
}

TEST_CASE("Malformed COBS data", "[cobs]") {
  // decode a cut-off packet
  std::vector<uint8_t> input{std::begin(to2), std::end(to2)};
  input.resize(to2.size() / 2);

  std::vector<uint8_t> output;
  output.resize(ti2.size() * 2);
  size_t encodedLen;
  size_t readLen;

  auto result = CobsEncoder::decode(input.data(), input.size(), &readLen, output.data(),
                                    output.size(), &encodedLen);
  CHECK(result == CobsEncoder::STREAM_TOO_SHORT);

  // inject a zero byte somewhere in the middle of a list of nonzero bytes
  input = {10, 1, 2, 3, 4, 0, 6, 7, 8, 9, 0};
  result = CobsEncoder::decode(input.data(), input.size(), &readLen, output.data(), output.size(),
                               &encodedLen);
  CHECK(result == CobsEncoder::DECODING_ERROR);
  // the whole broken frame is skipped, up to and including its delimiter
  CHECK(readLen == 6);
}

// a corrupt frame must not hold back the intact frames queued behind it
TEST_CASE("Resynchronisation after a corrupt frame", "[cobs]") {
  // a block claiming far more bytes than the frame holds, terminated like a real frame
  std::vector<uint8_t> stream{0x40, 0x11, 0x22, 0x00};
  const size_t corruptLen = stream.size();
  stream.insert(stream.end(), to1.begin(), to1.end());

  std::vector<uint8_t> output(ti1.size());
  size_t readLen = 0;
  size_t decodedLen = 0;

  auto result = CobsEncoder::decode(stream.data(), stream.size(), &readLen, output.data(),
                                    output.size(), &decodedLen);
  REQUIRE(result == CobsEncoder::DECODING_ERROR);
  REQUIRE(readLen == corruptLen);

  // the next frame decodes from where the corrupt one ended
  result = CobsEncoder::decode(stream.data() + readLen, stream.size() - readLen, &readLen,
                               output.data(), output.size(), &decodedLen);
  REQUIRE(result == returnvalue::OK);
  CHECK(readLen == to1.size());
  REQUIRE(decodedLen == ti1.size());
  CHECK(std::equal(output.begin(), output.end(), ti1.begin()));
}

// without a delimiter the frame may simply still be arriving, so nothing may be consumed
TEST_CASE("Incomplete frame is not consumed", "[cobs]") {
  std::vector<uint8_t> output(16);
  size_t readLen = 0xdead;
  size_t decodedLen = 0xdead;

  // a self-consistent frame that is merely missing its delimiter
  std::vector<uint8_t> input{3, 1, 2, 2, 3};
  auto result = CobsEncoder::decode(input.data(), input.size(), &readLen, output.data(),
                                    output.size(), &decodedLen);
  CHECK(result == CobsEncoder::STREAM_TOO_SHORT);
  CHECK(readLen == 0);
  CHECK(decodedLen == 0);

  // once the delimiter arrives, the very same bytes decode
  input.push_back(0);
  result = CobsEncoder::decode(input.data(), input.size(), &readLen, output.data(), output.size(),
                               &decodedLen);
  REQUIRE(result == returnvalue::OK);
  CHECK(readLen == input.size());
  CHECK(decodedLen == ti1.size());
}

// an empty payload survives a round trip and is reported as a zero length frame
TEST_CASE("Empty frame", "[cobs]") {
  std::vector<uint8_t> encoded(CobsEncoder::worstCaseEncodedLen(0));
  size_t encodedLen = 0;

  auto result = CobsEncoder::encode(nullptr, 0, encoded.data(), encoded.size(), &encodedLen);
  REQUIRE(result == returnvalue::OK);
  REQUIRE(encodedLen == 2);
  CHECK(encoded[0] == 0x01);
  CHECK(encoded[1] == 0x00);

  std::vector<uint8_t> output(4);
  size_t readLen = 0;
  size_t decodedLen = 0;
  result = CobsEncoder::decode(encoded.data(), encodedLen, &readLen, output.data(), output.size(),
                               &decodedLen);
  CHECK(result == returnvalue::OK);
  CHECK(readLen == 2);
  CHECK(decodedLen == 0);
}