/*
 * kl_transport.c
 * Transport layer: packet framing, byte-by-byte complement ack, counter.
 *
 * All operations are non-blocking state machines driven by
 * KL_Transport_Service() called from the main loop.
 */

#include "kl_transport.h"
#include "kl_phy.h"
#include "main.h"
#include <string.h>

static KL_Transport tl;

/* ================================================================
 * Init / Reset
 * ================================================================ */
void KL_Transport_Init(void)
{
    memset(&tl, 0, sizeof(tl));
    tl.tx_state = KL_TL_TX_IDLE;
    tl.rx_state = KL_TL_RX_IDLE;
    tl.raw_state = KL_TL_RAW_IDLE;
}

void KL_Transport_ResetCounter(void)
{
    tl.pkt_counter = 0;
    tl.counter_initialized = 0;
}

/* ================================================================
 * Packet TX
 * ================================================================ */
int KL_Transport_SendPacket(const uint8_t *payload, uint8_t payload_len)
{
    if (tl.tx_state != KL_TL_TX_IDLE) return 0;
    if (payload_len == 0 || payload_len > KL_MAX_PACKET_PAYLOAD) return 0;

    /* Serialize: [LENGTH][COUNTER][payload...] */
    uint8_t length_field = payload_len + 2;  /* +2 for COUNTER and LENGTH convention */
    uint8_t idx = 0;
    tl.tx_buf[idx++] = length_field;
    tl.tx_buf[idx++] = tl.pkt_counter++;
    memcpy(&tl.tx_buf[idx], payload, payload_len);
    idx += payload_len;

    tl.tx_len       = idx;
    tl.tx_idx       = 0;
    tl.tx_deadline  = HAL_GetTick() + KL_PKT_TMO_MS;
    tl.tx_next_tick = HAL_GetTick();
    tl.tx_state     = KL_TL_TX_SEND_BYTE;
    return 1;
}

void KL_Transport_SendAck(void)
{
    uint8_t payload = (uint8_t)KL_CMD_ACK;
    KL_Transport_SendPacket(&payload, 1);
}

void KL_Transport_SendNak(void)
{
    uint8_t payload = (uint8_t)KL_CMD_NAK;
    KL_Transport_SendPacket(&payload, 1);
}

int KL_Transport_TxBusy(void)
{
    return (tl.tx_state != KL_TL_TX_IDLE &&
            tl.tx_state != KL_TL_TX_DONE &&
            tl.tx_state != KL_TL_TX_ERROR) ? 1 : 0;
}

int KL_Transport_TxDone(void)  { return (tl.tx_state == KL_TL_TX_DONE)  ? 1 : 0; }
int KL_Transport_TxError(void) { return (tl.tx_state == KL_TL_TX_ERROR) ? 1 : 0; }

void KL_Transport_TxReset(void)
{
    tl.tx_state = KL_TL_TX_IDLE;
}

/* ================================================================
 * TX state machine
 * ================================================================ */
static void tl_tx_service(void)
{
    uint32_t now = HAL_GetTick();
    uint8_t  byte_val;

    switch (tl.tx_state) {

    case KL_TL_TX_SEND_BYTE:
        if (kl_tick_diff(now, tl.tx_next_tick) < 0) break;
        if (tl.tx_idx >= tl.tx_len) {
            /* All data bytes sent, send terminator */
            tl.tx_state = KL_TL_TX_SEND_TERM;
            break;
        }
        byte_val = tl.tx_buf[tl.tx_idx];
        if (KL_Phy_TxByte(byte_val)) {
            tl.tx_last_byte = byte_val;
            tl.tx_idx++;
            tl.tx_state = KL_TL_TX_WAIT_PHY;
        }
        break;

    case KL_TL_TX_WAIT_PHY:
        if (kl_tick_diff(now, tl.tx_deadline) >= 0) {
            tl.tx_state = KL_TL_TX_ERROR;
            break;
        }
        if (!KL_Phy_TxBusy()) {
            /* Echo consumed by PHY, now wait for master's complement */
            tl.tx_state = KL_TL_TX_WAIT_COMPLEMENT;
            tl.tx_next_tick = now;  /* Record when we started waiting */
        }
        break;

    case KL_TL_TX_WAIT_COMPLEMENT:
        if (kl_tick_diff(now, tl.tx_deadline) >= 0) {
            tl.tx_state = KL_TL_TX_ERROR;
            break;
        }
        {
            uint8_t rx_byte;
            if (KL_Phy_RxPop(&rx_byte)) {
                if (rx_byte == (uint8_t)~tl.tx_last_byte) {
                    /* Valid complement received, advance to next byte */
                    tl.tx_next_tick = HAL_GetTick() + KL_P1_GAP_MS;
                    tl.tx_state = KL_TL_TX_SEND_BYTE;
                } else {
                    /* Unexpected byte -- protocol error */
                    tl.tx_state = KL_TL_TX_ERROR;
                }
            }
        }
        break;

    case KL_TL_TX_SEND_TERM:
        if (KL_Phy_TxByte(KL_PACKET_END)) {
            tl.tx_state = KL_TL_TX_WAIT_TERM_PHY;
        }
        break;

    case KL_TL_TX_WAIT_TERM_PHY:
        if (kl_tick_diff(now, tl.tx_deadline) >= 0) {
            tl.tx_state = KL_TL_TX_ERROR;
            break;
        }
        if (!KL_Phy_TxBusy()) {
            tl.tx_state = KL_TL_TX_DONE;
        }
        break;

    case KL_TL_TX_IDLE:
    case KL_TL_TX_DONE:
    case KL_TL_TX_ERROR:
        break;
    }
}

