#include "PublicApp.h"

#include <stddef.h>
#include <string.h>

#include <EEPROM.h>
#include <Utils.h>
#include <avr/io.h>
#include <target.h>

extern "C" void MeshCoreWatchdogStage(uint8_t stage);

namespace {
const int kPrefsEEPROMAddr = 0;
const uint32_t kPrefsMagic = 0x50585253UL; // SRXP
const uint8_t kPrefsVersion = 2;

#if defined(USE_SX1262)
const int8_t kMaxTxPower = 22;
#else
const int8_t kMaxTxPower = 17;
#endif

struct StoredPrefs {
  uint32_t magic;
  uint8_t version;
  uint8_t size;
  uint16_t crc;
  PublicApp::RadioSettings settings;
  char node_name[16];
};

void buildTaggedText(char* dest, size_t dest_size, const char* name, const char* text) {
  size_t di = 0;
  if (dest_size == 0) {
    return;
  }
  while (name && *name && di + 1 < dest_size) {
    dest[di++] = *name++;
  }
  if (di + 3 < dest_size) {
    dest[di++] = ':';
    dest[di++] = ' ';
  }
  while (text && *text && di + 1 < dest_size) {
    dest[di++] = *text++;
  }
  dest[di] = 0;
}

uint16_t crc16(const uint8_t* data, size_t len) {
  uint16_t crc = 0xFFFF;
  for (size_t i = 0; i < len; ++i) {
    crc ^= data[i];
    for (uint8_t bit = 0; bit < 8; ++bit) {
      if (crc & 1) {
        crc = (crc >> 1) ^ 0xA001;
      } else {
        crc >>= 1;
      }
    }
  }
  return crc;
}
}

const uint8_t PublicApp::kPublicSecret[PUB_KEY_SIZE] = {
  0x8B, 0x33, 0x87, 0xE9, 0xC5, 0xCD, 0xEA, 0x6A,
  0xC9, 0xE5, 0xED, 0xBA, 0xA1, 0x15, 0xCD, 0x72,
  0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00,
  0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00,
};

PublicApp::PublicApp(mesh::Radio& radio, mesh::MillisecondClock& ms, mesh::PacketManager& mgr, PublicRTCClock& rtc, FlashStore& store)
  : mesh::Dispatcher(radio, ms, mgr),
    _rtc(&rtc),
    _store(&store),
    _next_seen(0),
    _next_outbound(0),
    _revision(0),
    _raw_rx_count(0),
    _parsed_rx_count(0),
    _reject_count(0),
    _tx_done_count(0),
    _tx_fail_count(0),
    _last_raw_len(0),
    _last_rx_type(0),
    _last_reject_code('-'),
    _last_rssi(0),
    _last_snr(0),
    _reject_type_count(0),
    _reject_hash_count(0),
    _reject_dup_count(0),
    _reject_mac_count(0),
    _reset_cause(0),
    _wdt_stage('-'),
    _public_sent_count(0),
    _public_recv_count(0),
    _public_recv_line_count(0),
    _started(false) {
  _radio_settings = defaultRadioSettings();
  memset(_node_name, 0, sizeof(_node_name));
  memset(_seen, 0, sizeof(_seen));
  memset(_recent_outbound, 0, sizeof(_recent_outbound));
}

void PublicApp::sanitizeAscii(char* dest, size_t dest_size, const char* src) {
  size_t di = 0;
  if (dest_size == 0) {
    return;
  }
  while (src && *src && di + 1 < dest_size) {
    char c = *src++;
    if (c >= 32 && c <= 126) {
      dest[di++] = c;
    }
  }
  dest[di] = 0;
}

void PublicApp::sanitizeDisplayText(char* dest, size_t dest_size, const uint8_t* src, size_t src_len) {
  size_t di = 0;
  if (dest_size == 0) {
    return;
  }
  for (size_t si = 0; src && si < src_len && di + 1 < dest_size; ++si) {
    uint8_t c = src[si];
    if (c == 0) {
      break;
    }
    if (c >= 32 && c <= 127) {
      dest[di++] = (char)c;
    } else if (c >= 0x80) {
      dest[di++] = (char)127;
      while (si + 1 < src_len && (src[si + 1] & 0xC0) == 0x80) {
        ++si;
      }
    }
  }
  dest[di] = 0;
}

