#include "buscfg.h"
#include "driver/uart.h"
#include "esp_log.h"
#include "esp_timer.h"
#include "freertos/FreeRTOS.h"
#include "freertos/task.h"
#include "mbedtls/base64.h"
#include "packetstructs.h"
#include "string.h"
#include "sx1262.h"
#include <ctype.h>
#include <stdint.h>
#include <stdio.h>
#include <string.h>
#include "mbedtls/aes.h"
#include "mbedtls/sha256.h"
#include "mbedtls/md.h"

#define TAG "canZem"

// requires at least a 256 byte data
FrameStruct decodeFrame(unsigned char *data, unsigned char dataLen) {
  FrameStruct frame;
  memset(&frame, 0, sizeof(frame));
  unsigned char index = 0;
  frame.header = data[index++];
  if ((frame.header & ROUTE_TYPE_MASK) == ROUTE_TYPE_TRANSPORT_DIRECT ||
      (frame.header & ROUTE_TYPE_MASK) == ROUTE_TYPE_TRANSPORT_FLOOD) {
    memcpy(frame.transportCodes, data + index, 4);
    index += 4;
  }
  frame.pathLen = data[index++];

  memcpy(frame.path, data + index, frame.pathLen);
  index += frame.pathLen;
  frame.payloadLen = dataLen - index;
  memcpy(frame.payload, data + index, frame.payloadLen);

  return frame;
}


#define KEY_SIZE    16   // 128-bit AES
#define HMAC_SIZE   2   // SHA256 output size

int aes_encrypt_ecb(const uint8_t *key,
                    const uint8_t *input, size_t ilen,
                    uint8_t *output)
{
    if (ilen % 16 != 0) {
        return -1; // must be multiple of 16
    }

    mbedtls_aes_context ctx;
    mbedtls_aes_init(&ctx);
    mbedtls_aes_setkey_enc(&ctx, key, 128);

    for (size_t offset = 0; offset < ilen; offset += 16) {
        mbedtls_aes_crypt_ecb(&ctx, MBEDTLS_AES_ENCRYPT,
                              input + offset, output + offset);
    }

    mbedtls_aes_free(&ctx);
    return 0;
}

// AES-ECB decrypt (same as Arduino's aes.decryptBlock)
int aes_decrypt_ecb(const uint8_t *key,
                    const uint8_t *input, size_t ilen,
                    uint8_t *output)
{
    if (ilen % 16 != 0) {
        return -1; // must be multiple of 16
    }

    mbedtls_aes_context ctx;
    mbedtls_aes_init(&ctx);
    mbedtls_aes_setkey_dec(&ctx, key, 128);

    for (size_t offset = 0; offset < ilen; offset += 16) {
        mbedtls_aes_crypt_ecb(&ctx, MBEDTLS_AES_DECRYPT,
                              input + offset, output + offset);
    }

    mbedtls_aes_free(&ctx);
    return 0;
}

// HMAC-SHA256
int hmac_sha256(const uint8_t *key, size_t keylen,
                const uint8_t *input, size_t ilen,
                uint8_t *output)
{
    const mbedtls_md_info_t *md_info = mbedtls_md_info_from_type(MBEDTLS_MD_SHA256);
    return mbedtls_md_hmac(md_info, key, keylen, input, ilen, output);
}



// Verify MAC + Decrypt
int mac_then_decrypt(const uint8_t *aes_key,
                     const uint8_t *input, size_t ilen,
                     uint8_t *plaintext, size_t *plen)
{
    if (ilen <= HMAC_SIZE) {
        fprintf(stderr, "[mac_then_decrypt] input too short (ilen=%zu)\n", ilen);
        return -1;
    }

    const uint8_t *mac = input;
    const uint8_t *ciphertext = input + HMAC_SIZE;
    size_t clen = ilen - HMAC_SIZE;

    uint8_t calc_mac[32]; // full SHA256
    int ret = hmac_sha256(aes_key, KEY_SIZE, ciphertext, clen, calc_mac);
    if (ret != 0) {
        fprintf(stderr, "[mac_then_decrypt] hmac_sha256() failed: %d\n", ret);
        return -4;
    }

    // Debug dump of MACs
    //fprintf(stderr, "[mac_then_decrypt] expected MAC (from input): ");
    //for (size_t i = 0; i < HMAC_SIZE; i++) fprintf(stderr, "%02X", mac[i]);
    //fprintf(stderr, "\n");

    //fprintf(stderr, "[mac_then_decrypt] calculated MAC: ");
    //for (size_t i = 0; i < HMAC_SIZE; i++) fprintf(stderr, "%02X", calc_mac[i]);
    //fprintf(stderr, "\n");

    // compare only first HMAC_SIZE bytes
    if (memcmp(mac, calc_mac, HMAC_SIZE) != 0) {
        fprintf(stderr, "[mac_then_decrypt] MAC check failed\n");
        return -2;
    }

    // AES-ECB decrypt (parity with Arduino)
    ret = aes_decrypt_ecb(aes_key, ciphertext, clen, plaintext);
    if (ret != 0) {
        fprintf(stderr, "[mac_then_decrypt] aes_decrypt_ecb() failed: %d\n", ret);
        return -3;
    }

    *plen = clen;
    //fprintf(stderr, "[mac_then_decrypt] success, decrypted %zu bytes\n", *plen);
    return 0;
}