/* ================================================================
 * Packet RX
 * ================================================================ */
int KL_Transport_StartReceive(uint32_t timeout_ms)
{
    if (tl.rx_state != KL_TL_RX_IDLE) return 0;

    memset(&tl.rx_packet, 0, sizeof(tl.rx_packet));
    tl.rx_idx           = 0;
    tl.rx_pkt_len       = 0;
    tl.rx_expected_total = 0;
    tl.rx_deadline      = HAL_GetTick() + timeout_ms;
    tl.rx_byte_deadline = HAL_GetTick() + timeout_ms;
    tl.rx_state         = KL_TL_RX_WAIT_LEN;
    return 1;
}

int KL_Transport_RxBusy(void)
{
    return (tl.rx_state != KL_TL_RX_IDLE &&
            tl.rx_state != KL_TL_RX_DONE &&
            tl.rx_state != KL_TL_RX_ERROR) ? 1 : 0;
}

int KL_Transport_RxDone(void)  { return (tl.rx_state == KL_TL_RX_DONE)  ? 1 : 0; }
int KL_Transport_RxError(void) { return (tl.rx_state == KL_TL_RX_ERROR) ? 1 : 0; }

const KL_Packet* KL_Transport_GetRxPacket(void) { return &tl.rx_packet; }

void KL_Transport_RxReset(void)
{
    tl.rx_state = KL_TL_RX_IDLE;
}

int KL_Transport_RxHadData(void)
{
    return (tl.rx_idx > 0) ? 1 : 0;
}

/* ================================================================
 * RX helper: send complement and advance to wait-for-PHY state
 * ================================================================ */
static int rx_send_complement(uint8_t byte, KL_TlRxState next_after_phy)
{
    uint8_t compl = (uint8_t)~byte;
    if (KL_Phy_TxByte(compl)) {
        tl.rx_last_byte = byte;
        tl.rx_state     = KL_TL_RX_WAIT_COMPL_PHY;
        /* Store the state to transition to after PHY completes.
         * We encode the "next state after complement" in rx_expected_total
         * as a side channel -- but cleaner to use a dedicated field.
         * Instead, we use a convention: after complement PHY done,
         * we look at rx_idx to determine what to do next. */
        (void)next_after_phy;
        return 1;
    }
    return 0;
}

/* ================================================================
 * RX state machine
 * ================================================================ */