PublicApp::RadioSettings PublicApp::defaultRadioSettings() {
  RadioSettings settings;
  settings.freq_khz = (uint32_t)(LORA_FREQ * 1000.0f + 0.5f);
  settings.bw_x10 = (uint16_t)(LORA_BW * 10.0f + 0.5f);
  settings.sf = LORA_SF;
  settings.cr = LORA_CR;
  settings.tx_power = LORA_TX_POWER;
  return settings;
}

bool PublicApp::validateRadioSettings(const RadioSettings& settings) {
  if (settings.freq_khz < 902000UL || settings.freq_khz > 928000UL) {
    return false;
  }
  if (settings.sf < 7 || settings.sf > 12) {
    return false;
  }
  if (settings.cr < 5 || settings.cr > 8) {
    return false;
  }
  if (settings.tx_power < 2 || settings.tx_power > kMaxTxPower) {
    return false;
  }

  static const uint16_t allowed_bw_x10[] = { 625, 1250, 2500, 5000 };
  for (uint8_t i = 0; i < sizeof(allowed_bw_x10) / sizeof(allowed_bw_x10[0]); ++i) {
    if (settings.bw_x10 == allowed_bw_x10[i]) {
      return true;
    }
  }
  return false;
}

const char* PublicApp::getResetCauseText() const {
#ifdef WDRF
  if (_reset_cause & _BV(WDRF)) {
    return "watchdog";
  }
#endif
#ifdef BORF
  if (_reset_cause & _BV(BORF)) {
    return "brownout";
  }
#endif
#ifdef EXTRF
  if (_reset_cause & _BV(EXTRF)) {
    return "external pin";
  }
#endif
#ifdef JTRF
  if (_reset_cause & _BV(JTRF)) {
    return "JTAG";
  }
#endif
#ifdef PORF
  if (_reset_cause & _BV(PORF)) {
    return "power-on";
  }
#endif
  return "unknown";
}

char PublicApp::getWatchdogStage() const {
  return (_wdt_stage >= 32 && _wdt_stage <= 126) ? (char)_wdt_stage : '?';
}

void PublicApp::loadRadioSettings() {
  StoredPrefs prefs;
  EEPROM.get(kPrefsEEPROMAddr, prefs);
  if (prefs.magic == kPrefsMagic &&
      prefs.version == kPrefsVersion &&
      prefs.size == sizeof(StoredPrefs) &&
      prefs.crc == crc16((const uint8_t*)&prefs.settings, sizeof(prefs) - offsetof(StoredPrefs, settings)) &&
      validateRadioSettings(prefs.settings)) {
    _radio_settings = prefs.settings;
    sanitizeAscii(_node_name, sizeof(_node_name), prefs.node_name);
    if (_node_name[0] == 0) {
      strncpy(_node_name, "MSXE", sizeof(_node_name) - 1);
      _node_name[sizeof(_node_name) - 1] = 0;
    }
    return;
  }

  _radio_settings = defaultRadioSettings();
  strncpy(_node_name, "MSXE", sizeof(_node_name) - 1);
  _node_name[sizeof(_node_name) - 1] = 0;
  saveRadioSettings();
}

void PublicApp::saveRadioSettings() {
  StoredPrefs prefs;
  memset(&prefs, 0, sizeof(prefs));
  prefs.magic = kPrefsMagic;
  prefs.version = kPrefsVersion;
  prefs.size = sizeof(StoredPrefs);
  prefs.settings = _radio_settings;
  strncpy(prefs.node_name, _node_name, sizeof(prefs.node_name) - 1);
  prefs.node_name[sizeof(prefs.node_name) - 1] = 0;
  prefs.crc = crc16((const uint8_t*)&prefs.settings, sizeof(prefs) - offsetof(StoredPrefs, settings));
  EEPROM.put(kPrefsEEPROMAddr, prefs);
}

