#include "esp_now_proto.h"
#include "client_registry.h"
#include "pb_decode.h"
#include "pb_encode.h"
#include <string.h>

static bool encode_mac(pb_ostream_t *stream, const pb_field_t *field,
                       void *const *arg) {
  const uint8_t *mac = (const uint8_t *)*arg;
  if (mac == NULL) {
    return true;
  }
  if (!pb_encode_tag_for_field(stream, field)) {
    return false;
  }
  return pb_encode_string(stream, mac, CLIENT_MAC_LEN);
}

static bool decode_mac(pb_istream_t *stream, const pb_field_t *field,
                       void **arg) {
  (void)field;
  uint8_t *mac = (uint8_t *)*arg;
  if (mac == NULL) {
    return false;
  }
  if (stream->bytes_left > CLIENT_MAC_LEN) {
    return false;
  }
  return pb_read(stream, mac, stream->bytes_left);
}

esp_err_t esp_now_proto_encode(const alox_EspNowMessage *msg, uint8_t *buf,
                               size_t buf_size, size_t *out_len) {
  if (msg == NULL || buf == NULL || out_len == NULL || buf_size == 0) {
    return ESP_ERR_INVALID_ARG;
  }

  pb_ostream_t stream = pb_ostream_from_buffer(buf, buf_size);
  if (!pb_encode(&stream, alox_EspNowMessage_fields, msg)) {
    return ESP_FAIL;
  }

  *out_len = stream.bytes_written;
  return ESP_OK;
}

void esp_now_proto_setup_presence_encode(alox_EspNowSlavePresence *presence,
                                         const uint8_t mac[6]) {
  if (presence == NULL) {
    return;
  }
  presence->mac.funcs.encode = encode_mac;
  presence->mac.arg = (void *)mac;
}

void esp_now_proto_setup_message_decode(alox_EspNowMessage *msg,
                                        uint8_t mac_out[6]) {
  if (msg == NULL || mac_out == NULL) {
    return;
  }
  msg->payload.slave_info.mac.funcs.decode = decode_mac;
  msg->payload.slave_info.mac.arg = mac_out;
  msg->payload.heartbeat.mac.funcs.decode = decode_mac;
  msg->payload.heartbeat.mac.arg = mac_out;
}

esp_err_t esp_now_proto_decode(const uint8_t *data, size_t len,
                               alox_EspNowMessage *msg) {
  if (data == NULL || msg == NULL || len == 0) {
    return ESP_ERR_INVALID_ARG;
  }

  pb_istream_t stream = pb_istream_from_buffer(data, len);
  if (!pb_decode(&stream, alox_EspNowMessage_fields, msg)) {
    return ESP_FAIL;
  }
  return ESP_OK;
}

const alox_EspNowSlavePresence *
esp_now_proto_get_presence(const alox_EspNowMessage *msg) {
  if (msg == NULL) {
    return NULL;
  }

  switch (msg->which_payload) {
  case alox_EspNowMessage_slave_info_tag:
    return &msg->payload.slave_info;
  case alox_EspNowMessage_heartbeat_tag:
    return &msg->payload.heartbeat;
  default:
    return NULL;
  }
}