#include "lib/config.h"
#include "meshcore/meshframing.h"
#include "meshcore/packets/ack.h"
#include "meshcore/packetstructs.h"
#include "meshcore/stats.h"
#include "util/hexdump.h"
#include "util/log.h"
#include <stdio.h>
#include <string.h>
#include "encrypted.h"
#include "FreeRTOS.h"
#include "task.h"


#define TAG "EncryptedMessage"

void sendEncryptedFrame (NodeEntry *targetNode, uint8_t payloadType, const uint8_t *plain, size_t plainLen) {
    FrameStruct frame;
    uint8_t offset = 0;

    // 1. Header
    frame.header =
        ROUTE_TYPE_FLOOD | //currently flood
        payloadType |
        PAYLOAD_VERSION_0;

    // 2. Destination + source
    frame.payload[offset++] = targetNode->pubKey[0];
    frame.payload[offset++] = persistent.pubkey[0];

    // 4. Encrypt + MAC
    size_t encLen;
    encrypt_then_mac (
        targetNode->secret,
        32,
        plain,
        plainLen,
        frame.payload + offset,
        &encLen);

    offset += encLen;

    // 5. Finalize
    frame.payloadLen = offset;
    memcpy (&frame.path, &targetNode->path, sizeof (frame.path));

    hexdump ("Encrypted frame", frame.payload, frame.payloadLen);
    sendFrame (frame);
}

void sendEncryptedTextMessage (NodeEntry *targetNode, const PlainTextMessagePayload *msg) {
    if (targetNode == NULL) {
        ESP_LOGW(TAG, "Node is null");
        return;
    }
    if (targetNode->last_seen_lt == 0) {
        ESP_LOGW(TAG, "Node is not populated");
        return;
    }
    uint8_t buf[256];
    uint8_t index = 0;

    uint8_t msgLen = strlen(msg->message) + 1;
    buf[index++] = msg->timestamp;
    buf[index++] = msg->timestamp >> 8;
    buf[index++] = msg->timestamp >> 16;
    buf[index++] = msg->timestamp >> 24;

    buf[index++] = (msg->textType << 2) | (msg->attempt & 0x03);
    memcpy (&buf[index], msg->message, msgLen);

    sendEncryptedFrame (
        targetNode,
        PAYLOAD_TYPE_TXT_MSG,
        buf,
        index + msgLen);
}

void sendEncryptedResponse (NodeEntry *targetNode, const Response *resp) {
    uint8_t buf[256];
    uint8_t index = 0;

    buf[index++] = resp->tag;
    buf[index++] = resp->tag >> 8;
    buf[index++] = resp->tag >> 16;
    buf[index++] = resp->tag >> 24;

    memcpy (&buf[index], resp->data, resp->dataLen);

    sendEncryptedFrame (
        targetNode,
        PAYLOAD_TYPE_RESPONSE,
        buf,
        index);
}

void sendEncryptedRequest (NodeEntry *targetNode, const Request *req) {
    uint8_t buf[256];
    uint8_t index = 0;

    buf[index++] = req->timestamp;
    buf[index++] = req->timestamp >> 8;
    buf[index++] = req->timestamp >> 16;
    buf[index++] = req->timestamp >> 24;

    buf[index++] = req->requestType;
    memcpy (&buf[index], req->data, req->dataLen);

    sendEncryptedFrame (
        targetNode,
        PAYLOAD_TYPE_REQ,
        buf,
        index);
}

void sendEncryptedPathPayload (NodeEntry *targetNode, const ReturnedPathPayload *path) {
    uint8_t buf[256];
    uint8_t index = 0;

    buf[index++] = path->path.pathLen;
    memcpy (&buf[index], path->path.path, path->path.pathLen);
    index += path->path.pathLen;

    buf[index++] = path->extra.type;
    memcpy (&buf[index], path->extra.data,
            sizeof (path->extra.data));

    sendEncryptedFrame (
        targetNode,
        PAYLOAD_TYPE_PATH,
        buf,
        index);
}

void printRequest (const Request *req) {
    printf ("Request:\n");
    printf ("  Timestamp: %u\n", req->timestamp);
    printf ("  Type: 0x%02X\n", req->requestType);
    printf ("  Data: ");
    for (int i = 0; i < req->dataLen; i++) {
        printf ("%02X ", req->data[i]);
    }
    printf ("\n");
}

void printResponse (const Response *resp) {
    printf ("Response:\n");
    printf ("  Tag: %u\n", resp->tag);
    printf ("  Data: ");
    for (int i = 0; i < resp->dataLen; i++) {
        printf ("%02X ", resp->data[i]);
    }
    printf ("\n");
}

void printPlainTextMessage (const PlainTextMessagePayload *msg) {
    printf ("PlainTextMessage:\n");
    printf ("  Timestamp: %u\n", msg->timestamp);
    printf ("  Attempt: %u\n", msg->attempt);
    printf ("  TextType: %u\n", msg->textType);
    printf ("  Message: %.*s\n", (int)strlen (msg->message), msg->message);
}

void printReturnedPathPayload (const ReturnedPathPayload *path) {
    printf ("ReturnedPathPayload:\n");
    printf ("  Path Length: %u\n", path->path.pathLen);
    printf ("  Path: ");
    for (int i = 0; i < path->path.pathLen; i++) {
        printf ("%02X ", path->path.path[i]);
    }
    printf ("\n");
    printf ("  Extra Type: %u\n", path->extra.type);
    printf ("  Extra Data: ");
    for (int i = 0; i < sizeof (path->extra.data); i++) {
        printf ("%02X ", path->extra.data[i]);
    }
    printf ("\n");
}