void PublicApp::applyRadioSettings(bool reinit) {
  if (_started && reinit) {
    radio_reinit();
  }
  radio_set_params(_radio_settings.freq_khz / 1000.0f,
                   _radio_settings.bw_x10 / 10.0f,
                   _radio_settings.sf,
                   _radio_settings.cr);
  radio_set_tx_power(_radio_settings.tx_power);
  if (_started) {
    _radio->begin();
  }
}

bool PublicApp::resetRadio() {
  bool ok = radio_reinit();
  radio_set_params(_radio_settings.freq_khz / 1000.0f,
                   _radio_settings.bw_x10 / 10.0f,
                   _radio_settings.sf,
                   _radio_settings.cr);
  radio_set_tx_power(_radio_settings.tx_power);
  if (_started) {
    _radio->begin();
  }
  bumpRevision();
  return ok;
}

bool PublicApp::setRadioSettings(const RadioSettings& settings, bool persist) {
  if (!validateRadioSettings(settings)) {
    return false;
  }
  _radio_settings = settings;
  applyRadioSettings(true);
  if (persist) {
    saveRadioSettings();
  }
  bumpRevision();
  return true;
}

void PublicApp::resetRadioSettingsToDefaults() {
  _radio_settings = defaultRadioSettings();
  applyRadioSettings();
  saveRadioSettings();
  bumpRevision();
}

void PublicApp::begin() {
  loadRadioSettings();
  _public_sent_count = _store->getSentCount();
  _public_recv_count = _store->getRecvCount();
  _public_recv_line_count = 0;
  applyRadioSettings();
  mesh::Dispatcher::begin();
  _started = true;
  bumpRevision();
}

void PublicApp::loop() {
  MeshCoreWatchdogStage('D');
  mesh::Dispatcher::loop();
  MeshCoreWatchdogStage('A');
}

bool PublicApp::getMessage(uint16_t idx, PublicMessage& out) const {
  if (!_store->getMessage(idx, out)) {
    return false;
  }
  out.rebroadcasted = out.outbound && isRebroadcasted(out.hash);
  return true;
}

void PublicApp::rememberOutbound(const uint8_t* hash) {
  if (!hash) {
    return;
  }
  RecentOutboundState& state = _recent_outbound[_next_outbound];
  state.valid = true;
  state.rebroadcasted = false;
  memcpy(state.hash, hash, sizeof(state.hash));
  _next_outbound = (_next_outbound + 1) % kRecentOutboundCount;
}

bool PublicApp::isRebroadcasted(const uint8_t* hash) const {
  if (!hash) {
    return false;
  }
  for (uint8_t i = 0; i < kRecentOutboundCount; ++i) {
    const RecentOutboundState& state = _recent_outbound[i];
    if (state.valid && memcmp(state.hash, hash, sizeof(state.hash)) == 0) {
      return state.rebroadcasted;
    }
  }
  return false;
}

bool PublicApp::hasSeen(const mesh::Packet& pkt) {
  uint8_t hash[MAX_HASH_SIZE];
  pkt.calculatePacketHash(hash);
  for (uint8_t i = 0; i < kSeenCount; ++i) {
    if (memcmp(_seen[i], hash, sizeof(hash)) == 0) {
      return true;
    }
  }
  return false;
}

void PublicApp::markSeen(const mesh::Packet& pkt) {
  pkt.calculatePacketHash(_seen[_next_seen]);
  _next_seen = (_next_seen + 1) % kSeenCount;
}

