#include "payload_codec.h"
#include <limits.h>

static constexpr uint16_t kMagic = 0xDDB3;
static constexpr uint8_t kSchema = 2;
static constexpr uint8_t kFlags = 0x01;
static constexpr size_t kMaxSamples = 30;

static void write_u16_le(uint8_t *dst, uint16_t value) {
  dst[0] = static_cast<uint8_t>(value & 0xFF);
  dst[1] = static_cast<uint8_t>((value >> 8) & 0xFF);
}

static void write_u32_le(uint8_t *dst, uint32_t value) {
  dst[0] = static_cast<uint8_t>(value & 0xFF);
  dst[1] = static_cast<uint8_t>((value >> 8) & 0xFF);
  dst[2] = static_cast<uint8_t>((value >> 16) & 0xFF);
  dst[3] = static_cast<uint8_t>((value >> 24) & 0xFF);
}

static uint16_t read_u16_le(const uint8_t *src) {
  return static_cast<uint16_t>(src[0]) | (static_cast<uint16_t>(src[1]) << 8);
}

static uint32_t read_u32_le(const uint8_t *src) {
  return static_cast<uint32_t>(src[0]) |
         (static_cast<uint32_t>(src[1]) << 8) |
         (static_cast<uint32_t>(src[2]) << 16) |
         (static_cast<uint32_t>(src[3]) << 24);
}

size_t uleb128_encode(uint32_t v, uint8_t *out, size_t cap) {
  size_t i = 0;
  do {
    if (i >= cap) {
      return 0;
    }
    uint8_t byte = static_cast<uint8_t>(v & 0x7F);
    v >>= 7;
    if (v != 0) {
      byte |= 0x80;
    }
    out[i++] = byte;
  } while (v != 0);
  return i;
}

bool uleb128_decode(const uint8_t *in, size_t len, size_t *pos, uint32_t *v) {
  if (!in || !pos || !v) {
    return false;
  }
  uint32_t result = 0;
  uint8_t shift = 0;
  size_t p = *pos;
  for (uint8_t i = 0; i < 5; ++i) {
    if (p >= len) {
      return false;
    }
    uint8_t byte = in[p++];
    if (i == 4 && (byte & 0xF0) != 0) {
      return false;
    }
    result |= static_cast<uint32_t>(byte & 0x7F) << shift;
    if ((byte & 0x80) == 0) {
      *pos = p;
      *v = result;
      return true;
    }
    shift = static_cast<uint8_t>(shift + 7);
  }
  return false;
}

uint32_t zigzag32(int32_t x) {
  return (static_cast<uint32_t>(x) << 1) ^ static_cast<uint32_t>(x >> 31);
}

int32_t unzigzag32(uint32_t u) {
  return static_cast<int32_t>((u >> 1) ^ (static_cast<uint32_t>(-static_cast<int32_t>(u & 1))));
}

size_t svarint_encode(int32_t x, uint8_t *out, size_t cap) {
  uint32_t zz = zigzag32(x);
  return uleb128_encode(zz, out, cap);
}

bool svarint_decode(const uint8_t *in, size_t len, size_t *pos, int32_t *x) {
  uint32_t u = 0;
  if (!uleb128_decode(in, len, pos, &u)) {
    return false;
  }
  *x = unzigzag32(u);
  return true;
}

static bool ensure_capacity(size_t needed, size_t cap, size_t pos) {
  return pos + needed <= cap;
}

bool encode_batch(const BatchInput &in, uint8_t *out, size_t out_cap, size_t *out_len) {
  if (!out || !out_len) {
    return false;
  }
  if (in.n > kMaxSamples) {
    return false;
  }
  if (in.dt_s == 0) {
    return false;
  }
  size_t pos = 0;
  if (!ensure_capacity(21, out_cap, pos)) {
    return false;
  }
  write_u16_le(&out[pos], kMagic);
  pos += 2;
  out[pos++] = kSchema;
  out[pos++] = kFlags;
  write_u16_le(&out[pos], in.sender_id);
  pos += 2;
  write_u16_le(&out[pos], in.batch_id);
  pos += 2;
  write_u32_le(&out[pos], in.t_last);
  pos += 4;
  out[pos++] = in.dt_s;
  out[pos++] = in.n;
  write_u16_le(&out[pos], in.battery_mV);
  pos += 2;
  out[pos++] = in.err_m;
  out[pos++] = in.err_d;
  out[pos++] = in.err_tx;
  out[pos++] = in.err_last;
  out[pos++] = in.err_rx_reject;

  if (in.n == 0) {
    *out_len = pos;
    return true;
  }

  if (!ensure_capacity(4, out_cap, pos)) {
    return false;
  }
  write_u32_le(&out[pos], in.energy_wh[0]);
  pos += 4;
  for (uint8_t i = 1; i < in.n; ++i) {
    if (in.energy_wh[i] < in.energy_wh[i - 1]) {
      return false;
    }
    uint32_t delta = in.energy_wh[i] - in.energy_wh[i - 1];
    size_t wrote = uleb128_encode(delta, &out[pos], out_cap - pos);
    if (wrote == 0) {
      return false;
    }
    pos += wrote;
  }

  auto encode_phase = [&](const int16_t *phase) -> bool {
    if (!ensure_capacity(2, out_cap, pos)) {
      return false;
    }
    write_u16_le(&out[pos], static_cast<uint16_t>(phase[0]));
    pos += 2;
    for (uint8_t i = 1; i < in.n; ++i) {
      int32_t delta = static_cast<int32_t>(phase[i]) - static_cast<int32_t>(phase[i - 1]);
      size_t wrote = svarint_encode(delta, &out[pos], out_cap - pos);
      if (wrote == 0) {
        return false;
      }
      pos += wrote;
    }
    return true;
  };

  if (!encode_phase(in.p1_w)) {
    return false;
  }
  if (!encode_phase(in.p2_w)) {
    return false;
  }
  if (!encode_phase(in.p3_w)) {
    return false;
  }

  *out_len = pos;
  return true;
}

