#include "meshcore.h"
#include "FreeRTOS.h"
#include "lib/monocypher/monocypher-ed25519.h"
#include "task.h"
#include "lib/base64.h"
#include "lib/cifra/aes.h"
#include "lib/cifra/sha2.h"
#include "lib/cifra/hmac.h"
#include "lib/config.h"

#define AESKeyCount 8

const uint8_t aesKeys[AESKeyCount][17] = {
    {0x11, 0x8b, 0x33, 0x87, 0xe9, 0xc5,
     0xcd, 0xea, 0x6a, 0xc9, 0xe5, 0xed,
     0xba, 0xa1, 0x15, 0xcd, 0x72},

    {0x0a, 0x44, 0x81, 0xda, 0x0e, 0x4e,
     0x03, 0xc4, 0x9e, 0x84, 0x77, 0x25,
     0xd8, 0x3a, 0x93, 0xbf, 0x80}
};

#define TAG "MeshCore"

// requires at least a 256 byte data
FrameStruct decodeFrame (unsigned char *data, unsigned char dataLen) {
    hexdump ("RxDump", data, 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;
}

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

    cf_aes_context ctx;
    cf_aes_init (&ctx, key, 16);

    for (size_t i = 0; i < ilen; i += 16) {
        cf_aes_encrypt (&ctx, input + i, output + i);
    }

    cf_aes_finish (&ctx);
    return 0;
}

int aes_decrypt_ecb (const uint8_t *key, const uint8_t *input, size_t ilen,
                     uint8_t *output) {
    if (ilen % 16 != 0)
        return -1;

    cf_aes_context ctx;
    cf_aes_init (&ctx, key, 16);

    for (size_t i = 0; i < ilen; i += 16) {
        cf_aes_decrypt (&ctx, input + i, output + i);
    }


    cf_aes_finish (&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) {
    cf_hmac_ctx ctx;
    cf_hmac_init (&ctx, &cf_sha256, key, keylen);
    cf_hmac_update (&ctx, input, ilen);
    cf_hmac_finish (&ctx, output);
    return 0;
}

// Verify MAC + Decrypt
#define KEY_SIZE 16  // AES-128

int encrypt_then_mac (const uint8_t *aes_key,
                      const uint8_t *plaintext, size_t plen,
                      uint8_t *output, size_t *olen) {
    if (plen == 0)
        return -1;

    // ciphertext will go right after HMAC
    uint8_t *ciphertext = output + HMAC_SIZE;

    // encrypt plaintext
    aes_encrypt_ecb (aes_key, plaintext, plen, ciphertext);

    // compute HMAC over ciphertext
    uint8_t mac[32];  // full SHA-256
    hmac_sha256 (aes_key, KEY_SIZE, ciphertext, plen, mac);

    // copy only HMAC_SIZE bytes of MAC
    memcpy (output, mac, HMAC_SIZE);

    // return total length = HMAC + ciphertext
    *olen = HMAC_SIZE + plen;
    return 0;
}

int mac_then_decrypt (const uint8_t *aes_key, const uint8_t *input, size_t ilen,
                      uint8_t *plaintext) {
    if (ilen <= HMAC_SIZE)
        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 SHA-256
    hmac_sha256 (aes_key, KEY_SIZE, ciphertext, clen, calc_mac);

    if (memcmp (mac, calc_mac, HMAC_SIZE) != 0)
        return -2;

    return aes_decrypt_ecb (aes_key, ciphertext, clen, plaintext);
}

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");
    }
}

// EncryptedPayloadStruct decodeEncryptedPayload(FrameStruct frame) {
//   EncryptedPayloadStruct enc;
//   memset(&enc, 0, sizeof(enc));
//   if ((frame.header & PAYLOAD_TYPE_MASK) != PAYLOAD_TYPE_GRP_TXT) {
//     return enc;
//   }
//   unsigned char index = 0;
//   enc.destinationHash = frame.payload[index++];
//   enc.sourceHash = frame.payload[index++];


//  index = 0;

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

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

