#include "meshframing.h"
#include "ch32v30x_gpio.h"
#include "lib/config.h"
#include "meshcore/stats.h"
#include "string.h"
#include "stdio.h"
#include "sx1262.h"
#include "util/hexdump.h"
#include "util/log.h"
#include "FreeRTOS.h"
#include "task.h"
#include "lib/cifra/sha2.h"

#define TAG "Meshframing"

int ReadFrame (FrameStruct *frame, int8_t *rssiPacket, int8_t *snrPacket, int8_t *rawSnr) {
    uint16_t irqRegs = GetIrqStatus();
    // uint8_t status = GetStatus();

    if (irqRegs & SX126X_IRQ_RX_DONE) {
        // ClearIrqStatus(SX126X_IRQ_RX_DONE);
        ClearIrqStatus (SX126X_IRQ_ALL);

        uint8_t offset = 0;
        uint8_t payloadLength = 0;
        GetRxBufferStatus (&payloadLength, &offset);

        if (payloadLength == 0) {
            return 0;
        }

        GetPacketStatus (rssiPacket, snrPacket, rawSnr);

        memset (frame, 0, sizeof (FrameStruct));
        WaitForIdle (BUSY_WAIT, "start ReadBuffer", 1);

        uint8_t cmd[3] = {SX126X_CMD_READ_BUFFER, offset, SX126X_CMD_NOP};

        GPIO_WriteBit (GPIOA, GPIO_Pin_4, 0);

        uint16_t curPayloadIndex = 0;
        uint8_t pathIndex = 0;
        uint8_t transportIndex = 0;
        uint8_t state = 0;  // 0=header, 1=transport, 2=pathLen, 3=path, 4=payload

        for (uint16_t i = 0; i < payloadLength + sizeof (cmd); i++) {
            uint8_t out = (i < sizeof (cmd)) ? cmd[i] : 0xFF;

            // Wait TX ready
            while (!SPI_I2S_GetFlagStatus (SPI1, SPI_I2S_FLAG_TXE));
            SPI_I2S_SendData (SPI1, out);

            // Wait RX ready
            while (!SPI_I2S_GetFlagStatus (SPI1, SPI_I2S_FLAG_RXNE));
            uint8_t in = (uint8_t)SPI_I2S_ReceiveData (SPI1);

            // Only process payload bytes
            if (i >= sizeof (cmd)) {
                switch (state) {
                case 0:  // header
                    frame->header = in;
                    state = ((frame->header & ROUTE_TYPE_MASK) == ROUTE_TYPE_TRANSPORT_DIRECT ||
                             (frame->header & ROUTE_TYPE_MASK) == ROUTE_TYPE_TRANSPORT_FLOOD)
                                ? 1
                                : 2;
                    break;

                case 1:  // transportCodes[4]
                    frame->transportCodes[transportIndex++] = in;
                    if (transportIndex >= 4)
                        state = 2;
                    break;

                case 2:  // pathLen
                    frame->path.pathLen = in;
                    if (frame->path.pathLen > 64) {
                        frame->path.pathLen = 64;
                    }
                    pathIndex = 0;
                    state = (frame->path.pathLen > 0) ? 3 : 4;
                    break;

                case 3:  // path
                    frame->path.path[pathIndex++] = in;
                    if (pathIndex >= frame->path.pathLen)
                        state = 4;
                    break;

                case 4:  // payload
                    frame->payload[curPayloadIndex++] = in;
                    break;
                }
            }
        }

        frame->payloadLen = curPayloadIndex;

        GPIO_WriteBit (GPIOA, GPIO_Pin_4, 1);

        WaitForIdle (BUSY_WAIT, "end ReadBuffer", 0);
        return payloadLength;
    }
    return 0;
}

/*
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.path.pathLen = data[index++];

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

    return frame;
}
*/
void printFrameHeader (const 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_CONTROL:
        printf ("control");
        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->path.pathLen);

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

void LoRaTransmit (const FrameStruct *frame) {
    uint8_t len = 2;  // header + path_len

    if ((frame->header & ROUTE_TYPE_MASK) == ROUTE_TYPE_TRANSPORT_DIRECT ||
        (frame->header & ROUTE_TYPE_MASK) == ROUTE_TYPE_TRANSPORT_FLOOD) {
        len += 4;
    }

    len += frame->path.pathLen;
    len += frame->payloadLen;

    uint16_t irqStatus;
    char rv = 0;

    if (txActive == 0) {
        txActive = 1;

        if (PacketParams[2] == 0x00) {  // explicit header, variable length
            PacketParams[3] = len;
        }

        WriteCommand (SX126X_CMD_SET_PACKET_PARAMS, PacketParams, 6);
        ClearIrqStatus (SX126X_IRQ_ALL);

        WaitForIdle (BUSY_WAIT, "start WriteBuffer", 1);

        uint8_t cmdBuf[2] = {SX126X_CMD_WRITE_BUFFER, 0x00};
        uint8_t cmdLen = sizeof (cmdBuf);

        GPIO_WriteBit (GPIOA, GPIO_Pin_4, 0);

        uint16_t payloadIndex = 0;
        uint8_t state = 0;
        uint8_t pathIndex = 0;
        uint8_t transportIndex = 0;

        for (uint16_t i = 0; i < len + cmdLen; i++) {
            uint8_t out = 0xFF;

            if (i < cmdLen) {
                out = cmdBuf[i];
            } else {
                switch (state) {
                case 0:  // header
                    out = frame->header;
                    state = ((frame->header & ROUTE_TYPE_MASK) == ROUTE_TYPE_TRANSPORT_DIRECT ||
                             (frame->header & ROUTE_TYPE_MASK) == ROUTE_TYPE_TRANSPORT_FLOOD)
                                ? 1
                                : 2;
                    break;

                case 1:  // transport codes
                    out = frame->transportCodes[transportIndex++];
                    if (transportIndex >= 4)
                        state = 2;
                    break;

                case 2:  // path length
                    out = frame->path.pathLen;
                    pathIndex = 0;
                    state = (frame->path.pathLen > 0) ? 3 : 4;
                    break;

                case 3:  // path
                    out = frame->path.path[pathIndex++];
                    if (pathIndex >= frame->path.pathLen)
                        state = 4;
                    break;

                case 4:  // payload
                    out = frame->payload[payloadIndex++];
                    break;
                }
            }

            while (SPI_I2S_GetFlagStatus (SPI1, SPI_I2S_FLAG_TXE) == RESET);
            SPI_I2S_SendData (SPI1, out);

            while (SPI_I2S_GetFlagStatus (SPI1, SPI_I2S_FLAG_RXNE) == RESET);
            (void)SPI_I2S_ReceiveData (SPI1);  // discard
        }

        GPIO_WriteBit (GPIOA, GPIO_Pin_4, 1);

        WaitForIdle (BUSY_WAIT, "end WriteBuffer", 0);

        SetTx (5000);

        irqStatus = GetIrqStatus();
        while (!(irqStatus & (SX126X_IRQ_TX_DONE | SX126X_IRQ_TIMEOUT))) {
            vTaskDelay (1);
            irqStatus = GetIrqStatus();
        }

        txActive = 0;
        SetRx (0xFFFFFF);

        if (irqStatus & SX126X_IRQ_TX_DONE) {
            rv = 1;
        }
    }

    if (rv == 0) {
        txLost++;
    }
}

