// Secure serial link — ties rs232, packet, and security into one API
//
// Handshake protocol:
//   1. Both sides send their DH public key (128 bytes) as a single packet
//   2. On receiving the remote key, compute shared secret immediately
//   3. Derive separate TX/RX cipher keys based on public key ordering
//   4. Transition to READY — all subsequent encrypted packets use the keys
//
// The handshake uses the packet layer's reliable delivery, so lost packets
// are automatically retransmitted. Both sides can send their public key
// simultaneously — there's no initiator/responder distinction.
//
// Directionality: the side with the lexicographically lower public key
// uses master XOR 0xAA for TX and master XOR 0x55 for RX. The other
// side uses the reverse. This is critical: if both sides used the same
// key and counter for CTR mode, they'd produce identical keystreams,
// and XOR'ing two ciphertexts would reveal the XOR of the plaintexts.
// The XOR-derived directional keys ensure each direction has a unique
// keystream even though both start their counters at zero.
//
// Channel header byte: bit 7 = encrypted, bits 6..0 = channel (0-127)

#include <stdint.h>
#include <stdbool.h>
#include <stdlib.h>
#include <string.h>
#include "secLink.h"
#include "../rs232/rs232.h"
#include "../packet/packet.h"
#include "../security/security.h"


// ========================================================================
// Internal defines
// ========================================================================

#define STATE_INIT          0
#define STATE_HANDSHAKE     1
#define STATE_READY         2

#define TX_KEY_XOR          0xAA
#define RX_KEY_XOR          0x55

#define ENCRYPT_FLAG        0x80
#define CHANNEL_MASK        0x7F


// ========================================================================
// Types
// ========================================================================

struct SecLinkS {
    int           com;
    PktConnT     *pkt;
    SecDhT       *dh;
    SecCipherT   *txCipher;
    SecCipherT   *rxCipher;
    SecLinkRecvT  userCallback;
    void         *userCtx;
    int           state;
    uint8_t       myPub[SEC_DH_KEY_SIZE];
    uint8_t       remoteKey[SEC_DH_KEY_SIZE];
    bool          gotRemoteKey;
};


// ========================================================================
// Static prototypes (alphabetical)
// ========================================================================

static void completeHandshake(SecLinkT *link);
static void internalRecv(void *ctx, const uint8_t *data, int len);


// ========================================================================
// Static functions (alphabetical)
// ========================================================================

// Called when we've received the remote's public key. Computes the DH
// shared secret, derives directional cipher keys, and transitions to READY.
// After this, the DH context (containing the private key) is destroyed
// immediately — forward secrecy principle: even if the long-term state is
// compromised later, past session keys can't be recovered.
static void completeHandshake(SecLinkT *link) {
    uint8_t masterKey[SEC_XTEA_KEY_SIZE];
    uint8_t txKey[SEC_XTEA_KEY_SIZE];
    uint8_t rxKey[SEC_XTEA_KEY_SIZE];
    bool    weAreLower;
    secDhComputeSecret(link->dh, link->remoteKey, SEC_DH_KEY_SIZE);
    secDhDeriveKey(link->dh, masterKey, SEC_XTEA_KEY_SIZE);

    // Derive directional keys so each side encrypts with a unique key
    weAreLower = (memcmp(link->myPub, link->remoteKey, SEC_DH_KEY_SIZE) < 0);
    for (int i = 0; i < SEC_XTEA_KEY_SIZE; i++) {
        txKey[i] = masterKey[i] ^ (weAreLower ? TX_KEY_XOR : RX_KEY_XOR);
        rxKey[i] = masterKey[i] ^ (weAreLower ? RX_KEY_XOR : TX_KEY_XOR);
    }

    link->txCipher = secCipherCreate(txKey);
    link->rxCipher = secCipherCreate(rxKey);

    memset(masterKey, 0, sizeof(masterKey));
    memset(txKey, 0, sizeof(txKey));
    memset(rxKey, 0, sizeof(rxKey));

    // Destroy DH context — private key no longer needed
    secDhDestroy(link->dh);
    link->dh = 0;

    link->state = STATE_READY;
}