void hexdump(const char *label, const uint8_t *data, size_t len) {
    if (label)
        printf("%s (len=%zu):\n", label, len);

    for (size_t i = 0; i < len; i += 16) {
        printf("%04zx  ", i);  // offset
        for (size_t j = 0; j < 16; j++) {
            if (i + j < len)
                printf("%02X ", data[i + j]);
            else
                printf("   ");  // pad spacing
        }
        printf(" ");
        for (size_t j = 0; j < 16 && i + j < len; j++) {
            uint8_t c = data[i + j];
            printf("%c", isprint(c) ? c : '.');
        }
        printf("\n");
    }
}

GroupTextMessage decodeGroupMessage(FrameStruct frame) {
  GroupTextMessage msg;
  memset(&msg, 0, sizeof(msg));
  if ((frame.header & PAYLOAD_TYPE_MASK) != PAYLOAD_TYPE_GRP_TXT) {
    return msg;
  }
  unsigned char index = 0;
  msg.channelHash = frame.payload[index++];
  unsigned char tmp[200];

  const uint8_t pubAes[16] = {0x8b, 0x33, 0x87, 0xe9, 0xc5, 0xcd, 0xea, 0x6a, 0xc9, 0xe5,
     0xed, 0xba, 0xa1, 0x15, 0xcd, 0x72};

  uint8_t aes_key[16];

  if (msg.channelHash == 17) {
    memcpy(aes_key, pubAes, sizeof(aes_key));
  } else {
    return msg;
  }

  size_t plaintextLen = 0;

  unsigned char res = mac_then_decrypt(aes_key, frame.payload + index, frame.payloadLen - index, tmp, &plaintextLen);
  if (plaintextLen == 0 || res != 0) {
    printf("error decrypting");
  }

  index = 0;
  
  memcpy(&msg.timestamp, tmp + index, 4);
  index += 4;
  msg.flags = tmp[index++];

  memcpy(msg.text, tmp + index, plaintextLen - index);
  return msg;
}

void printGroupMessage(GroupTextMessage msg) {
  printf("Message with channel hash %d, flags %d: %s\n", msg.channelHash, msg.flags, msg.text);
}

AdvertisementPayload decodeAdvertisement(FrameStruct frame) {
  AdvertisementPayload advert;
  memset(&advert, 0, sizeof(advert));

  if ((frame.header & PAYLOAD_TYPE_MASK) != PAYLOAD_TYPE_ADVERT) {
    return advert;
  }

  unsigned char index = 0;

  memcpy(advert.pubKey, frame.payload + index, 32);
  index += 32;

  memcpy(&advert.timestamp, frame.payload + index, 4);
  index += 4;

  memcpy(advert.signature, frame.payload + index, 64);
  index += 64;

  advert.dataFlags = frame.payload[index++];

  if (advert.dataFlags & ADVERTISEMENT_FLAG_HAS_LOCATION) {
    memcpy(&advert.latitude, frame.payload + index, 4);
    index += 4;
    memcpy(&advert.longitude, frame.payload + index, 4);
    index += 4;
  }

  if (advert.dataFlags & ADVERTISEMENT_FLAG_RFU1) {
    memcpy(&advert.rfu1, frame.payload + index, 2);
    index += 2;
  }
  if (advert.dataFlags & ADVERTISEMENT_FLAG_RFU2) {
    memcpy(&advert.rfu2, frame.payload + index, 2);
    index += 2;
  }
  memcpy(advert.nodeName, frame.payload + index, frame.payloadLen - index);
  advert.nodeName[frame.payloadLen - index] = 0;

  return advert;
}

