/*
 * Author: Kilian
 * Date: 2025-04-20
 * Description: This file contains a TX-only implementation of the ISO-TP protocol.
 *              
 * References:
 * https://munich.dissec.to/kb/chapters/isotp/isotp.html#protocol-basics-iso-15765-2
 * https://www.kernel.org/doc/html/next//networking/iso15765-2.html
 * https://en.wikipedia.org/wiki/ISO_15765-2
 * https://ramn.readthedocs.io/en/latest/userguide/isotp_tutorial.html
 */
#include "isotp.h"

#include <stdint.h>
#include <stddef.h>
#include <string.h>

#define MAX_CAN_DATA_SIZE 8
#define MAX_ISOTP_DATA_SIZE_MESSAGE 4095
#define MAX_ISOTP_DATA_SINGLE 7
#define MAX_ISOTP_DATA_FIRST 6
#define MAX_ISOTP_DATA_CONSECUTIVE 7

typedef union [[gnu::packed]] {
  struct {
    uint8_t header;
    uint8_t data[MAX_ISOTP_DATA_SINGLE];
  } single;
  struct {
    uint8_t header_1;
    uint8_t header_2;
    uint8_t data[MAX_ISOTP_DATA_FIRST];
  } first;
  struct {
    uint8_t header;
    uint8_t data[MAX_ISOTP_DATA_CONSECUTIVE];
  } consecutive;
} isotp_message;

enum {
    ISOTP_SINGLE_FRAME = 0 << 4,
    ISOTP_FIRST_FRAME = 1 << 4,
    ISOTP_CONSECUTIVE_FRAME = 2 << 4,
    ISOTP_FLOW_CONTROL = 3 << 4
};

struct isotp_message_state {
    const uint8_t * data;
    size_t remaining;
    uint16_t id;
    uint32_t timestamp;    // Timestamp for timeout handling
    enum : uint8_t {
        ISOTP_NEW_MESSAGE = 0,
        ISOTP_WAITING_FLOW_CONTROL = 1,
        ISOTP_READY = 2,
        ISOTP_DONE = 3,
        ISOTP_FAILED = 4
    } state;
    struct {
        enum : uint8_t {
            ISOTP_CONTINUE = 0,
            ISOTP_WAIT = 1,
            ISOTP_ABORT = 2
        } state;
        uint8_t block_size;
        uint8_t st_min; // in ms, we round up to 1 ms if sub 1 ms
        uint8_t st_min_counter;
    } flow_control;
    uint8_t index;
};

static struct isotp_message_state tx_queue[MAX_ISOTP_MESSAGES_IN_FLIGHT] = {[0 ... MAX_ISOTP_MESSAGES_IN_FLIGHT - 1] = {.state=ISOTP_DONE }};

[[gnu::weak]] bool isotp_can_transmit(uint16_t, const uint8_t *, size_t) { return false; }
[[gnu::weak]] uint32_t isotp_can_get_time() { return 0; }

uint32_t last_called;

isotp_status_t isotp_init() {
    // platform-specific initialization
    last_called = isotp_can_get_time() - 100;
    return (isotp_status_t) {ISOTP_OK};
}

isotp_status_t isotp_add_message(uint16_t id, const uint8_t * data, size_t datalen) {
    if (datalen > MAX_ISOTP_DATA_SIZE_MESSAGE) {
        return ISOTP_DATA_TOO_LONG;
    }

    if (data == NULL || datalen == 0) {
        return ISOTP_INVALID_PARAMETER;
    }

    for (size_t i = 0; i < MAX_ISOTP_MESSAGES_IN_FLIGHT; i++) {
        if (!(tx_queue[i].state == ISOTP_DONE || tx_queue[i].state == ISOTP_FAILED) && tx_queue[i].id == id) {
            return ISOTP_MESSAGE_ALREADY_IN_FLIGHT;
        }
    }

    for (size_t i = 0; i < MAX_ISOTP_MESSAGES_IN_FLIGHT; i++) {
        if (tx_queue[i].state == ISOTP_DONE || tx_queue[i].state == ISOTP_FAILED) {
            tx_queue[i] = (struct isotp_message_state) {data, datalen, id, 0, ISOTP_NEW_MESSAGE, {}, 1};
            return ISOTP_OK;
        }
    }

    return ISOTP_TOO_MANY_MESSAGES;
}