GroupTextMessage decodeGroupMessage (FrameStruct frame) {
    GroupTextMessage msg;
    memset (&msg, 0, sizeof (msg));
    if ((frame.header & PAYLOAD_TYPE_MASK) != PAYLOAD_TYPE_GRP_TXT) {
        ESP_LOGW (TAG, "Not a group text");
        return msg;
    }
    unsigned char index = 0;
    msg.channelHash = frame.payload[index++];
    unsigned char tmp[184];


    unsigned char decrypted = 0;
    for (unsigned char i = 0; i < AESKeyCount; i++) {

        if (msg.channelHash != aesKeys[i][0]) {
            ESP_LOGW (TAG, "Hash %d does not equal %d", aesKeys[i][0], msg.channelHash);
            continue;
        }

        ESP_LOGW (TAG, "Hash does equal %d", msg.channelHash);

        if (mac_then_decrypt (aesKeys[i] + 1, frame.payload + index, frame.payloadLen - index, tmp) != 0) {
            ESP_LOGW (TAG, "HMAC failed on grouphash key %d not matching %d", aesKeys[i][0], msg.channelHash);
            continue;
        }
        hexdump ("RxDumpDec", tmp, frame.payloadLen - index);
        decrypted = 1;
        break;
    }


    if (!decrypted) {
        return msg;
    }

    unsigned char plaintextLen = frame.payloadLen - index;
    index = 0;

    ESP_LOGI (TAG, "Starting memcpy");
    vTaskDelay (pdMS_TO_TICKS (10));
    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) {
    unsigned char keyBuf[50];
    unsigned char sigBuf[90];
    memset (keyBuf, 0, sizeof (keyBuf));
    memset (sigBuf, 0, sizeof (sigBuf));
    base64_encode (advert.pubKey, 32, keyBuf);
    base64_encode (advert.signature, 64, sigBuf);

    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) {
    switch (frame.header & ROUTE_TYPE_MASK) {
    case ROUTE_TYPE_TRANSPORT_FLOOD:
        printf ("transport flood");
        break;

    case ROUTE_TYPE_FLOOD:
        printf ("flood");
        break;

    case ROUTE_TYPE_DIRECT:
        printf ("direct");
        break;

    case ROUTE_TYPE_TRANSPORT_DIRECT:
        printf ("transport direct");
        break;
    }

    printf (", payload type is ");

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

    case PAYLOAD_TYPE_RESPONSE:
        printf ("response");
        break;

    case PAYLOAD_TYPE_TXT_MSG:
        printf ("text message");
        break;

    case PAYLOAD_TYPE_ACK:
        printf ("acknowledgement");
        break;

    case PAYLOAD_TYPE_ADVERT:
        printf ("advert");
        break;

    case PAYLOAD_TYPE_GRP_TXT:
        printf ("group text");
        break;

    case PAYLOAD_TYPE_GRP_DATA:
        printf ("group data");
        break;

    case PAYLOAD_TYPE_ANON_REQ:
        printf ("anon request");
        break;

    case PAYLOAD_TYPE_PATH:
        printf ("path");
        break;

    case PAYLOAD_TYPE_TRACE:
        printf ("trace");
        break;

    case PAYLOAD_TYPE_MULTIPART:
        printf ("multipart");
        break;

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

    printf (", payload version is %s ", version);

    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');
}

void sendFrame (FrameStruct frame) {
    uint8_t txBuf[256];
    size_t offset = 0;

    txBuf[offset++] = frame.header;

    if ((frame.header & ROUTE_TYPE_MASK) == ROUTE_TYPE_TRANSPORT_DIRECT ||
        (frame.header & ROUTE_TYPE_MASK) == ROUTE_TYPE_TRANSPORT_FLOOD) {
        memcpy (txBuf + offset, frame.transportCodes, 4);
        offset += 4;
    }

    if (frame.pathLen > 64) {
        frame.pathLen = 64;
    }

    txBuf[offset++] = frame.pathLen;

    memcpy (txBuf + offset, frame.path, frame.pathLen);
    offset += frame.pathLen;

    uint16_t maxPayloadLen = 256 - offset;

    uint16_t payloadLen = frame.payloadLen > maxPayloadLen ? maxPayloadLen : frame.payloadLen;

    memcpy (txBuf + offset, frame.payload, payloadLen);
    offset += payloadLen;

    hexdump ("TxDump", txBuf, offset);
    LoRaSend (txBuf, offset, SX126x_TXMODE_SYNC);
}