void printAdvertisement(AdvertisementPayload advert) {
  size_t keyB64len = 0;
  size_t sigB64len = 0;
  unsigned char keyBuf[50];
  unsigned char sigBuf[90];
  memset(keyBuf, 0, sizeof(keyBuf));
  memset(sigBuf, 0, sizeof(sigBuf));
  mbedtls_base64_encode(keyBuf, sizeof(keyBuf), &keyB64len, advert.pubKey, 32);
  mbedtls_base64_encode(sigBuf, sizeof(sigBuf), &sigB64len, advert.signature,
                        64);

  printf("%s on %ld with type %s on %s location %ld %ld, public key %s and "
         "signature %s\n",
         advert.dataFlags & ADVERTISEMENT_FLAG_HAS_NAME ? advert.nodeName
                                                        : "nameless node",
         advert.timestamp,
         (advert.dataFlags & 0x07) == 0x04
             ? "sensor"
             : ((advert.dataFlags & 0x07) == 0x03
                    ? "room server"
                    : ((advert.dataFlags & 0x07) == 0x02 ? "repeater"
                                                         : "chat node")),
         advert.dataFlags & 0x80 ? "known" : "unknown", advert.latitude,
         advert.longitude, keyBuf, sigBuf);
}

void printFrameHeader(FrameStruct frame) {
  char strBuf[512] = "Frame route is ";
  switch (frame.header & ROUTE_TYPE_MASK) {
  case ROUTE_TYPE_TRANSPORT_FLOOD:
    strcat(strBuf, "transport flood");
    break;

  case ROUTE_TYPE_FLOOD:
    strcat(strBuf, "flood");
    break;

  case ROUTE_TYPE_DIRECT:
    strcat(strBuf, "direct");
    break;

  case ROUTE_TYPE_TRANSPORT_DIRECT:
    strcat(strBuf, "transport direct");
    break;
  }

  strcat(strBuf, ", payload type is ");

  switch (frame.header & PAYLOAD_TYPE_MASK) {
  case PAYLOAD_TYPE_REQ:
    strcat(strBuf, "request");
    break;

  case PAYLOAD_TYPE_RESPONSE:
    strcat(strBuf, "response");
    break;

  case PAYLOAD_TYPE_TXT_MSG:
    strcat(strBuf, "text message");
    break;

  case PAYLOAD_TYPE_ACK:
    strcat(strBuf, "acknowledgement");
    break;

  case PAYLOAD_TYPE_ADVERT:
    strcat(strBuf, "advert");
    break;

  case PAYLOAD_TYPE_GRP_TXT:
    strcat(strBuf, "group text");
    break;

  case PAYLOAD_TYPE_GRP_DATA:
    strcat(strBuf, "group data");
    break;

  case PAYLOAD_TYPE_ANON_REQ:
    strcat(strBuf, "anon request");
    break;

  case PAYLOAD_TYPE_PATH:
    strcat(strBuf, "path");
    break;

  case PAYLOAD_TYPE_TRACE:
    strcat(strBuf, "trace");
    break;

  case PAYLOAD_TYPE_MULTIPART:
    strcat(strBuf, "multipart");
    break;

  case PAYLOAD_TYPE_RAW_CUSTOM:
    strcat(strBuf, "raw");
    break;
  }
  char version[2];
  version[0] = (frame.header >> 6) + '0';
  version[1] = 0;

  strcat(strBuf, ", payload version is ");

  strcat(strBuf, version);

  puts(strBuf);

  if ((frame.header & ROUTE_TYPE_MASK) == ROUTE_TYPE_TRANSPORT_DIRECT ||
      (frame.header & ROUTE_TYPE_MASK) == ROUTE_TYPE_TRANSPORT_FLOOD) {
    printf("Transport codes: %d %d\n", *((uint16_t *)frame.transportCodes),
           *((uint16_t *)&(frame.transportCodes[2])));
  }
  printf("Path is %d nodes long", frame.pathLen);

  for (uint8_t pathIndex = 0; pathIndex < frame.pathLen; pathIndex++) {
    printf("node %d - %02X, ", pathIndex, frame.path[pathIndex]);
  }
  putchar('\n');

  char payloadBuf[185];
  memset(payloadBuf, 0, sizeof(payloadBuf));
  memcpy(payloadBuf, frame.payload, frame.payloadLen);

  //   unsigned char outBuf[345];
  //   memset(outBuf, 0, sizeof(outBuf));
  //   size_t output_len = 0; // will store the real output size
  //   mbedtls_base64_encode(outBuf, sizeof(outBuf), &output_len, frame.payload,
  //   frame.payloadLen); printf("Payload L%d: %s\n", output_len, outBuf);

  // printf("\npayload is %s \n", payloadBuf);
}