/* Macro to check CAN transmission result and handle errors */
#define TRY_CAN_TRANSMIT(id, data, size) \
    if (!isotp_can_transmit(id, data, size)) { \
        tx_queue[i].state = ISOTP_FAILED; \
        return ISOTP_CAN_TRANSMIT_ERROR; \
    }

isotp_status_t isotp_update() {
    uint32_t now = isotp_can_get_time();
    uint32_t delta = now - last_called;
    for (size_t i = 0; i < MAX_ISOTP_MESSAGES_IN_FLIGHT; i++) {

        isotp_message msg = {};

        switch (tx_queue[i].state) {
            case ISOTP_NEW_MESSAGE:
                if (tx_queue[i].remaining <= MAX_ISOTP_DATA_SINGLE) {
                    msg.single.header = ISOTP_SINGLE_FRAME | (tx_queue[i].remaining & 0x0F);
                    memcpy(msg.single.data, tx_queue[i].data, tx_queue[i].remaining);
                    #if ISOTP_PADDING
                    memset(msg.single.data + tx_queue[i].remaining, ISOTP_PADDING, MAX_ISOTP_DATA_SINGLE - tx_queue[i].remaining);
                    TRY_CAN_TRANSMIT(tx_queue[i].id, (const uint8_t *)&msg, MAX_CAN_DATA_SIZE);
                    #else
                    TRY_CAN_TRANSMIT(tx_queue[i].id, (const uint8_t *)&msg, tx_queue[i].remaining + 1);
                    #endif
                    tx_queue[i].remaining = 0;
                    tx_queue[i].state = ISOTP_DONE;
                } else {
                    msg.first.header_1 = ISOTP_FIRST_FRAME | ((tx_queue[i].remaining >> 8) & 0x0F);
                    msg.first.header_2 = tx_queue[i].remaining & 0xFF;
                    memcpy(msg.first.data, tx_queue[i].data, MAX_ISOTP_DATA_FIRST);
                    TRY_CAN_TRANSMIT(tx_queue[i].id, (const uint8_t *)&msg, MAX_CAN_DATA_SIZE);
                    tx_queue[i].remaining -= MAX_ISOTP_DATA_FIRST;
                    tx_queue[i].data += MAX_ISOTP_DATA_FIRST;
                    tx_queue[i].timestamp = now; // Set timestamp when entering WAITING_FLOW_CONTROL state
                    tx_queue[i].state = ISOTP_WAITING_FLOW_CONTROL;
                }
                break;
            case ISOTP_READY:
            ISOTP_SEND_AGAIN:
                if (tx_queue[i].flow_control.st_min_counter <= delta) {
                    tx_queue[i].flow_control.st_min_counter = tx_queue[i].flow_control.st_min;
                } else {
                    tx_queue[i].flow_control.st_min_counter -= delta;
                    break;
                }
                if (tx_queue[i].flow_control.block_size == 1) {
                    tx_queue[i].state = ISOTP_WAITING_FLOW_CONTROL;
                }
                if (tx_queue[i].flow_control.block_size != 0) {
                    tx_queue[i].flow_control.block_size--;
                }

                msg.consecutive.header = ISOTP_CONSECUTIVE_FRAME | (tx_queue[i].index & 0x0F);
                size_t bytes_to_send = tx_queue[i].remaining <= MAX_ISOTP_DATA_CONSECUTIVE ? tx_queue[i].remaining : MAX_ISOTP_DATA_CONSECUTIVE;
                memcpy(msg.consecutive.data, tx_queue[i].data, bytes_to_send);
                
                #if ISOTP_PADDING
                if (bytes_to_send < MAX_ISOTP_DATA_CONSECUTIVE) {
                    memset(msg.consecutive.data + bytes_to_send, ISOTP_PADDING, MAX_ISOTP_DATA_CONSECUTIVE - bytes_to_send);
                    TRY_CAN_TRANSMIT(tx_queue[i].id, (const uint8_t *)&msg, MAX_CAN_DATA_SIZE);
                } else {
                    TRY_CAN_TRANSMIT(tx_queue[i].id, (const uint8_t *)&msg, MAX_CAN_DATA_SIZE);
                }
                #else
                TRY_CAN_TRANSMIT(tx_queue[i].id, (const uint8_t *)&msg, bytes_to_send + 1);
                #endif
                
                tx_queue[i].data += bytes_to_send;
                tx_queue[i].remaining -= bytes_to_send;
                tx_queue[i].index++;
                
                if (tx_queue[i].remaining == 0) {
                    tx_queue[i].state = ISOTP_DONE;
                } 
                
                if (tx_queue[i].state == ISOTP_READY) {
                    if (tx_queue[i].flow_control.st_min == 0) { // no delay needed, and we are not waiting for flow control -> send again
                        goto ISOTP_SEND_AGAIN;
                    }
                }
                break;
            case ISOTP_WAITING_FLOW_CONTROL:
                if (now - tx_queue[i].timestamp > ISOTP_FC_WAIT_TIMEOUT) {
                    tx_queue[i].state = ISOTP_FAILED; // Timeout handling
                }
                break;
            case ISOTP_DONE:
            case ISOTP_FAILED:
                break;
        }
    }
    last_called = now;
    return ISOTP_OK;
}

