#include "modbus_rtu.h"

#include <string.h>

#include "esp_log.h"

#define TAG "MODBUS_RTU"

static uint16_t get_u16_be(const uint8_t *src)
{
    return (uint16_t)(((uint16_t)src[0] << 8) | (uint16_t)src[1]);
}

static void put_u16_be(uint8_t *dst, uint16_t value)
{
    dst[0] = (uint8_t)((value >> 8) & 0xFF);
    dst[1] = (uint8_t)(value & 0xFF);
}

uint16_t modbus_rtu_crc16(const uint8_t *data, size_t len)
{
    uint16_t crc = 0xFFFF;
    size_t i;
    int bit;

    for (i = 0; i < len; i++)
    {
        crc ^= data[i];

        for (bit = 0; bit < 8; bit++)
        {
            if (crc & 0x0001U)
                crc = (uint16_t)((crc >> 1) ^ 0xA001U);
            else
                crc >>= 1;
        }
    }

    return crc;
}

static size_t build_exception_response(uint8_t unit_id,
                                       uint8_t function_code,
                                       uint8_t exception_code,
                                       uint8_t *resp,
                                       size_t resp_size)
{
    uint16_t crc;

    if (resp == NULL || resp_size < 5U)
        return 0U;

    resp[0] = unit_id;
    resp[1] = (uint8_t)(function_code | 0x80U);
    resp[2] = exception_code;

    crc = modbus_rtu_crc16(resp, 3U);
    resp[3] = (uint8_t)(crc & 0xFFU);
    resp[4] = (uint8_t)((crc >> 8) & 0xFFU);

    return 5U;
}

void modbus_rtu_init(modbus_rtu_ctx_t *ctx, uart_port_t uart_num, uint32_t baud_rate)
{
    if (ctx == NULL)
        return;

    memset(ctx, 0, sizeof(*ctx));
    ctx->uart_num = uart_num;
    ctx->baud_rate = baud_rate;
}

bool modbus_rtu_add_device(modbus_rtu_ctx_t *ctx,
                           uint8_t unit_id,
                           modbus_memory_t *mem,
                           const char *name)
{
    uint8_t i;
    modbus_rtu_device_t *dev;

    if (ctx == NULL || mem == NULL)
        return false;

    if (unit_id == 0U || unit_id > 247U)
        return false;

    if (ctx->unit_count >= MODBUS_RTU_MAX_DEVICES)
        return false;

    for (i = 0; i < ctx->unit_count; i++)
    {
        if (ctx->devices[i].unit_id == unit_id)
            return false;
    }

    dev = &ctx->devices[ctx->unit_count++];
    dev->unit_id = unit_id;
    dev->mem = mem;
    dev->name = name;

    ESP_LOGI(TAG, "Registered RTU device unit_id=%u name=%s",
             unit_id, (name != NULL) ? name : "(unnamed)");

    return true;
}

modbus_rtu_device_t *modbus_rtu_find_device(modbus_rtu_ctx_t *ctx, uint8_t unit_id)
{
    uint8_t i;

    if (ctx == NULL)
        return NULL;

    for (i = 0; i < ctx->unit_count; i++)
    {
        if (ctx->devices[i].unit_id == unit_id)
            return &ctx->devices[i];
    }

    return NULL;
}

int modbus_rtu_read_frame(modbus_rtu_ctx_t *ctx,
                          uint8_t *buf,
                          size_t buf_size,
                          TickType_t first_byte_timeout)
{
    int total;
    int n;

    if (ctx == NULL || buf == NULL || buf_size == 0U)
        return -1;

    total = uart_read_bytes(ctx->uart_num, buf, 1, first_byte_timeout);
    if (total <= 0)
        return 0;

    while ((size_t)total < buf_size)
    {
        n = uart_read_bytes(ctx->uart_num,
                            buf + total,
                            (uint32_t)(buf_size - (size_t)total),
                            pdMS_TO_TICKS(10));
        if (n <= 0)
            break;

        total += n;
    }

    return total;
}