void sendGroupMessage (GroupTextMessage msg) {

    msg.channelHash = aesKeys[msg.keyIndex][0];

    msg.flags = 0;
    msg.timestamp = RTC_GetCounter();

    FrameStruct frame;
    frame.header = ROUTE_TYPE_FLOOD | PAYLOAD_TYPE_GRP_TXT | PAYLOAD_VERSION_0;
    frame.pathLen = 0;
    size_t offset = 0;
    memset (frame.payload, 0, sizeof (frame.payload));
    frame.payload[offset++] = msg.channelHash;

    uint8_t cipherBuf[176];
    size_t offset2 = 0;
    memcpy (cipherBuf, (const void *)&(msg.timestamp), 4);
    offset2 += 4;
    cipherBuf[offset2++] = msg.flags;
    size_t textSize = offset2 + strlen ((const char *)msg.text);
    if (textSize > 175) {
        textSize = 175;
    }
    memcpy (cipherBuf + offset2, msg.text, textSize);
    offset2 += textSize;

    offset2 = (((offset2 / 16) + 1) * 16);

    size_t olen = 0;
    hexdump ("TxDumpDec", cipherBuf, offset2);
    encrypt_then_mac (&(aesKeys[msg.keyIndex][1]), cipherBuf, offset2, &(frame.payload[offset]), &olen);

    frame.payloadLen = olen + 1;

    sendFrame (frame);
    return;
}

void makeSendGroupMessage (char *txt, uint8_t keyIndex) {
    GroupTextMessage msg;
    strcpy ((char *)msg.text, txt);
    msg.keyIndex = keyIndex;
    sendGroupMessage (msg);
    return;
}

void sendAdvert() {
    AdvertisementPayload ad;
    memcpy (ad.pubKey, persistent.pubkey, sizeof (ad.pubKey));
    ad.dataFlags = ADVERTISEMENT_FLAG_IS_CHAT_NODE | ADVERTISEMENT_FLAG_HAS_NAME;
    strcpy (ad.nodeName, persistent.nodeName);
    ad.timestamp = RTC_GetCounter();

    // 1. Build app_data exactly like AdvertDataBuilder::encodeTo
    uint8_t app_data[40];
    size_t app_len = 0;

    app_data[app_len++] = ad.dataFlags;

    if (ad.dataFlags & ADVERTISEMENT_FLAG_HAS_LOCATION) {
        memcpy (app_data + app_len, &ad.latitude, sizeof (ad.latitude));
        app_len += sizeof (ad.latitude);
        memcpy (app_data + app_len, &ad.longitude, sizeof (ad.longitude));
        app_len += sizeof (ad.longitude);
    }

    if (ad.dataFlags & ADVERTISEMENT_FLAG_RFU1) {
        memcpy (app_data + app_len, &ad.rfu1, sizeof (ad.rfu1));
        app_len += sizeof (ad.rfu1);
    }

    if (ad.dataFlags & ADVERTISEMENT_FLAG_RFU2) {
        memcpy (app_data + app_len, &ad.rfu2, sizeof (ad.rfu2));
        app_len += sizeof (ad.rfu2);
    }

    if (ad.dataFlags & ADVERTISEMENT_FLAG_HAS_NAME) {
        size_t nodenameLen = strlen (ad.nodeName);
        memcpy (app_data + app_len, ad.nodeName, nodenameLen);
        app_len += nodenameLen;
    }

    // 2. Reserve frame and build payload header
    FrameStruct frame;
    frame.header = ROUTE_TYPE_FLOOD | PAYLOAD_TYPE_ADVERT | PAYLOAD_VERSION_0;

    size_t offset = 0;
    memcpy (frame.payload + offset, ad.pubKey, sizeof (ad.pubKey));
    offset += sizeof (ad.pubKey);

    memcpy (frame.payload + offset, &ad.timestamp, sizeof (ad.timestamp));
    offset += sizeof (ad.timestamp);

    // reserve signature space
    uint8_t *signature_pos = frame.payload + offset;
    offset += 64;

    // append app_data after signature
    memcpy (frame.payload + offset, app_data, app_len);
    offset += app_len;

    // 3. Sign pubKey + timestamp + app_data
    uint8_t message[76];
    size_t msg_len = 0;
    memcpy (message + msg_len, ad.pubKey, sizeof (ad.pubKey));
    msg_len += sizeof (ad.pubKey);
    memcpy (message + msg_len, &ad.timestamp, sizeof (ad.timestamp));
    msg_len += sizeof (ad.timestamp);
    memcpy (message + msg_len, app_data, app_len);
    msg_len += app_len;

    crypto_ed25519_meshcore_sign (signature_pos, persistent.privkey, persistent.pubkey, message, msg_len);

    hexdump ("Complete advert", frame.payload, offset);

    hexdump ("Public key", ad.pubKey, 32);

    hexdump ("Signature", signature_pos, 64);

    hexdump ("Timestamp", &ad.timestamp, 4);

    hexdump ("NodeName", ad.nodeName, 20);

    hexdump ("Appdata", app_data, app_len);

    // 5. Set payload length and send
    frame.payloadLen = offset;
    frame.pathLen = 0;
    sendFrame (frame);
}