void printEncryptedPayload (const EncryptedPayloadStruct *enc) {
    printf ("EncryptedPayload:\n");
    printf ("  Type: 0x%02X\n", enc->type);
    printf ("  DestinationHash: 0x%02X\n", enc->destinationHash);
    printf ("  SourceHash: 0x%02X\n", enc->sourceHash);
    printf ("  CipherMAC: 0x%04X\n", enc->cipherMAC);
    printf ("  PayloadLen: %zu\n", enc->payloadLen);
    printf ("  Payload: ");
    for (size_t i = 0; i < enc->payloadLen; i++) {
        printf ("%02X ", enc->payload[i]);
    }
    printf ("\n");
}

EncryptedPayloadStruct decodeEncryptedPayload (FrameStruct frame) {
    EncryptedPayloadStruct enc;
    memset (&enc, 0, sizeof (enc));
    enc.type = frame.header & PAYLOAD_TYPE_MASK;
    unsigned char index = 0;

    enc.destinationHash = frame.payload[index++];
    enc.sourceHash = frame.payload[index++];
    enc.cipherMAC = frame.payload[index];
    enc.cipherMAC |= frame.payload[index + 1] << 8;

    if (enc.destinationHash != persistent.pubkey[0]) {
        return enc;
    }

    ESP_LOGI(TAG, "Finding remote node, sourceHash is %d", enc.sourceHash);

    NodeEntry *remNode = getNode (enc.sourceHash);

    enc.remNode = remNode;


    if (remNode == NULL) {
        ESP_LOGW(TAG, "Node not in DB");
        return enc;
    }
    ESP_LOGI(TAG, "Found node with index %d", remNode - persistent.contacts);


    if (mac_then_decrypt (remNode->secret, 32, &(frame.payload[index]), frame.payloadLen - index, enc.payload) != 0) {
        ESP_LOGW (TAG, "HMAC failed on encrypted message %s", remNode->name);
    } else {
        enc.payloadLen = frame.payloadLen - HMAC_SIZE;
        ESP_LOGI(TAG, "HMAC success from %s, %u bytes long", remNode->name, enc.payloadLen);
        sendDiscreteAck(enc.payload, 5 + strlen((char *)&enc.payload[5]), remNode->pubKey);
    }

    return enc;
}

void parseEncryptedPayload (EncryptedPayloadStruct enc) {
    // printEncryptedPayload(&enc);

    printf ("EncryptedPayload:\n");
    printf ("  Type: 0x%02X\n", enc.type);
    printf ("  DestinationHash: 0x%02X\n", enc.destinationHash);
    printf ("  SourceHash: 0x%02X\n", enc.sourceHash);
    printf ("  CipherMAC: 0x%04X\n", enc.cipherMAC);
    printf ("  PayloadLen: %u\n", enc.payloadLen);
    printf ("  Payload: ");
    hexdump("Full payload buffer", enc.payload, sizeof(enc.payload));
    printf ("\n");

    uint8_t index = 0;
    if (enc.type == PAYLOAD_TYPE_PATH) {

        ReturnedPathPayload retPath;
        retPath.path.pathLen = enc.payload[index++];
        if (retPath.path.pathLen > 64) {
            ESP_LOGW (TAG, "Path too long\n");
            return;
        }
        memcpy (retPath.path.path, &(enc.payload[index]), retPath.path.pathLen);
        index += retPath.path.pathLen;
        retPath.extra.type = enc.payload[index++];
        memcpy (retPath.extra.data, &(enc.payload[index]), enc.payloadLen - index);

    } else if (enc.type == PAYLOAD_TYPE_REQ) {
        Request req;
        req.timestamp = enc.payload[index++];
        req.timestamp |= enc.payload[index++] << 8;
        req.timestamp |= enc.payload[index++] << 16;
        req.timestamp |= enc.payload[index++] << 24;
        req.requestType = enc.payload[index++];
        req.dataLen = enc.payloadLen - index;
        memcpy (req.data, &(enc.payload[index]), req.dataLen);
        printRequest (&req);
        switch (req.requestType) {
            Response resp;
            resp.tag = RTC_GetCounter();
            memcpy(resp.data, &stats, sizeof(stats));
            resp.dataLen = sizeof(stats);
            sendEncryptedResponse(enc.remNode, &resp);
            break;
        case REQUEST_KEEPALIVE:
            break;
        case REQUEST_GET_TELEMETRY_DATA:
            break;
        case REQUEST_GET_MIN_MAX_AVG:
            break;
        case REQUEST_GET_ACCESS_LIST:
            break;
        }

    } else if (enc.type == PAYLOAD_TYPE_RESPONSE) {
        Response resp;
        resp.tag = enc.payload[index++];
        resp.tag |= enc.payload[index++] << 8;
        resp.tag |= enc.payload[index++] << 16;
        resp.tag |= enc.payload[index++] << 24;
        resp.dataLen = enc.payloadLen - index;
        memcpy (resp.data, &(enc.payload[index]), resp.dataLen);
        printResponse (&resp);

    } else if (enc.type == PAYLOAD_TYPE_TXT_MSG) {
        PlainTextMessagePayload plaintext;
        plaintext.timestamp = enc.payload[index++];
        plaintext.timestamp |= enc.payload[index++] << 8;
        plaintext.timestamp |= enc.payload[index++] << 16;
        plaintext.timestamp |= enc.payload[index++] << 24;
        plaintext.attempt = enc.payload[index] & 0x03;
        plaintext.textType = enc.payload[index++] >> 2;
        memcpy (plaintext.message, &(enc.payload[index]), enc.payloadLen - index);
        printPlainTextMessage (&plaintext);
    }
}