static bool handle_read_bits(uint8_t unit_id,
                             uint8_t function_code,
                             modbus_memory_t *mem,
                             const uint8_t *req,
                             size_t req_len,
                             uint8_t *resp,
                             size_t resp_size,
                             size_t *resp_len)
{
    uint16_t start_addr;
    uint16_t quantity;
    uint8_t byte_count;
    uint16_t i;
    uint16_t crc;

    if (req == NULL || mem == NULL || resp == NULL || resp_len == NULL || req_len < 8U)
        return false;

    start_addr = get_u16_be(&req[2]);
    quantity = get_u16_be(&req[4]);

    if (quantity == 0U || quantity > 2000U)
    {
        *resp_len = build_exception_response(unit_id, function_code, 0x03, resp, resp_size);
        return true;
    }

    byte_count = (uint8_t)((quantity + 7U) / 8U);

    if (resp_size < (size_t)(3U + byte_count + 2U))
        return false;

    resp[0] = unit_id;
    resp[1] = function_code;
    resp[2] = byte_count;
    memset(&resp[3], 0, byte_count);

    for (i = 0; i < quantity; i++)
    {
        uint8_t bit = 0U;
        bool ok;

        if (function_code == 0x01U)
            ok = modbus_memory_read_coil(mem, (uint16_t)(start_addr + i), &bit);
        else
            ok = modbus_memory_read_discrete_input(mem, (uint16_t)(start_addr + i), &bit);

        if (!ok)
        {
            *resp_len = build_exception_response(unit_id, function_code, 0x02, resp, resp_size);
            return true;
        }

        if (bit != 0U)
            resp[3 + (i / 8U)] |= (uint8_t)(1U << (i % 8U));
    }

    *resp_len = (size_t)(3U + byte_count + 2U);
    crc = modbus_rtu_crc16(resp, *resp_len - 2U);
    resp[*resp_len - 2U] = (uint8_t)(crc & 0xFFU);
    resp[*resp_len - 1U] = (uint8_t)((crc >> 8) & 0xFFU);

    return true;
}

static bool handle_read_registers(uint8_t unit_id,
                                  uint8_t function_code,
                                  modbus_memory_t *mem,
                                  const uint8_t *req,
                                  size_t req_len,
                                  uint8_t *resp,
                                  size_t resp_size,
                                  size_t *resp_len)
{
    uint16_t start_addr;
    uint16_t quantity;
    uint8_t byte_count;
    uint16_t i;
    uint16_t crc;

    if (req == NULL || mem == NULL || resp == NULL || resp_len == NULL || req_len < 8U)
        return false;

    start_addr = get_u16_be(&req[2]);
    quantity = get_u16_be(&req[4]);

    if (quantity == 0U || quantity > 125U)
    {
        *resp_len = build_exception_response(unit_id, function_code, 0x03, resp, resp_size);
        return true;
    }

    byte_count = (uint8_t)(quantity * 2U);

    if (resp_size < (size_t)(3U + byte_count + 2U))
        return false;

    resp[0] = unit_id;
    resp[1] = function_code;
    resp[2] = byte_count;

    for (i = 0; i < quantity; i++)
    {
        uint16_t value = 0U;
        bool ok;

        if (function_code == 0x03U)
            ok = modbus_memory_read_holding_register(mem, (uint16_t)(start_addr + i), &value);
        else
            ok = modbus_memory_read_input_register(mem, (uint16_t)(start_addr + i), &value);

        if (!ok)
        {
            *resp_len = build_exception_response(unit_id, function_code, 0x02, resp, resp_size);
            return true;
        }

        put_u16_be(&resp[3U + (i * 2U)], value);
    }

    *resp_len = (size_t)(3U + byte_count + 2U);
    crc = modbus_rtu_crc16(resp, *resp_len - 2U);
    resp[*resp_len - 2U] = (uint8_t)(crc & 0xFFU);
    resp[*resp_len - 1U] = (uint8_t)((crc >> 8) & 0xFFU);

    return true;
}