// Internal packet-layer callback. Routes incoming packets based on state:
//   - During handshake: expects a 128-byte DH public key
//   - When ready: strips the channel header, decrypts if flagged, and
//     forwards plaintext to the user callback
// The channel header is always one byte, so the minimum valid data packet
// is 1 byte (header only, zero-length payload).
static void internalRecv(void *ctx, const uint8_t *data, int len) {
    SecLinkT *link = (SecLinkT *)ctx;

    if (link->state == STATE_HANDSHAKE && !link->gotRemoteKey) {
        // During handshake, expect exactly one 128-byte public key
        if (len == SEC_DH_KEY_SIZE) {
            memcpy(link->remoteKey, data, len);
            link->gotRemoteKey = true;
            completeHandshake(link);
        }
    } else if (link->state == STATE_READY || link->state == STATE_INIT) {
        // Need at least the channel header byte
        if (len < 1) {
            return;
        }

        uint8_t       hdr        = data[0];
        uint8_t       channel    = hdr & CHANNEL_MASK;
        bool          encrypted  = (hdr & ENCRYPT_FLAG) != 0;
        const uint8_t *payload   = data + 1;
        int            payloadLen = len - 1;

        if (encrypted && link->state == STATE_READY) {
            uint8_t plaintext[PKT_MAX_PAYLOAD];
            if (payloadLen > (int)sizeof(plaintext)) {
                payloadLen = (int)sizeof(plaintext);
            }
            memcpy(plaintext, payload, payloadLen);
            secCipherCrypt(link->rxCipher, plaintext, payloadLen);
            if (link->userCallback) {
                link->userCallback(link->userCtx, plaintext, payloadLen, channel);
            }
        } else if (!encrypted) {
            if (link->userCallback) {
                link->userCallback(link->userCtx, payload, payloadLen, channel);
            }
        }
        // encrypted but not READY: silently drop
    }
}


// ========================================================================
// Public functions (alphabetical)
// ========================================================================

void secLinkClose(SecLinkT *link) {
    if (!link) {
        return;
    }

    if (link->txCipher) {
        secCipherDestroy(link->txCipher);
    }
    if (link->rxCipher) {
        secCipherDestroy(link->rxCipher);
    }
    if (link->dh) {
        secDhDestroy(link->dh);
    }
    if (link->pkt) {
        pktClose(link->pkt);
    }
    rs232Close(link->com);

    memset(link, 0, sizeof(SecLinkT));
    free(link);
}


int secLinkGetPending(SecLinkT *link) {
    if (!link) {
        return SECLINK_ERR_PARAM;
    }

    return pktGetPending(link->pkt);
}


int secLinkHandshake(SecLinkT *link) {
    int len;
    int rc;

    if (!link) {
        return SECLINK_ERR_PARAM;
    }

    // Generate DH keypair
    link->dh = secDhCreate();
    if (!link->dh) {
        return SECLINK_ERR_ALLOC;
    }

    rc = secDhGenerateKeys(link->dh);
    if (rc != SEC_SUCCESS) {
        return SECLINK_ERR_HANDSHAKE;
    }

    // Export our public key
    len = SEC_DH_KEY_SIZE;
    secDhGetPublicKey(link->dh, link->myPub, &len);

    // Enter handshake state and send our key
    link->state        = STATE_HANDSHAKE;
    link->gotRemoteKey = false;

    rc = pktSend(link->pkt, link->myPub, SEC_DH_KEY_SIZE, true);
    if (rc != PKT_SUCCESS) {
        return SECLINK_ERR_HANDSHAKE;
    }

    // Poll until the callback completes the handshake
    while (link->state != STATE_READY) {
        if (pktPoll(link->pkt) == PKT_ERR_DISCONNECTED) {
            return SECLINK_ERR_HANDSHAKE;
        }
    }

    return SECLINK_SUCCESS;
}


