#include <stdint.h>
#include <stdbool.h>
#include "pp/payload_parser.h"
#include "pp/payload_builder.h"
#include "modbus.h"

/**
 * Function to calculate MODBUS CRC.
 *
 * https://github.com/starnight/MODBUS-CRC/blob/master/modbuscrc.c
 **/
static uint16_t crc16_update(uint16_t crc, uint8_t a)
{
    int i;
    crc ^= (uint16_t) a;
    for (i = 0; i < 8; ++i) {
        if (crc & 1) {
            crc = (crc >> 1) ^ 0xA001;
        } else {
            crc = (crc >> 1);
        }
    }
    return crc;
}

static inline uint16_t crc16_init()
{
    return 0xFFFF;
}

/**
 * Handle a modbus request
 */
ModbusError_t mb_handleRequest(
        ModbusSlave_t *ms,
        const uint8_t *req,
        size_t req_size,
        uint8_t *resp,
        size_t resp_capacity,
        size_t *resp_size
)
{
    uint16_t txn_id, protocol_id, ref;

#ifdef MB_SUPPORT_FC22
    uint16_t and_mask, or_mask;
#endif

#ifdef MB_SUPPORT_FC23
    uint16_t ref2, count2;
#endif

#if defined(MB_SUPPORT_FC01) || defined(MB_SUPPORT_FC02)
    bool bvalue;
#endif

#if defined(MB_SUPPORT_FC03) || defined(MB_SUPPORT_FC04) \
    || defined(MB_SUPPORT_FC16) || defined(MB_SUPPORT_FC22) || defined(MB_SUPPORT_FC23) \
    || defined(MB_SUPPORT_FC05) || defined(MB_SUPPORT_FC06)
    uint16_t value;
#endif

#if defined(MB_SUPPORT_FC01) || defined(MB_SUPPORT_FC02) || defined(MB_SUPPORT_FC15)
    uint8_t scratch, bytecount, bitcnt;
#endif

#if defined(MB_SUPPORT_FC01) || defined(MB_SUPPORT_FC02) || defined(MB_SUPPORT_FC03) \
    || defined(MB_SUPPORT_FC04) || defined(MB_SUPPORT_FC15) || defined(MB_SUPPORT_FC16) \
    || defined(MB_SUPPORT_FC23)
    uint16_t count;
#endif

    ModbusException_t exc = MB_EXCEPTION_OK;
    size_t numbytes;
    uint8_t fcx;

    const size_t TCP_RESP_OVERHEAD = 8;
    const size_t RTU_RESP_OVERHEAD = 4;
    const size_t overhead = ms->proto == MB_PROTO_TCP ? TCP_RESP_OVERHEAD : RTU_RESP_OVERHEAD;
    pb_mark_t resp_fc_mark = NULL, resp_len_mark = NULL, resp_pld_start_mark = NULL;

    const bool tcp = ms->proto == MB_PROTO_TCP;
    PayloadParser pp = pp_start_be(req, req_size, NULL);
    PayloadBuilder pb = pb_start_be(resp, resp_capacity, NULL);

    if (tcp) {
        /* Parse header */
        txn_id = pp_u16(&pp);
        protocol_id = pp_u16(&pp);
        pp_skip(&pp, 2); // msglen
        if (protocol_id != 0) {
            return MB_ERR_BADPROTO;
        }
    } else {
        /* check  CRC */
        uint16_t crc = crc16_init();
        for (int pos = 0; pos < req_size /* size of CRC */; pos++) {
            crc = crc16_update(crc, pp_u8(&pp));
        }
        if (crc != 0) {
            return MB_ERR_CHECKSUM;
        }
        pp_rewind(&pp);
    }

    const uint8_t slave_id = pp_u8(&pp);

    /* check addressing (don't check for TCP - UnitID is used e.g. for gateway target devices) */
    if (!tcp && slave_id != ms->addr) {
        return MB_ERR_NOTFORME;
    }

    fcx = pp_u8(&pp);

    if (!pp_ok(&pp)) {
        return MB_ERR_NEEDMORE;
    }

    /* start building the response */
    if (tcp) {
        pb_u16(&pb, txn_id);
        pb_u16(&pb, protocol_id);
        resp_len_mark = pb_save(&pb);
        pb_u16(&pb, 0); // Placeholder for LEN
    }
    pb_u8(&pb, slave_id);
    resp_fc_mark = pb_save(&pb);
    pb_u8(&pb, fcx); // Exceptions will add 0x80 to this
    resp_pld_start_mark = pb_save(&pb);

    if (ms->startOfAccess) {
        exc = ms->startOfAccess(ms, fcx, slave_id);
        if (exc != 0) {
            goto exception;
        }
    }

    switch (fcx) {
#ifdef MB_SUPPORT_FC05
        case FC05_WRITE_SINGLE_COIL:
            if (!ms->writeCoil) {
                exc = MB_EXCEPTION_ILLEGAL_FUNCTION;
                goto exception;
            }
            ref = pp_u16(&pp);
            value = pp_u16(&pp);
            if (!pp_ok(&pp)) {
                return MB_ERR_NEEDMORE;
            }
            if (resp_capacity < 4 + overhead) {
                exc = MB_EXCEPTION_ILLEGAL_DATA_ADDRESS; // TODO?
                goto exception;
            }
            if (value == 0xFF00) {
                exc = ms->writeCoil(ms, ref, true);
            } else {
                exc = ms->writeCoil(ms, ref, false);
            }
            if (exc != 0) {
                goto exception;
            }
            pb_u16(&pb, ref);
            pb_u16(&pb, value);
            break;
#endif

#if defined(MB_SUPPORT_FC01) || defined(MB_SUPPORT_FC02)
        case FC01_READ_COILS:
        case FC02_READ_DISCRETES:
            /* check we have the needed function */
#ifdef MB_SUPPORT_FC01
            if (fcx == FC01_READ_COILS && !ms->readCoil) {
                exc = MB_EXCEPTION_ILLEGAL_FUNCTION;
                goto exception;
            }
#endif
#ifdef MB_SUPPORT_FC02
            if (fcx == FC02_READ_DISCRETES && !ms->readDiscrete) {
                exc = MB_EXCEPTION_ILLEGAL_FUNCTION;
                goto exception;
            }
#endif
            ref = pp_u16(&pp);
            count = pp_u16(&pp);
            if (!pp_ok(&pp)) {
                return MB_ERR_NEEDMORE;
            }
            bytecount = (count + 7) / 8;
            if (count > 255 * 8 || resp_capacity < bytecount + overhead) {
                exc = MB_EXCEPTION_ILLEGAL_DATA_ADDRESS; // TODO?
                goto exception;
            }
            pb_u8(&pb, bytecount);
            bitcnt = 0;
            scratch = 0;
            while (count-- > 0) {
#ifdef MB_SUPPORT_FC01
                if (fcx == FC01_READ_COILS) {
                    exc = ms->readCoil(ms, ref++, &bvalue);
                }
#endif
#ifdef MB_SUPPORT_FC02
                if (fcx == FC02_READ_DISCRETES) {
                    exc = ms->readDiscrete(ms, ref++, &bvalue);
                }
#endif
                if (exc != 0) {
                    goto exception;
                }
                scratch |= ((bvalue & 1) << bitcnt);
                bitcnt++;
                if (bitcnt == 8) {
                    pb_u8(&pb, scratch);
                    scratch = 0;
                    bitcnt = 0;
                }
            }
            if (bitcnt != 0) {
                pb_u8(&pb, scratch);
            }
            break;
#endif

#if defined(MB_SUPPORT_FC03) || defined(MB_SUPPORT_FC04)
        case FC03_READ_HOLDING_REGISTERS:
        case FC04_READ_INPUT_REGISTERS:
#ifdef MB_SUPPORT_FC03
            /* check we have the needed function */
            if (fcx == FC03_READ_HOLDING_REGISTERS && !ms->readHolding) {
                exc = MB_EXCEPTION_ILLEGAL_FUNCTION;
                goto exception;
            }
#endif
#ifdef MB_SUPPORT_FC04
            if (fcx == FC04_READ_INPUT_REGISTERS && !ms->readInput) {
                exc = MB_EXCEPTION_ILLEGAL_FUNCTION;
                goto exception;
            }
#endif
            ref = pp_u16(&pp);
            count = pp_u16(&pp);
            if (!pp_ok(&pp)) {
                return MB_ERR_NEEDMORE;
            }
            if (count > 255 || resp_capacity < count * 2 + overhead) {
                exc = MB_EXCEPTION_ILLEGAL_DATA_ADDRESS; // TODO?
                goto exception;
            }
            pb_u8(&pb, count*2);
            while (count-- > 0) {
#ifdef MB_SUPPORT_FC03
                if (fcx == FC03_READ_HOLDING_REGISTERS) {
                    exc = ms->readHolding(ms, ref++, &value);
                }
#endif
#ifdef MB_SUPPORT_FC04
                if (fcx == FC04_READ_INPUT_REGISTERS) {
                    exc = ms->readInput(ms, ref++, &value);
                }
#endif
                if (exc != 0) {
                    goto exception;
                }
                pb_u16(&pb, value);
            }
            break;
#endif

#ifdef MB_SUPPORT_FC06
        case FC06_WRITE_SINGLE_REGISTER:
            if (!ms->writeHolding) {
                exc = MB_EXCEPTION_ILLEGAL_FUNCTION;
                goto exception;
            }
            ref = pp_u16(&pp);
            value = pp_u16(&pp);
            if (!pp_ok(&pp)) {
                return MB_ERR_NEEDMORE;
            }
            if (resp_capacity < 4 + overhead) {
                exc = MB_EXCEPTION_ILLEGAL_DATA_ADDRESS; // TODO?
                goto exception;
            }
            exc = ms->writeHolding(ms, ref, value);
            if (exc != 0) {
                goto exception;
            }
            pb_u16(&pb, ref);
            pb_u16(&pb, value);
            break;
#endif

#ifdef MB_SUPPORT_FC16
        case FC16_WRITE_MULTIPLE_REGISTERS:
            if (!ms->writeHolding) {
                exc = MB_EXCEPTION_ILLEGAL_FUNCTION;
                goto exception;
            }
            ref = pp_u16(&pp);
            count = pp_u16(&pp);
            pp_skip(&pp, 1); // this is always count * 2
            if (!pp_ok(&pp)) {
                return MB_ERR_NEEDMORE;
            }
            if (count > 255 || resp_capacity < 4 + overhead) {
                exc = MB_EXCEPTION_ILLEGAL_DATA_ADDRESS; // TODO?
                goto exception;
            }
            if (pp_remains(&pp) < count * 2 + (tcp ? 0 : 2)) {
                return MB_ERR_NEEDMORE;
            }
            pb_u16(&pb, ref);
            pb_u16(&pb, count);
            while (count-- > 0) {
                value = pp_u16(&pp);
                exc = ms->writeHolding(ms, ref++, value);
                if (exc != 0) {
                    goto exception;
                }
            }
            break;
#endif

#ifdef MB_SUPPORT_FC15
        case FC15_WRITE_MULTIPLE_COILS:
            if (!ms->writeCoil) {
                exc = MB_EXCEPTION_ILLEGAL_FUNCTION;
                goto exception;
            }
            ref = pp_u16(&pp);
            count = pp_u16(&pp);
            bytecount = pp_u8(&pp);
            if (!pp_ok(&pp)) {
                return MB_ERR_NEEDMORE;
            }
            if (count > 255 || resp_capacity < 4 + overhead) {
                exc = MB_EXCEPTION_ILLEGAL_DATA_ADDRESS; // TODO?
                goto exception;
            }
            if (pp_remains(&pp) < bytecount + (tcp ? 0 : 2)) {
                return MB_ERR_NEEDMORE;
            }
            pb_u16(&pb, ref);
            pb_u16(&pb, count);
            bitcnt = 8;
            scratch = 0;
            while (count-- > 0) {
                if (bitcnt == 8) {
                    scratch = pp_u8(&pp);
                    bitcnt = 0;
                }
                exc = ms->writeCoil(ms, ref++, (scratch >> bitcnt) & 1);
                bitcnt++;
                if (exc != 0) {
                    goto exception;
                }
            }
            break;
#endif

#ifdef MB_SUPPORT_FC22
        case FC22_MASK_WRITE_REGISTER:
            if (!ms->writeHolding || !ms->readHolding) {
                exc = MB_EXCEPTION_ILLEGAL_FUNCTION;
                goto exception;
            }
            ref = pp_u16(&pp);
            and_mask = pp_u16(&pp);
            or_mask = pp_u16(&pp);
            if (!pp_ok(&pp)) {
                return MB_ERR_NEEDMORE;
            }
            if (resp_capacity < 4 + overhead) {
                exc = MB_EXCEPTION_ILLEGAL_DATA_ADDRESS; // TODO?
                goto exception;
            }
            exc = ms->readHolding(ms, ref, &value);
            if (exc != 0) {
                goto exception;
            }

            value = (value & and_mask) | (or_mask & ~and_mask);
            exc = ms->writeHolding(ms, ref, value);
            if (exc != 0) {
                goto exception;
            }

            pb_u16(&pb, ref);
            pb_u16(&pb, and_mask);
            pb_u16(&pb, or_mask);
            break;
#endif

#ifdef MB_SUPPORT_FC23
        case FC23_READ_WRITE_MULTIPLE_REGISTERS:
            if (!ms->writeHolding || !ms->readHolding) {
                exc = MB_EXCEPTION_ILLEGAL_FUNCTION;
                goto exception;
            }
            // read
            ref = pp_u16(&pp);
            count = pp_u16(&pp);
            // write
            ref2 = pp_u16(&pp);
            count2 = pp_u16(&pp);
            pp_skip(&pp, 1); // qty of bytes
            if (!pp_ok(&pp)) {
                return MB_ERR_NEEDMORE;
            }
            if (count > 255 || count2 > 255 || resp_capacity < count * 2 + overhead) {
                exc = MB_EXCEPTION_ILLEGAL_DATA_ADDRESS; // TODO?
                goto exception;
            }
            if (pp_remains(&pp) < count2 * 2 + (tcp ? 0 : 2)) {
                return MB_ERR_NEEDMORE;
            }
            pb_u8(&pb, 2 * count);
            // First, write
            while (count2-- > 0) {
                value = pp_u16(&pp);
                exc = ms->writeHolding(ms, ref2++, value);
                if (exc != 0) {
                    goto exception;
                }
            }
            // Second, read
            while (count-- > 0) {
                exc = ms->readHolding(ms, ref++, &value);
                if (exc != 0) {
                    goto exception;
                }
                pb_u16(&pb, value);
            }
            break;
#endif

        default:
            exc = MB_EXCEPTION_ILLEGAL_FUNCTION;
            goto exception;
    }
    goto checksum;

    exception:;
    *resp_fc_mark |= 0x80;
    pb_restore(&pb, resp_pld_start_mark);
    pb_u8(&pb, (uint8_t) exc);
    goto checksum;

    checksum:
    numbytes = pb_length(&pb);
    if (tcp) {
        pb_restore(&pb, resp_len_mark);
        pb_u16(&pb, numbytes - 6);
        *resp_size = numbytes;
    } else {
        uint8_t *m = pb.start;
        uint16_t crc = crc16_init();
        *resp_size = numbytes + 2;
        while (numbytes > 0) {
            crc = crc16_update(crc, *m++);
            numbytes--;
        }
        pb.bigendian = 0; // CRC is sent as little endian?
        pb_u16(&pb, crc);
    }

    if (!pb_ok(&pb)) {
        return MB_ERROR;
    }

    if (ms->endOfAccess) {
        ms->endOfAccess(ms);
    }
    return MB_OK;
}