bool PublicApp::markSentRebroadcast(const mesh::Packet& pkt) {
  if (!pkt.isRouteFlood() || pkt.getPathHashCount() == 0) {
    return false;
  }
  uint8_t hash[MAX_HASH_SIZE];
  pkt.calculatePacketHash(hash);
  for (uint8_t i = 0; i < kRecentOutboundCount; ++i) {
    RecentOutboundState& state = _recent_outbound[i];
    if (!state.valid || memcmp(state.hash, hash, sizeof(state.hash)) != 0) {
      continue;
    }
    if (!state.rebroadcasted) {
      state.rebroadcasted = true;
      _last_reject_code = '+';
      bumpRevision();
    }
    return true;
  }
  return false;
}

mesh::Packet* PublicApp::createPublicPacket(const char* text) {
  char safe_text[kDraftMax + 1];
  sanitizeAscii(safe_text, sizeof(safe_text), text);
  if (safe_text[0] == 0) {
    return NULL;
  }

  uint8_t plain[5 + kPublicMaxTaggedText + 1];
  uint32_t timestamp = _rtc->getCurrentTime();
  memcpy(&plain[0], &timestamp, sizeof(timestamp));
  plain[4] = 0;

  buildTaggedText((char*)&plain[5], sizeof(plain) - 5, _node_name, safe_text);

  int plain_len = 5 + (int)strnlen((const char*)&plain[5], sizeof(plain) - 5);
  int enc_len = ((plain_len + CIPHER_BLOCK_SIZE - 1) / CIPHER_BLOCK_SIZE) * CIPHER_BLOCK_SIZE;
  if (plain_len > kPublicMaxPlain || 1 + CIPHER_MAC_SIZE + enc_len > MAX_PACKET_PAYLOAD) {
    return NULL;
  }

  mesh::Packet* packet = obtainNewPacket();
  if (!packet) {
    return NULL;
  }

  packet->header = (PAYLOAD_TYPE_GRP_TXT << PH_TYPE_SHIFT);
  packet->payload[0] = kPublicHash;
  packet->payload_len = 1 + mesh::Utils::encryptThenMAC(kPublicSecret, &packet->payload[1], plain, plain_len);
  return packet;
}

void PublicApp::sendFlood(mesh::Packet* packet, uint32_t delay_millis, uint8_t path_hash_size) {
  if (!packet || path_hash_size == 0 || path_hash_size > 3) {
    return;
  }
  packet->header &= ~PH_ROUTE_MASK;
  packet->header |= ROUTE_TYPE_FLOOD;
  packet->setPathHashSizeAndCount(path_hash_size, 0);
  markSeen(*packet);
  sendPacket(packet, 1, delay_millis);
}

bool PublicApp::sendPublic(const char* text) {
  char safe_text[kDraftMax + 1];
  char display_text[sizeof(PublicMessage::text)];
  uint8_t hash[MAX_HASH_SIZE];
  sanitizeAscii(safe_text, sizeof(safe_text), text);

  mesh::Packet* packet = createPublicPacket(text);
  if (!packet) {
    return false;
  }
  uint32_t timestamp = _rtc->getCurrentTime();
  packet->calculatePacketHash(hash);
  buildTaggedText(display_text, sizeof(display_text), _node_name, safe_text);
  if (!_store->appendMessage(timestamp, display_text, true, hash)) {
    releasePacket(packet);
    return false;
  }
  rememberOutbound(hash);
  sendFlood(packet);
  ++_public_sent_count;
  bumpRevision();
  return true;
}

bool PublicApp::setNodeName(const char* name) {
  char sanitized[sizeof(_node_name)];
  sanitizeAscii(sanitized, sizeof(sanitized), name);
  if (sanitized[0] == 0) {
    strncpy(sanitized, "MSXE", sizeof(sanitized) - 1);
    sanitized[sizeof(sanitized) - 1] = 0;
  }
  if (strcmp(_node_name, sanitized) == 0) {
    return true;
  }
  strncpy(_node_name, sanitized, sizeof(_node_name) - 1);
  _node_name[sizeof(_node_name) - 1] = 0;
  saveRadioSettings();
  bumpRevision();
  return true;
}