bool decode_batch(const uint8_t *buf, size_t len, BatchInput *out) {
  if (!buf || !out) {
    return false;
  }
  size_t pos = 0;
  if (len < 21) {
    return false;
  }
  uint16_t magic = read_u16_le(&buf[pos]);
  pos += 2;
  uint8_t schema = buf[pos++];
  uint8_t flags = buf[pos++];
  if (magic != kMagic || schema != kSchema || (flags & 0x01) == 0) {
    return false;
  }
  out->sender_id = read_u16_le(&buf[pos]);
  pos += 2;
  out->batch_id = read_u16_le(&buf[pos]);
  pos += 2;
  out->t_last = read_u32_le(&buf[pos]);
  pos += 4;
  out->dt_s = buf[pos++];
  out->n = buf[pos++];
  out->battery_mV = read_u16_le(&buf[pos]);
  pos += 2;
  out->err_m = buf[pos++];
  out->err_d = buf[pos++];
  out->err_tx = buf[pos++];
  out->err_last = buf[pos++];
  out->err_rx_reject = buf[pos++];

  if (out->n > kMaxSamples || out->dt_s == 0) {
    return false;
  }
  if (out->n == 0) {
    for (uint8_t i = 0; i < kMaxSamples; ++i) {
      out->energy_wh[i] = 0;
      out->p1_w[i] = 0;
      out->p2_w[i] = 0;
      out->p3_w[i] = 0;
    }
    return pos == len;
  }
  if (pos + 4 > len) {
    return false;
  }
  out->energy_wh[0] = read_u32_le(&buf[pos]);
  pos += 4;
  for (uint8_t i = 1; i < out->n; ++i) {
    uint32_t delta = 0;
    if (!uleb128_decode(buf, len, &pos, &delta)) {
      return false;
    }
    uint64_t sum = static_cast<uint64_t>(out->energy_wh[i - 1]) + static_cast<uint64_t>(delta);
    if (sum > UINT32_MAX) {
      return false;
    }
    out->energy_wh[i] = static_cast<uint32_t>(sum);
  }

  auto decode_phase = [&](int16_t *phase) -> bool {
    if (pos + 2 > len) {
      return false;
    }
    phase[0] = static_cast<int16_t>(read_u16_le(&buf[pos]));
    pos += 2;
    int32_t prev = static_cast<int32_t>(phase[0]);
    for (uint8_t i = 1; i < out->n; ++i) {
      int32_t delta = 0;
      if (!svarint_decode(buf, len, &pos, &delta)) {
        return false;
      }
      int32_t value = prev + delta;
      if (value < INT16_MIN || value > INT16_MAX) {
        return false;
      }
      phase[i] = static_cast<int16_t>(value);
      prev = value;
    }
    return true;
  };

  if (!decode_phase(out->p1_w)) {
    return false;
  }
  if (!decode_phase(out->p2_w)) {
    return false;
  }
  if (!decode_phase(out->p3_w)) {
    return false;
  }

  for (uint8_t i = out->n; i < kMaxSamples; ++i) {
    out->energy_wh[i] = 0;
    out->p1_w[i] = 0;
    out->p2_w[i] = 0;
    out->p3_w[i] = 0;
  }

  return pos == len;
}

#ifdef PAYLOAD_CODEC_TEST
bool payload_codec_self_test() {
  BatchInput in = {};
  in.sender_id = 1;
  in.batch_id = 42;
  in.t_last = 1700000000;
  in.dt_s = 1;
  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;

  uint8_t buf[256];
  size_t len = 0;
  if (!encode_batch(in, buf, sizeof(buf), &len)) {
    Serial.println("payload_codec_self_test: encode failed");
    return false;
  }

  BatchInput out = {};
  if (!decode_batch(buf, len, &out)) {
    Serial.println("payload_codec_self_test: decode failed");
    return false;
  }

  if (out.sender_id != in.sender_id || out.batch_id != in.batch_id || out.t_last != in.t_last ||
      out.dt_s != in.dt_s || out.n != in.n || out.battery_mV != in.battery_mV ||
      out.err_m != in.err_m || out.err_d != in.err_d || out.err_tx != in.err_tx || out.err_last != in.err_last ||
      out.err_rx_reject != in.err_rx_reject) {
    Serial.println("payload_codec_self_test: header mismatch");
    return false;
  }

  for (uint8_t i = 0; i < in.n; ++i) {
    if (out.energy_wh[i] != in.energy_wh[i] || out.p1_w[i] != in.p1_w[i] || out.p2_w[i] != in.p2_w[i] ||
        out.p3_w[i] != in.p3_w[i]) {
      Serial.println("payload_codec_self_test: sample mismatch");
      return false;
    }
  }

  Serial.printf("payload_codec_self_test: ok len=%u\n", static_cast<unsigned>(len));
  return true;
}
#endif