bool secLinkIsReady(SecLinkT *link) {
    if (!link) {
        return false;
    }
    return link->state == STATE_READY;
}


SecLinkT *secLinkOpen(int com, int32_t bps, int dataBits, char parity, int stopBits, int handshake, SecLinkRecvT callback, void *ctx) {
    SecLinkT *link;
    int       rc;

    link = (SecLinkT *)calloc(1, sizeof(SecLinkT));
    if (!link) {
        return 0;
    }

    link->com          = com;
    link->userCallback = callback;
    link->userCtx      = ctx;
    link->state        = STATE_INIT;

    // Open serial port
    rc = rs232Open(com, bps, dataBits, parity, stopBits, handshake);
    if (rc != RS232_SUCCESS) {
        free(link);
        return 0;
    }

    // Open packet layer with our internal callback
    link->pkt = pktOpen(com, 0, internalRecv, link);
    if (!link->pkt) {
        rs232Close(com);
        free(link);
        return 0;
    }

    return link;
}


int secLinkPoll(SecLinkT *link) {
    if (!link) {
        return SECLINK_ERR_PARAM;
    }

    return pktPoll(link->pkt);
}


int secLinkSend(SecLinkT *link, const uint8_t *data, int len, uint8_t channel, bool encrypt, bool block) {
    uint8_t buf[PKT_MAX_PAYLOAD];
    int     rc;

    if (!link || !data) {
        return SECLINK_ERR_PARAM;
    }
    if (channel > SECLINK_MAX_CHANNEL) {
        return SECLINK_ERR_PARAM;
    }
    if (encrypt && link->state != STATE_READY) {
        return SECLINK_ERR_NOT_READY;
    }
    if (len <= 0 || len > SECLINK_MAX_PAYLOAD) {
        return SECLINK_ERR_PARAM;
    }

    // CRITICAL: check window space BEFORE encrypting. If we encrypted first
    // and then the send failed, the cipher counter would have advanced but
    // the data wouldn't have been sent, permanently desynchronizing the
    // TX cipher state from the remote's RX cipher state.
    if (!block && !pktCanSend(link->pkt)) {
        return SECLINK_ERR_SEND;
    }

    // Build channel header byte
    buf[0] = channel & CHANNEL_MASK;
    if (encrypt) {
        buf[0] |= ENCRYPT_FLAG;
    }

    // Copy payload after header
    memcpy(buf + 1, data, len);

    // Encrypt the payload portion only (not the header)
    if (encrypt) {
        secCipherCrypt(link->txCipher, buf + 1, len);
    }

    rc = pktSend(link->pkt, buf, len + 1, block);
    if (rc != PKT_SUCCESS) {
        return SECLINK_ERR_SEND;
    }

    return SECLINK_SUCCESS;
}


// Convenience function for sending large buffers. Splits the data into
// SECLINK_MAX_PAYLOAD (254 byte) chunks and sends each one, blocking until
// the send window has room. The receiver sees multiple packets on the same
// channel and must reassemble if needed. Always blocking because the caller
// expects the entire buffer to be sent when this returns.
int secLinkSendBuf(SecLinkT *link, const uint8_t *data, int len, uint8_t channel, bool encrypt) {
    int offset = 0;
    int rc;

    if (!link || !data) {
        return SECLINK_ERR_PARAM;
    }
    if (len <= 0) {
        return SECLINK_ERR_PARAM;
    }

    while (offset < len) {
        int chunk = len - offset;
        if (chunk > SECLINK_MAX_PAYLOAD) {
            chunk = SECLINK_MAX_PAYLOAD;
        }

        rc = secLinkSend(link, data + offset, chunk, channel, encrypt, true);
        if (rc != SECLINK_SUCCESS) {
            return rc;
        }

        offset += chunk;
    }

    return SECLINK_SUCCESS;
}