void PublicApp::reject(char code) {
  _last_reject_code = code;
  ++_reject_count;
  switch (code) {
    case 'T':
      ++_reject_type_count;
      break;
    case 'H':
      ++_reject_hash_count;
      break;
    case 'D':
      ++_reject_dup_count;
      break;
    case 'M':
      ++_reject_mac_count;
      break;
  }
}

void PublicApp::syncClockFromTimestamp(uint32_t timestamp) {
  if (timestamp != 0xFFFFFFFFUL && timestamp > _rtc->getCurrentTime()) {
    _rtc->setCurrentTime(timestamp + 1);
  }
}

void PublicApp::logRxRaw(float snr, float rssi, const uint8_t raw[], int len) {
  (void)raw;
  ++_raw_rx_count;
  _last_raw_len = (len > 255) ? 255 : (uint8_t)len;
  _last_snr = (int8_t)snr;
  _last_rssi = (int8_t)rssi;
}

void PublicApp::logRx(mesh::Packet* packet, int len, float score) {
  (void)len;
  (void)score;
  ++_parsed_rx_count;
  _last_rx_type = packet ? packet->getPayloadType() : 0xFF;
}

void PublicApp::logTx(mesh::Packet* packet, int len) {
  (void)packet;
  (void)len;
  ++_tx_done_count;
}

void PublicApp::logTxFail(mesh::Packet* packet, int len) {
  (void)packet;
  (void)len;
  ++_tx_fail_count;
}

mesh::DispatcherAction PublicApp::onRecvPacket(mesh::Packet* pkt) {
  MeshCoreWatchdogStage('C');
  if (!pkt || pkt->getPayloadType() != PAYLOAD_TYPE_GRP_TXT) {
    reject('T');
    return ACTION_RELEASE;
  }
  if (pkt->payload_len <= (1 + CIPHER_MAC_SIZE)) {
    reject('L');
    return ACTION_RELEASE;
  }
  if (pkt->payload[0] != kPublicHash) {
    reject('H');
    return ACTION_RELEASE;
  }
  if (markSentRebroadcast(*pkt)) {
    return ACTION_RELEASE;
  }
  if (hasSeen(*pkt)) {
    reject('D');
    return ACTION_RELEASE;
  }
  markSeen(*pkt);

  int mac_data_len = pkt->payload_len - 1;
  int cipher_len = mac_data_len - CIPHER_MAC_SIZE;
  if (cipher_len <= 0 || (cipher_len % CIPHER_BLOCK_SIZE) != 0 || cipher_len > MAX_PACKET_PAYLOAD) {
    reject('L');
    return ACTION_RELEASE;
  }

  uint8_t plain[MAX_PACKET_PAYLOAD];
  MeshCoreWatchdogStage('c');
  int plain_len = mesh::Utils::MACThenDecrypt(kPublicSecret, plain, &pkt->payload[1], mac_data_len);
  MeshCoreWatchdogStage('p');
  if (plain_len <= 5) {
    reject('M');
    return ACTION_RELEASE;
  }
  MeshCoreWatchdogStage('v');
  if ((plain[4] >> 2) != 0) {
    reject('Y');
    return ACTION_RELEASE;
  }

  uint32_t timestamp = 0;
  memcpy(&timestamp, &plain[0], sizeof(timestamp));

  MeshCoreWatchdogStage('x');
  char* text = (char*)&plain[5];
  size_t max_len = plain_len - 5;
  sanitizeDisplayText(text, sizeof(PublicMessage::text), &plain[5], max_len);
  if (text[0] != 0) {
    syncClockFromTimestamp(timestamp);
    MeshCoreWatchdogStage('m');
    MeshCoreWatchdogStage('w');
    if (!_store->appendMessage(timestamp, text, false)) {
      reject('W');
      return ACTION_RELEASE;
    }
    ++_public_recv_count;
    _public_recv_line_count += wrappedLineCount(text, PublicMessage::kWrapCols);
    _last_reject_code = 'O';
    MeshCoreWatchdogStage('o');
    bumpRevision();
  } else {
    reject('E');
  }
  return ACTION_RELEASE;
}