static bool handle_write_single_coil(uint8_t unit_id,
                                     modbus_memory_t *mem,
                                     const uint8_t *req,
                                     size_t req_len,
                                     uint8_t *resp,
                                     size_t resp_size,
                                     size_t *resp_len,
                                     bool is_broadcast)
{
    uint16_t addr;
    uint16_t raw;
    uint16_t crc;

    if (req == NULL || mem == NULL || resp_len == NULL || req_len < 8U)
        return false;

    addr = get_u16_be(&req[2]);
    raw = get_u16_be(&req[4]);

    if (raw != 0x0000U && raw != 0xFF00U)
    {
        if (!is_broadcast)
            *resp_len = build_exception_response(unit_id, 0x05U, 0x03U, resp, resp_size);
        else
            *resp_len = 0U;
        return true;
    }

    if (!modbus_memory_write_coil(mem, addr, (raw == 0xFF00U) ? 1U : 0U))
    {
        if (!is_broadcast)
            *resp_len = build_exception_response(unit_id, 0x05U, 0x02U, resp, resp_size);
        else
            *resp_len = 0U;
        return true;
    }

    if (is_broadcast)
    {
        *resp_len = 0U;
        return true;
    }

    if (resp == NULL || resp_size < 8U)
        return false;

    memcpy(resp, req, 6U);
    resp[0] = unit_id;

    crc = modbus_rtu_crc16(resp, 6U);
    resp[6] = (uint8_t)(crc & 0xFFU);
    resp[7] = (uint8_t)((crc >> 8) & 0xFFU);

    *resp_len = 8U;
    return true;
}

static bool handle_write_single_register(uint8_t unit_id,
                                         modbus_memory_t *mem,
                                         const uint8_t *req,
                                         size_t req_len,
                                         uint8_t *resp,
                                         size_t resp_size,
                                         size_t *resp_len,
                                         bool is_broadcast)
{
    uint16_t addr;
    uint16_t value;
    uint16_t crc;

    if (req == NULL || mem == NULL || resp_len == NULL || req_len < 8U)
        return false;

    addr = get_u16_be(&req[2]);
    value = get_u16_be(&req[4]);

    if (!modbus_memory_write_holding_register(mem, addr, value))
    {
        if (!is_broadcast)
            *resp_len = build_exception_response(unit_id, 0x06U, 0x02U, resp, resp_size);
        else
            *resp_len = 0U;
        return true;
    }

    if (is_broadcast)
    {
        *resp_len = 0U;
        return true;
    }

    if (resp == NULL || resp_size < 8U)
        return false;

    memcpy(resp, req, 6U);
    resp[0] = unit_id;

    crc = modbus_rtu_crc16(resp, 6U);
    resp[6] = (uint8_t)(crc & 0xFFU);
    resp[7] = (uint8_t)((crc >> 8) & 0xFFU);

    *resp_len = 8U;
    return true;
}