isotp_status_t isotp_handle_flow_control(uint16_t id, const uint8_t * data, size_t datalen) {
    if (datalen < 3) {
        return ISOTP_NOT_FLOW_CONTROL;
    }

    if (data == NULL) {
        return ISOTP_INVALID_PARAMETER;
    }

    for (size_t i = 0; i < MAX_ISOTP_MESSAGES_IN_FLIGHT; i++) {
        if (tx_queue[i].state == ISOTP_DONE || tx_queue[i].state == ISOTP_FAILED || tx_queue[i].id != id) {
            continue;
        }

        uint8_t type = data[0] & 0xF0;
        uint8_t status = data[0] & 0x0F;
        uint8_t block_size = data[1];
        uint8_t st_min = data[2];

        if (type != ISOTP_FLOW_CONTROL) {
            return ISOTP_NOT_FLOW_CONTROL;
        }

        switch (status) {
            case 0: // Continue
                tx_queue[i].flow_control.state = ISOTP_CONTINUE;
                break;
            case 1: // Wait
                tx_queue[i].flow_control.state = ISOTP_WAIT;
                tx_queue[i].state = ISOTP_WAITING_FLOW_CONTROL;
                tx_queue[i].timestamp = isotp_can_get_time(); // Update timestamp for timeout handling
                return ISOTP_OK;
            case 2: // Abort
                tx_queue[i].flow_control.state = ISOTP_ABORT;
                tx_queue[i].state = ISOTP_FAILED;
                return ISOTP_OK;
            default:
                return ISOTP_NOT_FLOW_CONTROL;
        }

        tx_queue[i].flow_control.block_size = block_size;
        tx_queue[i].flow_control.st_min = st_min > 0x7F ? 1 : st_min; // round up to 1 ms
        tx_queue[i].flow_control.st_min_counter = st_min > 0x7F ? 1 : st_min;
        tx_queue[i].state = ISOTP_READY;

        return ISOTP_OK;
    }
    return ISOTP_INVALID_PARAMETER;
}