/*
void sendFrame (const 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->header & ROUTE_TYPE_MASK) == ROUTE_TYPE_FLOOD ||
        (frame->header & ROUTE_TYPE_MASK) == ROUTE_TYPE_TRANSPORT_FLOOD) {
        stats.sentFloodCount++;
    }

    if ((frame->header & ROUTE_TYPE_MASK) == ROUTE_TYPE_DIRECT ||
        (frame->header & ROUTE_TYPE_MASK) == ROUTE_TYPE_TRANSPORT_DIRECT) {
        stats.sentDirectCount++;
    }

    stats.packetsSentCount++;

    uint8_t pathLen = frame->path.pathLen;


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

    txBuf[offset++] = pathLen;

    memcpy (txBuf + offset, frame->path.path, pathLen);
    hexdump ("TxDump Path", frame->path.path, frame->path.pathLen);
    offset += 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);
    hexdump ("TxDumpPayload", frame->payload, frame->payloadLen);
    TickType_t start, end;
    start = xTaskGetTickCount();
    LoRaSend (txBuf, offset, SX126x_TXMODE_SYNC);
    end = xTaskGetTickCount();
    tickAirtime += end - start;
}
*/

void retransmitFrame (FrameStruct *frame) {

    /* -------- FLOOD -------- */
    if (frame->header & ROUTE_TYPE_FLOOD ||
        frame->header & ROUTE_TYPE_TRANSPORT_FLOOD) {

        if (frame->header != DONT_RETRANSMIT_HEADER &&
            frame->path.pathLen + 1 < MAX_FLOOD_TTL) {

            // append self to END of path
            frame->path.path[frame->path.pathLen++] = persistent.pubkey[0];
            LoRaTransmit (frame);
        }
    }

    /* -------- DIRECT -------- */
    if (frame->header & ROUTE_TYPE_DIRECT ||
        frame->header & ROUTE_TYPE_TRANSPORT_DIRECT) {

        // are we the next hop?
        if (frame->path.pathLen > 0 &&
            frame->path.path[0] == persistent.pubkey[0]) {

            // remove self from START of path
            frame->path.pathLen--;
            memmove (frame->path.path,
                     frame->path.path + 1,
                     frame->path.pathLen);

            // forward only if there's another hop
            if (frame->path.pathLen > 0) {
                LoRaTransmit (frame);
            }
        }
    }
}

// Verify MAC + Decrypt

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

    size_t padded_len = ((plen + 15) / 16) * 16;

    // prepare padded buffer
    uint8_t padded[padded_len];
    memset (padded, 0, padded_len);    // zero padding
    memcpy (padded, plaintext, plen);  // copy plaintext

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

    // encrypt plaintext
    aes_encrypt_ecb (aes_key, 16, padded, padded_len, ciphertext);

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

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

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

int mac_then_decrypt (const uint8_t *aes_key, const uint8_t keySize, 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;

    if (clen % 16 != 0)
        return -2;         // must be multiple of block size

    uint8_t calc_mac[32];  // full SHA-256
    hmac_sha256 (aes_key, keySize, ciphertext, clen, calc_mac);

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

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

uint16_t getTransportCode (const FrameStruct *frame) {
    uint16_t code;
    /*
    // compute HMAC over ciphertext
    uint8_t mac[32];  // full SHA-256
    cf_hmac_ctx ctx;
    cf_hmac_init (&ctx, &cf_sha256, key, sizeof(key));
    cf_hmac_update (&ctx, &(frame->header), sizeof (frame->header));
    cf_hmac_update (&ctx, frame->payload, frame->payloadLen);
    cf_hmac_finish (&ctx, mac);

    // copy only HMAC_SIZE bytes of MAC
    memcpy (&code, mac, HMAC_SIZE);
    */
    if (code == 0) {  // reserve codes 0000 and FFFF
        code++;
    } else if (code == 0xFFFF) {
        code--;
    }
    return code;
}