#include "ugv_comms.hh"
#include "ugv_config.h"

#include <esp_log.h>
#include <mbedtls/base64.h>
#include <rom/crc.h>

#include "messages.pb.h"

#ifdef COMMS_SX127X
#include "sx127x_registers.h"
#endif

namespace ugv {
namespace comms {

static const char *TAG = "ugv_comms";
static const TickType_t SEND_PERIOD = pdMS_TO_TICKS(15 * 1000);

CommsClass::CommsClass()
    : last_packet_tick(0),
      last_packet_rssi(INT32_MIN),
      last_packet_snr(INT8_MIN),
      status(),
      drive_heading(),
      new_target(nullptr),
      new_config(nullptr) {
  status.set_state(messages::STATE_IDLE);
  mutex = xSemaphoreCreateMutex();
}

void CommsClass::Init() {
  esp_err_t ret;

  auto *loc = status.mutable_location();
  loc->set_fix_quality(0);
  loc->set_latitude(43.65);
  loc->set_longitude(-116.20);
  loc->set_altitude(2730);
  status.set_state(messages::UGV_State::STATE_IDLE);

#ifdef COMMS_SX127X
  sx127x_config_t lora_config = sx127x_config_default();
  lora_config.sck_io_num = (gpio_num_t)LORA_SCK;
  lora_config.miso_io_num = (gpio_num_t)LORA_MISO;
  lora_config.mosi_io_num = (gpio_num_t)LORA_MOSI;
  lora_config.cs_io_num = (gpio_num_t)LORA_CS;
  lora_config.rst_io_num = (gpio_num_t)LORA_RST;
  lora_config.irq_io_num = (gpio_num_t)LORA_IRQ;
  lora_config.frequency = LORA_FREQ;
  lora_config.tx_power = 17;
  lora_config.spreading_factor = 12;
  lora_config.signal_bandwidth = 10E3;
  lora_config.sync_word = 0x34;
  lora_config.crc = SX127X_CRC_ENABLED;

  ret = sx127x_init(&lora_config, &lora);
  if (ret != ESP_OK) {
    const char *err_name = esp_err_to_name(ret);
    ESP_LOGE(TAG, "LoRa init failed: %s", err_name);
    return;
  }
  ret = sx127x_start(lora);
  if (ret != ESP_OK) {
    ESP_LOGE(TAG, "LoRa start failed: %d", ret);
    return;
  }
  ESP_LOGI(TAG, "LoRa initialized");
#else
  e32::Config lora_config;
  lora_config.uart_port = LORA_UART;
  lora_config.uart_rx_pin = LORA_RX;
  lora_config.uart_tx_pin = LORA_TX;
  ret = lora.Init(lora_config);
  if (ret != ESP_OK) {
    const char *error_name = esp_err_to_name(ret);
    ESP_LOGE(TAG, "E32 LoRa Init failed: %s (%d)", error_name, ret);
    return;
  }
#endif

  xTaskCreate(CommsClass::CommsTask, "ugv_comms", 8 * 1024, this, 2,
              &task_handle);
  ESP_LOGD(TAG, "UGV comms started");
}

void CommsClass::Lock(TickType_t ticks_to_wait) {
  xSemaphoreTake(mutex, ticks_to_wait);
}

void CommsClass::Unlock() { xSemaphoreGive(mutex); }

int32_t CommsClass::ReadRssi() {
  int32_t rssi;
#ifdef COMMS_SX127X
  sx127x_read_rssi(lora, &rssi);
#else
  rssi = 0;
#endif
  return rssi;
}

uint8_t CommsClass::ReadLnaGain() {
  uint8_t lna_gain;
#ifdef COMMS_SX127X
  sx127x_read_lna_gain(lora, &lna_gain);
#else
  lna_gain = 0;
#endif
  return lna_gain;
}

void CommsClass::CommsTask(void *arg) {
  ((CommsClass *)arg)->RunTask();
  vTaskDelete(NULL);
}

void CommsClass::RunTask() {
  using messages::Location;
  using messages::UGV_Message;
  using messages::UGV_State;
  using messages::UGV_Status;

  TickType_t current_tick = xTaskGetTickCount();
  next_status_send = current_tick;

#ifdef COMMS_SX127X
  sx127x_rx_packet_t rx_packet;
#else
  std::string rx_buf;
  int rx_len;
#endif

  while (true) {
    TickType_t delay_ticks = next_status_send - current_tick;

#ifdef COMMS_SX127X
    ret = sx127x_recv_packet(lora, &rx_packet, delay_ticks);
    if (ret == ESP_OK) {
      HandlePacket(&rx_packet);

      sx127x_packet_rx_free(&rx_packet);
    }
#else
    // receiving packet data now
    rx_len = lora.ReadLn(rx_buf, delay_ticks);
    if (rx_len <= 0) {
      ESP_LOGV(TAG, "timeout for packet rx");
      // lora.Flush();
    } else {
      ESP_LOGV(TAG, "rx data: %.*s", rx_buf.size(), rx_buf.c_str());
      HandlePacket((uint8_t *)rx_buf.c_str(), rx_buf.size());
    }
    // TODO: checksum?
#endif

    current_tick = xTaskGetTickCount();

    if (current_tick < next_status_send) {
      continue;
    }

    SendStatus();

    current_tick = xTaskGetTickCount();
  }
}

void CommsClass::SendStatus() {
  Lock();
  status_message.mutable_status()->CopyFrom(this->status);
  Unlock();
  status_message.SerializeToString(&status_message_data);
  esp_err_t ret = SendPacket(status_message_data);
  if (ret != ESP_OK) {
    ESP_LOGE(TAG, "error sending packet: %d", ret);
  } else {
    ESP_LOGV(TAG, "lora wrote UGV_Message packet");
  }

  next_status_send = xTaskGetTickCount() + SEND_PERIOD;
}

#ifdef COMMS_SX127X
void CommsClass::HandlePacket(sx127x_rx_packet_t *rx_packet) {
  ESP_LOGI(TAG, "lora received packet (len %d, rssi: %d, snr: %f): %.*s\n",
           rx_packet->data_len, rx_packet->rssi, rx_packet->snr * 0.25f,
           rx_packet->data_len, rx_packet->data);

  xSemaphoreTake(mutex, portMAX_DELAY);
  last_packet_rssi = rx_packet->rssi;
  last_packet_snr = rx_packet->snr;
  xSemaphoreGive(mutex);

  HandlePacket(rx_packet->data, rx_packet->data_len);
}
#endif

void CommsClass::HandlePacket(const uint8_t *data, size_t data_size) {
  using messages::GroundMessage;

  xSemaphoreTake(mutex, pdMS_TO_TICKS(10));
  last_packet_tick = xTaskGetTickCount();
  xSemaphoreGive(mutex);
  ESP_LOGI(TAG, "rx base64 message: %.*s", data_size, data);

  int ret;
  size_t decoded_size = (data_size * 4 + 2) / 3;
  uint8_t *decoded = new uint8_t[decoded_size];
  size_t decoded_len;

  ret = mbedtls_base64_decode(decoded, decoded_size, &decoded_len, data,
                              data_size);
  if (ret != 0) {
    ESP_LOGE(TAG, "base64 decode error: %d", ret);
    delete[] decoded;
    return;
  }
  if (decoded_len < 4) {
    ESP_LOGE(TAG, "message too short (%d bytes)", decoded_len);
    delete[] decoded;
    return;
  }

  uint32_t calccrc = crc32_le(0, decoded, decoded_len - 4);
  uint32_t msgcrc;
  memcpy(&msgcrc, decoded + decoded_len - 4, 4);
  if (calccrc != msgcrc) {
    ESP_LOGW(TAG, "crc did not match (message %u, calculated %u)", msgcrc, calccrc);
    delete[] decoded;
    return;
  }

  GroundMessage ground_message;

  bool pb_ret = ground_message.ParseFromArray(decoded, decoded_len - 4);
  delete[] decoded;
  if (!pb_ret) {
    ESP_LOGE(TAG, "rx invalid protobuf");
    return;
  }

  switch (ground_message.ground_message_case()) {
    case GroundMessage::kCommand:
      HandleCommand(ground_message.command());
      break;
    default:
      ESP_LOGE(TAG, "invalid ground message: %d",
               ground_message.ground_message_case());
      break;
  }
}

void CommsClass::HandleCommand(const messages::GroundCommand &command) {
  ESP_LOGI(TAG, "rx command id %d type %d", command.id(), command.type());
  using messages::UGV_State;

  switch (command.type()) {
    case messages::CMD_DISABLE:
      xSemaphoreTake(mutex, pdMS_TO_TICKS(10));
      status.set_state(UGV_State::STATE_IDLE);
      xSemaphoreGive(mutex);
      break;
    case messages::CMD_DRIVE_TO_TARGET:
      xSemaphoreTake(mutex, pdMS_TO_TICKS(10));
      status.set_state(UGV_State::STATE_AQUIRING);
      xSemaphoreGive(mutex);
      break;
    case messages::CMD_TEST:
      xSemaphoreTake(mutex, pdMS_TO_TICKS(10));
      status.set_state(UGV_State::STATE_TEST);
      xSemaphoreGive(mutex);
      break;
    case messages::CMD_DRIVE_HEADING: {
      if (command.has_drive_heading()) {
        xSemaphoreTake(mutex, pdMS_TO_TICKS(10));
        status.set_state(UGV_State::STATE_DRIVE_HEADING);
        this->drive_heading = command.drive_heading();
        xSemaphoreGive(mutex);
      } else {
        status.set_state(UGV_State::STATE_IDLE);
        ESP_LOGE(TAG, "invalid CMD_DRIVE_HEADING");
      }
      break;
    }
    case messages::CMD_SET_TARGET: {
      if (command.has_target_location()) {
        xSemaphoreTake(mutex, pdMS_TO_TICKS(10));
        if (!new_target) {
          new_target = new messages::TargetLocation();
        }
        new_target->CopyFrom(command.target_location());
        xSemaphoreGive(mutex);
      }
      break;
    }
    case messages::CMD_SET_CONFIG: {
      if (command.has_config()) {
        xSemaphoreTake(mutex, pdMS_TO_TICKS(10));
        if (!new_config) {
          new_config = new config::Config();
        }
        new_config->CopyFrom(command.config());
        xSemaphoreGive(mutex);
      }
      break;
    }
    case messages::CMD_GET_STATUS: {
      SendStatus();
      break;
    }
    case messages::CMD_PING: {
      break;
    }
    default:
      ESP_LOGW(TAG, "unhandled command type: %d", command.type());
      return;  // early return, no ack
  }

  messages::UGV_Message ugv_message;
  ugv_message.set_command_ack(command.id());

  SendPacket(ugv_message);
}

esp_err_t CommsClass::SendPacket(const uint8_t *data, size_t size) {
#ifdef COMMS_SX127X
  return sx127x_send_packet(lora, data, size, 0);
#else
  if (size >= MAX_PACKET_LEN) {
    ESP_LOGE(TAG, "SendPacket size too long: %d", size);
    return ESP_ERR_INVALID_SIZE;
  }
  size_t fullsize = size + 4;
  uint8_t *data_with_crc = new uint8_t[fullsize];
  uint32_t crc = crc32_le(0, data, size);
  ESP_LOGV(TAG, "Calculated crc value: %u", crc);
  memcpy(data_with_crc, data, size);
  memcpy(data_with_crc + size, &crc, 4);
  size_t encoded_size = ((fullsize + 6) / 3) * 4;
  uint8_t *encoded = new uint8_t[encoded_size];
  size_t encoded_len;
  int ret =
      mbedtls_base64_encode(encoded, encoded_size, &encoded_len, data_with_crc, fullsize);
  delete[] data_with_crc;
  if (ret != 0) {
    delete[] encoded;
    ESP_LOGE(TAG, "base64 encode error: %d", ret);
    return ESP_FAIL;
  }
  lora.WriteLn(encoded, encoded_len);
  delete[] encoded;
  return ESP_OK;
#endif
}

esp_err_t CommsClass::SendPacket(const std::string &str) {
  return SendPacket((const uint8_t *)str.c_str(), str.size());
}

esp_err_t CommsClass::SendPacket(const google::protobuf::MessageLite &message) {
  std::string str = message.SerializeAsString();
  return SendPacket(str);
}

}  // namespace comms
}  // namespace ugv