static void tl_rx_service(void)
{
    uint32_t now = HAL_GetTick();
    uint8_t  byte_val;

    switch (tl.rx_state) {

    case KL_TL_RX_WAIT_LEN:
        if (kl_tick_diff(now, tl.rx_deadline) >= 0) {
            tl.rx_state = KL_TL_RX_ERROR;
            break;
        }
        if (KL_Phy_RxPop(&byte_val)) {
            tl.rx_buf[tl.rx_idx++] = byte_val;
            tl.rx_pkt_len = byte_val;
            tl.rx_packet.length = byte_val;
            tl.rx_byte_deadline = now + KL_BYTE_TMO_MS;
            /* Send complement for LENGTH */
            tl.rx_state = KL_TL_RX_COMPL_LEN;
        }
        break;

    case KL_TL_RX_COMPL_LEN:
        if (kl_tick_diff(now, tl.rx_deadline) >= 0) {
            tl.rx_state = KL_TL_RX_ERROR;
            break;
        }
        {
            uint8_t compl = (uint8_t)~tl.rx_buf[0];
            if (KL_Phy_TxByte(compl)) {
                tl.rx_state = KL_TL_RX_WAIT_COMPL_PHY;
            }
        }
        break;

    case KL_TL_RX_WAIT_CNT:
        if (kl_tick_diff(now, tl.rx_deadline) >= 0) {
            tl.rx_state = KL_TL_RX_ERROR;
            break;
        }
        if (KL_Phy_RxPop(&byte_val)) {
            tl.rx_buf[tl.rx_idx++] = byte_val;
            tl.rx_byte_deadline = now + KL_BYTE_TMO_MS;

            /* Validate/sync packet counter */
            if (!tl.counter_initialized) {
                tl.pkt_counter = byte_val;
                tl.counter_initialized = 1;
            }
            tl.rx_packet.counter = byte_val;
            tl.pkt_counter = byte_val + 1;

            /* Send complement for COUNTER */
            tl.rx_state = KL_TL_RX_COMPL_CNT;
        }
        break;

    case KL_TL_RX_COMPL_CNT:
        if (kl_tick_diff(now, tl.rx_deadline) >= 0) {
            tl.rx_state = KL_TL_RX_ERROR;
            break;
        }
        {
            uint8_t compl = (uint8_t)~tl.rx_buf[1];
            if (KL_Phy_TxByte(compl)) {
                tl.rx_state = KL_TL_RX_WAIT_COMPL_PHY;
            }
        }
        break;

    case KL_TL_RX_WAIT_BYTE:
        if (kl_tick_diff(now, tl.rx_deadline) >= 0) {
            tl.rx_state = KL_TL_RX_ERROR;
            break;
        }
        if (KL_Phy_RxPop(&byte_val)) {
            tl.rx_buf[tl.rx_idx++] = byte_val;
            tl.rx_byte_deadline = now + KL_BYTE_TMO_MS;
            /* Always send complement for every non-terminator byte.
             * WAIT_COMPL_PHY will decide whether to go to WAIT_TERM
             * or back to WAIT_BYTE based on rx_idx vs rx_pkt_len. */
            tl.rx_state = KL_TL_RX_COMPL_BYTE;
        }
        break;

    case KL_TL_RX_COMPL_BYTE:
        if (kl_tick_diff(now, tl.rx_deadline) >= 0) {
            tl.rx_state = KL_TL_RX_ERROR;
            break;
        }
        {
            uint8_t compl = (uint8_t)~tl.rx_buf[tl.rx_idx - 1];
            if (KL_Phy_TxByte(compl)) {
                tl.rx_state = KL_TL_RX_WAIT_COMPL_PHY;
            }
        }
        break;

    case KL_TL_RX_WAIT_COMPL_PHY:
        if (kl_tick_diff(now, tl.rx_deadline) >= 0) {
            tl.rx_state = KL_TL_RX_ERROR;
            break;
        }
        if (!KL_Phy_TxBusy()) {
            /* Complement echo done. Determine next state based on progress.
             * rx_idx tells us how many bytes we've stored so far:
             *   1 = just finished LEN complement  -> wait for CNT
             *   2 = just finished CNT complement  -> wait for CMD/data
             *   >=3 = CMD or data complement done -> wait for next byte or terminator
             */
            if (tl.rx_idx == 1) {
                tl.rx_state = KL_TL_RX_WAIT_CNT;
            } else if (tl.rx_idx == 2) {
                tl.rx_state = KL_TL_RX_WAIT_BYTE;
            } else {
                /* Check if all non-terminator bytes are collected */
                if (tl.rx_idx >= tl.rx_pkt_len) {
                    tl.rx_state = KL_TL_RX_WAIT_TERM;
                } else {
                    tl.rx_state = KL_TL_RX_WAIT_BYTE;
                }
            }
        }
        break;

    case KL_TL_RX_WAIT_TERM:

        if (kl_tick_diff(now, tl.rx_deadline) >= 0) {
            tl.rx_state = KL_TL_RX_ERROR;
            break;
        }
        if (KL_Phy_RxPop(&byte_val)) {
            if (byte_val != KL_PACKET_END) {
                tl.rx_state = KL_TL_RX_ERROR;
                break;
            }
            /* Parse into KL_Packet structure */
            tl.rx_packet.command  = tl.rx_buf[2];          /* CMD at index 2 */
            tl.rx_packet.data_len = tl.rx_pkt_len - 3;     /* pkt_len - 3 = data count */
            if (tl.rx_packet.data_len > KL_MAX_PACKET_PAYLOAD) {
                tl.rx_packet.data_len = KL_MAX_PACKET_PAYLOAD;
            }
            memcpy(tl.rx_packet.data, &tl.rx_buf[3], tl.rx_packet.data_len);
            tl.rx_state = KL_TL_RX_DONE;
        }
        break;

    case KL_TL_RX_IDLE:
    case KL_TL_RX_DONE:
    case KL_TL_RX_ERROR:
        break;
    }
}

/* ================================================================
 * Raw byte TX (for handshake -- no framing, no complement)
 * ================================================================ */
int KL_Transport_SendRawByte(uint8_t byte)
{
    if (tl.raw_state != KL_TL_RAW_IDLE) return 0;
    if (!KL_Phy_TxByte(byte)) return 0;
    tl.raw_state = KL_TL_RAW_WAIT_PHY;
    return 1;
}

int KL_Transport_RawTxDone(void)
{
    return (tl.raw_state == KL_TL_RAW_DONE) ? 1 : 0;
}

void KL_Transport_RawTxReset(void)
{
    tl.raw_state = KL_TL_RAW_IDLE;
}

static void tl_raw_service(void)
{
    if (tl.raw_state == KL_TL_RAW_WAIT_PHY) {
        if (!KL_Phy_TxBusy()) {
            tl.raw_state = KL_TL_RAW_DONE;
        }
    }
}

/* ================================================================
 * Combined service
 * ================================================================ */
void KL_Transport_Service(void)
{
    tl_tx_service();
    tl_rx_service();
    tl_raw_service();
}