static bool handle_write_multiple_coils(uint8_t unit_id,
                                        modbus_memory_t *mem,
                                        const uint8_t *req,
                                        size_t req_len,
                                        uint8_t *resp,
                                        size_t resp_size,
                                        size_t *resp_len,
                                        bool is_broadcast)
{
    uint16_t start_addr;
    uint16_t quantity;
    uint8_t byte_count;
    uint16_t i;
    uint8_t values[1968];
    uint16_t crc;

    if (req == NULL || mem == NULL || resp_len == NULL || req_len < 9U)
        return false;

    start_addr = get_u16_be(&req[2]);
    quantity = get_u16_be(&req[4]);
    byte_count = req[6];

    if (quantity == 0U || quantity > 0x07B0U)
    {
        if (!is_broadcast)
            *resp_len = build_exception_response(unit_id, 0x0FU, 0x03U, resp, resp_size);
        else
            *resp_len = 0U;
        return true;
    }

    if (byte_count != (uint8_t)((quantity + 7U) / 8U))
    {
        if (!is_broadcast)
            *resp_len = build_exception_response(unit_id, 0x0FU, 0x03U, resp, resp_size);
        else
            *resp_len = 0U;
        return true;
    }

    if (req_len < (size_t)(7U + byte_count + 2U))
        return false;

    for (i = 0; i < quantity; i++)
        values[i] = (uint8_t)((req[7U + (i / 8U)] >> (i % 8U)) & 0x01U);

    if (!modbus_memory_write_coils(mem, start_addr, quantity, values))
    {
        if (!is_broadcast)
            *resp_len = build_exception_response(unit_id, 0x0FU, 0x02U, resp, resp_size);
        else
            *resp_len = 0U;
        return true;
    }

    if (is_broadcast)
    {
        *resp_len = 0U;
        return true;
    }

    if (resp == NULL || resp_size < 8U)
        return false;

    resp[0] = unit_id;
    resp[1] = 0x0FU;
    resp[2] = req[2];
    resp[3] = req[3];
    resp[4] = req[4];
    resp[5] = req[5];

    crc = modbus_rtu_crc16(resp, 6U);
    resp[6] = (uint8_t)(crc & 0xFFU);
    resp[7] = (uint8_t)((crc >> 8) & 0xFFU);

    *resp_len = 8U;
    return true;
}

static bool handle_write_multiple_registers(uint8_t unit_id,
                                            modbus_memory_t *mem,
                                            const uint8_t *req,
                                            size_t req_len,
                                            uint8_t *resp,
                                            size_t resp_size,
                                            size_t *resp_len,
                                            bool is_broadcast)
{
    uint16_t start_addr;
    uint16_t quantity;
    uint8_t byte_count;
    uint16_t values[123];
    uint16_t i;
    uint16_t crc;

    if (req == NULL || mem == NULL || resp_len == NULL || req_len < 9U)
        return false;

    start_addr = get_u16_be(&req[2]);
    quantity = get_u16_be(&req[4]);
    byte_count = req[6];

    if (quantity == 0U || quantity > 123U)
    {
        if (!is_broadcast)
            *resp_len = build_exception_response(unit_id, 0x10U, 0x03U, resp, resp_size);
        else
            *resp_len = 0U;
        return true;
    }

    if (byte_count != (uint8_t)(quantity * 2U))
    {
        if (!is_broadcast)
            *resp_len = build_exception_response(unit_id, 0x10U, 0x03U, resp, resp_size);
        else
            *resp_len = 0U;
        return true;
    }

    if (req_len < (size_t)(7U + byte_count + 2U))
        return false;

    for (i = 0; i < quantity; i++)
        values[i] = get_u16_be(&req[7U + (i * 2U)]);

    if (!modbus_memory_write_holding_registers(mem, start_addr, quantity, values))
    {
        if (!is_broadcast)
            *resp_len = build_exception_response(unit_id, 0x10U, 0x02U, resp, resp_size);
        else
            *resp_len = 0U;
        return true;
    }

    if (is_broadcast)
    {
        *resp_len = 0U;
        return true;
    }

    if (resp == NULL || resp_size < 8U)
        return false;

    resp[0] = unit_id;
    resp[1] = 0x10U;
    resp[2] = req[2];
    resp[3] = req[3];
    resp[4] = req[4];
    resp[5] = req[5];

    crc = modbus_rtu_crc16(resp, 6U);
    resp[6] = (uint8_t)(crc & 0xFFU);
    resp[7] = (uint8_t)((crc >> 8) & 0xFFU);

    *resp_len = 8U;
    return true;
}