void app_main(void) {

  uart_set_baudrate(UART_NUM_0, 115200);
  ESP_LOGI("BOOT", "BRN Systems incorporated MeshCore firmware build: %s %s",
           __DATE__, __TIME__);

  spi_bus_config_t HighSpeedBusCfg = {
      // CONNECTED to LoRa and SD card
      .mosi_io_num = HSPI_MOSI_GPIO, .miso_io_num = HSPI_MISO_GPIO,
      .sclk_io_num = HSPI_SCK_GPIO,  .quadwp_io_num = -1,
      .quadhd_io_num = -1,
      .max_transfer_sz = 256, // probably change
  };
  ESP_ERROR_CHECK(
      spi_bus_initialize(SPI3_HOST, &HighSpeedBusCfg, SPI_DMA_CH_AUTO));

  const int64_t interval_us = 10000; // 10 ms
  int64_t start_time, end_time, elapsed;

  LoRaInit();
  int8_t txPowerInDbm = 20;
  uint32_t frequencyInHz = 869525000;

  ESP_LOGW(TAG, "Enable TCXO");
  // float tcxoVoltage = 2.2; //ebyte
  float tcxoVoltage = 1.8; // heltec
  bool useRegulatorLDO = true;

  LoRaDebugPrint(false);
  if (LoRaBegin(frequencyInHz, txPowerInDbm, tcxoVoltage, useRegulatorLDO) !=
      0) {
    ESP_LOGE(TAG, "Does not recognize the module");
    while (1) {
      vTaskDelay(1);
    }
  }

  uint8_t spreadingFactor = 11;
  uint8_t bandwidth = SX126X_LORA_BW_250_0;
  uint8_t codingRate = SX126X_LORA_CR_4_5;
  uint16_t preambleLength = 16;
  bool crcOn = true;
  bool invertIrq = false;

  LoRaConfig(spreadingFactor, bandwidth, codingRate, preambleLength, 0, crcOn,
             invertIrq);

  uint8_t bufIn[256];

  while (1) {
    start_time = esp_timer_get_time();

    uint8_t rxLen = LoRaReceive(bufIn, sizeof(bufIn));
    if (rxLen > 0) {
      // ESP_LOGI(TAG, "%d byte packet received", rxLen);

      int8_t rssi, snr;
      GetPacketStatus(&rssi, &snr);
      // ESP_LOGI(TAG, "rssi=%d[dBm] snr=%d[dB]", rssi, snr);

      FrameStruct frame = decodeFrame(bufIn, rxLen);

      printFrameHeader(frame);

      switch (frame.header & PAYLOAD_TYPE_MASK) {
      case PAYLOAD_TYPE_REQ:
        break;

      case PAYLOAD_TYPE_RESPONSE:
        break;

      case PAYLOAD_TYPE_TXT_MSG:
        break;

      case PAYLOAD_TYPE_ACK:
        unsigned char checkSumBuf[10];
        memset(checkSumBuf, 0, sizeof(checkSumBuf));
        size_t checksumLen = 0;
        mbedtls_base64_encode(checkSumBuf, sizeof(checkSumBuf), &checksumLen,
                              frame.payload, 4);
        printf("Checksum: %s\n", checkSumBuf);
        break;

      case PAYLOAD_TYPE_ADVERT:
        AdvertisementPayload advert = decodeAdvertisement(frame);
        printAdvertisement(advert);
        break;

      case PAYLOAD_TYPE_GRP_TXT:
        GroupTextMessage msg = decodeGroupMessage(frame);
        printGroupMessage(msg);
        break;

      case PAYLOAD_TYPE_GRP_DATA:
        break;

      case PAYLOAD_TYPE_ANON_REQ:
        break;

      case PAYLOAD_TYPE_PATH:
        break;

      case PAYLOAD_TYPE_TRACE:
        break;

      case PAYLOAD_TYPE_MULTIPART:
        break;

      case PAYLOAD_TYPE_RAW_CUSTOM:
        break;
      }
    }

    int lost = GetPacketLost();
    if (lost != 0) {
      ESP_LOGW(TAG, "%d packets lost", lost);
    }

    end_time = esp_timer_get_time();
    elapsed = end_time - start_time;

    if (elapsed < interval_us) {
      vTaskDelay(pdMS_TO_TICKS((interval_us - elapsed) / 1000));
    }
  }
}