bool modbus_rtu_process_request(modbus_rtu_ctx_t *ctx,
                                const uint8_t *req,
                                size_t req_len,
                                uint8_t *resp,
                                size_t resp_size,
                                size_t *resp_len)
{
    uint8_t unit_id;
    uint8_t function_code;
    bool is_broadcast;
    uint16_t rx_crc;
    uint16_t calc_crc;
    modbus_rtu_device_t *dev;
    uint8_t i;

    if (ctx == NULL || req == NULL || resp == NULL || resp_len == NULL)
        return false;

    *resp_len = 0U;

    if (req_len < 4U)
        return false;

    unit_id = req[0];
    function_code = req[1];
    is_broadcast = (unit_id == 0U);

    rx_crc = (uint16_t)req[req_len - 2U] |
             (uint16_t)((uint16_t)req[req_len - 1U] << 8);

    calc_crc = modbus_rtu_crc16(req, req_len - 2U);
    if (rx_crc != calc_crc)
    {
        ESP_LOGW(TAG, "CRC mismatch rx=%04X calc=%04X", rx_crc, calc_crc);
        return false;
    }

    if (is_broadcast)
    {
        bool handled = false;

        for (i = 0; i < ctx->unit_count; i++)
        {
            modbus_memory_t *mem = ctx->devices[i].mem;
            size_t ignored_len = 0U;

            if (mem == NULL)
                continue;

            switch (function_code)
            {
                case 0x05U:
                    handled |= handle_write_single_coil(ctx->devices[i].unit_id,
                                                        mem,
                                                        req, req_len,
                                                        resp, resp_size,
                                                        &ignored_len,
                                                        true);
                    break;

                case 0x06U:
                    handled |= handle_write_single_register(ctx->devices[i].unit_id,
                                                            mem,
                                                            req, req_len,
                                                            resp, resp_size,
                                                            &ignored_len,
                                                            true);
                    break;

                case 0x0FU:
                    handled |= handle_write_multiple_coils(ctx->devices[i].unit_id,
                                                           mem,
                                                           req, req_len,
                                                           resp, resp_size,
                                                           &ignored_len,
                                                           true);
                    break;

                case 0x10U:
                    handled |= handle_write_multiple_registers(ctx->devices[i].unit_id,
                                                               mem,
                                                               req, req_len,
                                                               resp, resp_size,
                                                               &ignored_len,
                                                               true);
                    break;

                default:
                    break;
            }
        }

        *resp_len = 0U;
        return handled;
    }

    dev = modbus_rtu_find_device(ctx, unit_id);
    if (dev == NULL || dev->mem == NULL)
        return false;

    switch (function_code)
    {
        case 0x01U:
        case 0x02U:
            return handle_read_bits(unit_id,
                                    function_code,
                                    dev->mem,
                                    req, req_len,
                                    resp, resp_size,
                                    resp_len);

        case 0x03U:
        case 0x04U:
            return handle_read_registers(unit_id,
                                         function_code,
                                         dev->mem,
                                         req, req_len,
                                         resp, resp_size,
                                         resp_len);

        case 0x05U:
            return handle_write_single_coil(unit_id,
                                            dev->mem,
                                            req, req_len,
                                            resp, resp_size,
                                            resp_len,
                                            false);

        case 0x06U:
            return handle_write_single_register(unit_id,
                                                dev->mem,
                                                req, req_len,
                                                resp, resp_size,
                                                resp_len,
                                                false);

        case 0x0FU:
            return handle_write_multiple_coils(unit_id,
                                               dev->mem,
                                               req, req_len,
                                               resp, resp_size,
                                               resp_len,
                                               false);

        case 0x10U:
            return handle_write_multiple_registers(unit_id,
                                                   dev->mem,
                                                   req, req_len,
                                                   resp, resp_size,
                                                   resp_len,
                                                   false);

        default:
            *resp_len = build_exception_response(unit_id, function_code, 0x01U, resp, resp_size);
            return true;
    }
}