#include "setp_tms2812_boot.h"

#include <string.h>

#define SHA256_BLOCK_SIZE 64U

typedef struct {
    uint32_t state[8];
    uint64_t bit_count;
    uint8_t block[SHA256_BLOCK_SIZE];
    uint16_t block_length;
} sha256_ctx_t;

static uint32_t rotr32(uint32_t value, uint8_t bits)
{
    return (value >> bits) | (value << (32U - bits));
}

static void sha256_transform(sha256_ctx_t *ctx, const uint8_t *block)
{
    static const uint32_t k[64] = {
        0x428A2F98UL, 0x71374491UL, 0xB5C0FBCFUL, 0xE9B5DBA5UL,
        0x3956C25BUL, 0x59F111F1UL, 0x923F82A4UL, 0xAB1C5ED5UL,
        0xD807AA98UL, 0x12835B01UL, 0x243185BEUL, 0x550C7DC3UL,
        0x72BE5D74UL, 0x80DEB1FEUL, 0x9BDC06A7UL, 0xC19BF174UL,
        0xE49B69C1UL, 0xEFBE4786UL, 0x0FC19DC6UL, 0x240CA1CCUL,
        0x2DE92C6FUL, 0x4A7484AAUL, 0x5CB0A9DCUL, 0x76F988DAUL,
        0x983E5152UL, 0xA831C66DUL, 0xB00327C8UL, 0xBF597FC7UL,
        0xC6E00BF3UL, 0xD5A79147UL, 0x06CA6351UL, 0x14292967UL,
        0x27B70A85UL, 0x2E1B2138UL, 0x4D2C6DFCUL, 0x53380D13UL,
        0x650A7354UL, 0x766A0ABBUL, 0x81C2C92EUL, 0x92722C85UL,
        0xA2BFE8A1UL, 0xA81A664BUL, 0xC24B8B70UL, 0xC76C51A3UL,
        0xD192E819UL, 0xD6990624UL, 0xF40E3585UL, 0x106AA070UL,
        0x19A4C116UL, 0x1E376C08UL, 0x2748774CUL, 0x34B0BCB5UL,
        0x391C0CB3UL, 0x4ED8AA4AUL, 0x5B9CCA4FUL, 0x682E6FF3UL,
        0x748F82EEUL, 0x78A5636FUL, 0x84C87814UL, 0x8CC70208UL,
        0x90BEFFFAUL, 0xA4506CEBUL, 0xBEF9A3F7UL, 0xC67178F2UL
    };
    uint32_t w[64];
    uint32_t a, b, c, d, e, f, g, h, s0, s1, ch, maj, temp1, temp2;
    uint16_t i;
    for (i = 0U; i < 16U; i++) {
        uint16_t p = (uint16_t)(i * 4U);
        w[i] = ((uint32_t)block[p] << 24U)
             | ((uint32_t)block[p + 1U] << 16U)
             | ((uint32_t)block[p + 2U] << 8U)
             | (uint32_t)block[p + 3U];
    }
    for (i = 16U; i < 64U; i++) {
        s0 = rotr32(w[i - 15U], 7U) ^ rotr32(w[i - 15U], 18U)
           ^ (w[i - 15U] >> 3U);
        s1 = rotr32(w[i - 2U], 17U) ^ rotr32(w[i - 2U], 19U)
           ^ (w[i - 2U] >> 10U);
        w[i] = w[i - 16U] + s0 + w[i - 7U] + s1;
    }
    a = ctx->state[0]; b = ctx->state[1]; c = ctx->state[2]; d = ctx->state[3];
    e = ctx->state[4]; f = ctx->state[5]; g = ctx->state[6]; h = ctx->state[7];
    for (i = 0U; i < 64U; i++) {
        s1 = rotr32(e, 6U) ^ rotr32(e, 11U) ^ rotr32(e, 25U);
        ch = (e & f) ^ ((~e) & g);
        temp1 = h + s1 + ch + k[i] + w[i];
        s0 = rotr32(a, 2U) ^ rotr32(a, 13U) ^ rotr32(a, 22U);
        maj = (a & b) ^ (a & c) ^ (b & c);
        temp2 = s0 + maj;
        h = g; g = f; f = e; e = d + temp1;
        d = c; c = b; b = a; a = temp1 + temp2;
    }
    ctx->state[0] += a; ctx->state[1] += b; ctx->state[2] += c; ctx->state[3] += d;
    ctx->state[4] += e; ctx->state[5] += f; ctx->state[6] += g; ctx->state[7] += h;
}

static void sha256_init(sha256_ctx_t *ctx)
{
    static const uint32_t initial[8] = {
        0x6A09E667UL, 0xBB67AE85UL, 0x3C6EF372UL, 0xA54FF53AUL,
        0x510E527FUL, 0x9B05688CUL, 0x1F83D9ABUL, 0x5BE0CD19UL
    };
    (void)memcpy(ctx->state, initial, sizeof(initial));
    ctx->bit_count = 0U;
    ctx->block_length = 0U;
}

static void sha256_update(sha256_ctx_t *ctx, const uint8_t *data, uint16_t length)
{
    uint16_t i;
    for (i = 0U; i < length; i++) {
        ctx->block[ctx->block_length++] = data[i];
        ctx->bit_count += 8U;
        if (ctx->block_length == SHA256_BLOCK_SIZE) {
            sha256_transform(ctx, ctx->block);
            ctx->block_length = 0U;
        }
    }
}

static void sha256_finish(sha256_ctx_t *ctx, uint8_t digest[SETP_SHA256_SIZE])
{
    uint16_t i;
    uint64_t bits = ctx->bit_count;
    ctx->block[ctx->block_length++] = 0x80U;
    if (ctx->block_length > 56U) {
        while (ctx->block_length < SHA256_BLOCK_SIZE) ctx->block[ctx->block_length++] = 0U;
        sha256_transform(ctx, ctx->block);
        ctx->block_length = 0U;
    }
    while (ctx->block_length < 56U) ctx->block[ctx->block_length++] = 0U;
    for (i = 0U; i < 8U; i++) {
        ctx->block[63U - i] = (uint8_t)(bits & 0xFFU);
        bits >>= 8U;
    }
    sha256_transform(ctx, ctx->block);
    for (i = 0U; i < 8U; i++) {
        digest[i * 4U] = (uint8_t)(ctx->state[i] >> 24U);
        digest[i * 4U + 1U] = (uint8_t)(ctx->state[i] >> 16U);
        digest[i * 4U + 2U] = (uint8_t)(ctx->state[i] >> 8U);
        digest[i * 4U + 3U] = (uint8_t)ctx->state[i];
    }
}

static void boot_status(const setp_tms2812_boot_t *boot, setp_fw_status_t *status)
{
    status->state = boot->state;
    status->active_slot = boot->config.active_slot;
    status->max_block_size = boot->config.max_block_size;
    status->next_offset = boot->next_offset;
    status->image_size = boot->manifest.image_size;
    status->last_error = boot->last_error;
    status->flags = 0U;
}

static bool boot_send_response(setp_tms2812_boot_t *boot,
                               const setp_frame_t *request,
                               const setp_can_id_t *request_id,
                               uint16_t status,
                               const uint8_t *body, uint16_t body_length)
{
    setp_frame_t response;
    setp_can_id_t response_id;
    size_t packet_length;
    if ((uint32_t)body_length + 2U > sizeof(boot->response_payload)) return false;
    setp_put_u16(boot->response_payload, status);
    if (body_length != 0U) (void)memcpy(&boot->response_payload[2], body, body_length);
    response.flags = SETP_FLAG_RESPONSE;
    if (status != SETP_STATUS_OK) response.flags |= SETP_FLAG_ERROR;
    if ((request->flags & SETP_FLAG_PRIORITY) != 0U) response.flags |= SETP_FLAG_PRIORITY;
    response.message_type = request->message_type;
    response.source = boot->config.node_id;
    response.destination = request->source;
    response.sequence = request->sequence;
    response.payload_length = (uint16_t)(body_length + 2U);
    response.payload = boot->response_payload;
    packet_length = setp_frame_encode(&response, boot->response_packet,
                                      sizeof(boot->response_packet));
    if (packet_length == 0U) return false;
    response_id.destination = request_id->source;
    response_id.source = boot->config.node_id;
    response_id.priority = (response.flags & SETP_FLAG_PRIORITY) != 0U ? 1U : 0U;
    response_id.channel = request_id->channel;
    return setp_can_segment(boot->response_packet, (uint16_t)packet_length,
                            setp_can_id_pack(&response_id), boot->port.send_can,
                            boot->port_user);
}

static uint16_t boot_verify_image(setp_tms2812_boot_t *boot)
{
    uint8_t data[SETP_TMS2812_MAX_BLOCK_SIZE];
    uint8_t digest[SETP_SHA256_SIZE];
    sha256_ctx_t sha;
    uint32_t crc = 0xFFFFFFFFUL;
    uint32_t offset = 0U;
    uint16_t i;
    sha256_init(&sha);
    while (offset < boot->manifest.image_size) {
        uint32_t remaining = boot->manifest.image_size - offset;
        uint16_t length = remaining > sizeof(data) ? (uint16_t)sizeof(data) : (uint16_t)remaining;
        if (!boot->port.read_image(boot->port_user, offset, data, length)) {
            return SETP_STATUS_INTERNAL;
        }
        sha256_update(&sha, data, length);
        for (i = 0U; i < length; i++) {
            uint8_t bit;
            crc ^= data[i];
            for (bit = 0U; bit < 8U; bit++)
                crc = (crc >> 1U) ^ (((crc & 1U) != 0U) ? 0xEDB88320UL : 0U);
        }
        offset += length;
    }
    sha256_finish(&sha, digest);
    crc ^= 0xFFFFFFFFUL;
    if ((crc != boot->manifest.image_crc32)
            || (memcmp(digest, boot->manifest.sha256, SETP_SHA256_SIZE) != 0)) {
        return SETP_STATUS_VERIFY_FAILED;
    }
    return SETP_STATUS_OK;
}

static uint16_t boot_fw_begin(setp_tms2812_boot_t *boot, const setp_frame_t *request)
{
    setp_fw_begin_t value;
    bool same_manifest;
    if (!setp_fw_begin_decode(request->payload, request->payload_length, &value))
        return SETP_STATUS_INVALID_LENGTH;
    if ((value.image_size == 0U) || (value.image_size > boot->config.max_image_size)
            || (value.slot != boot->config.active_slot)
            || ((value.base_address != 0U)
                && (value.base_address != boot->config.app_base_address))
            || (value.block_size == 0U)
            || (value.block_size > boot->config.max_block_size))
        return SETP_STATUS_INVALID_ARGUMENT;
    if (((value.flags & SETP_FW_FLAG_SIGNED) != 0U) || boot->config.require_signature) {
        if ((boot->port.authorize == NULL)
                || !boot->port.authorize(boot->port_user, &value))
            return SETP_STATUS_AUTH_FAILED;
    }
    same_manifest = boot->state == SETP_FW_RECEIVING
        && boot->manifest.image_size == value.image_size
        && boot->manifest.image_crc32 == value.image_crc32
        && boot->manifest.image_version == value.image_version
        && memcmp(boot->manifest.sha256, value.sha256, SETP_SHA256_SIZE) == 0;
    if (same_manifest && ((value.flags & SETP_FW_FLAG_RESUME) != 0U)) return SETP_STATUS_OK;
    if (!boot->port.erase_image(boot->port_user, value.image_size))
        return SETP_STATUS_INTERNAL;
    boot->manifest = value;
    boot->manifest.signature = NULL;
    boot->manifest.signature_length = 0U;
    boot->next_offset = 0U;
    boot->state = SETP_FW_RECEIVING;
    return SETP_STATUS_OK;
}

static uint16_t boot_fw_data(setp_tms2812_boot_t *boot, const setp_frame_t *request)
{
    setp_fw_data_t value;
    uint8_t current[SETP_TMS2812_MAX_BLOCK_SIZE];
    if (boot->state != SETP_FW_RECEIVING) return SETP_STATUS_WRONG_STATE;
    if (!setp_fw_data_decode(request->payload, request->payload_length, &value))
        return SETP_STATUS_CRC;
    if ((value.data_length > boot->manifest.block_size)
            || (value.data_length > boot->config.max_block_size)
            || (value.offset > boot->manifest.image_size)
            || ((uint32_t)value.data_length > boot->manifest.image_size - value.offset))
        return SETP_STATUS_INVALID_ARGUMENT;
    if (value.offset < boot->next_offset) {
        if ((value.offset + value.data_length > boot->next_offset)
                || !boot->port.read_image(boot->port_user, value.offset,
                                          current, value.data_length)
                || memcmp(current, value.data, value.data_length) != 0)
            return SETP_STATUS_SEQUENCE;
        return SETP_STATUS_OK;
    }
    if (value.offset != boot->next_offset) return SETP_STATUS_SEQUENCE;
    if (!boot->port.write_image(boot->port_user, value.offset,
                                value.data, value.data_length))
        return SETP_STATUS_INTERNAL;
    boot->next_offset += value.data_length;
    return SETP_STATUS_OK;
}

static uint16_t boot_fw_end(setp_tms2812_boot_t *boot, const setp_frame_t *request)
{
    setp_fw_end_t value;
    uint16_t status;
    if (boot->state != SETP_FW_RECEIVING) return SETP_STATUS_WRONG_STATE;
    if (!setp_fw_end_decode(request->payload, request->payload_length, &value))
        return SETP_STATUS_INVALID_LENGTH;
    if ((boot->next_offset != boot->manifest.image_size)
            || (value.image_size != boot->manifest.image_size)
            || (value.image_crc32 != boot->manifest.image_crc32)
            || (memcmp(value.sha256, boot->manifest.sha256, SETP_SHA256_SIZE) != 0))
        return SETP_STATUS_VERIFY_FAILED;
    boot->state = SETP_FW_VERIFYING;
    status = boot_verify_image(boot);
    boot->state = status == SETP_STATUS_OK ? SETP_FW_READY : SETP_FW_FAILED;
    return status;
}

static bool boot_dispatch(setp_tms2812_boot_t *boot, const setp_frame_t *request,
                          const setp_can_id_t *request_id, uint32_t now_ms)
{
    uint8_t body[SETP_TMS2812_RESPONSE_PAYLOAD_SIZE - 2U];
    uint16_t body_length = 0U;
    uint16_t status = SETP_STATUS_OK;
    uint8_t reboot = 0U;
    setp_fw_status_t fw_status;
    if ((request->flags & (SETP_FLAG_RESPONSE | SETP_FLAG_EVENT)) != 0U) return false;
    if ((request->source > 0xFFU) || (request->destination != boot->config.node_id)) return false;
    switch (request->message_type) {
    case SETP_MSG_PING:
        if (request->payload_length != 0U) status = SETP_STATUS_INVALID_LENGTH;
        else { setp_put_u32(body, now_ms); body_length = 4U; }
        break;
    case SETP_MSG_DEVICE_INFO:
        if (request->payload_length != 0U) status = SETP_STATUS_INVALID_LENGTH;
        else {
            setp_device_info_t info;
            info.schema_version = SETP_DEVICE_INFO_SCHEMA_VERSION;
            info.device_class = boot->config.device_class;
            info.hardware_version = boot->config.hardware_version;
            info.firmware_version = boot->config.firmware_version;
            info.dictionary_version = boot->config.dictionary_version;
            info.serial_number = boot->config.serial_number;
            info.model_length = boot->config.model_length;
            info.model = boot->config.model;
            body_length = (uint16_t)setp_device_info_encode(&info, body, sizeof(body));
            if (body_length == 0U) status = SETP_STATUS_INTERNAL;
        }
        break;
    case SETP_MSG_CAPABILITIES:
        if (request->payload_length != 0U) status = SETP_STATUS_INVALID_LENGTH;
        else {
            setp_capabilities_t caps;
            caps.schema_version = SETP_CAPABILITIES_SCHEMA_VERSION;
            caps.max_payload = (uint16_t)(SETP_FW_DATA_HEADER_SIZE + boot->config.max_block_size);
            caps.interface_mask = SETP_IFACE_MASK(SETP_IFACE_CAN);
            caps.feature_flags = SETP_FEATURE_FIRMWARE;
            caps.max_read_items = 0U; caps.max_write_items = 0U;
            caps.max_subscriptions = 0U; caps.max_publish_items = 0U;
            body_length = (uint16_t)setp_capabilities_encode(&caps, body, sizeof(body));
            if (body_length == 0U) status = SETP_STATUS_INTERNAL;
        }
        break;
    case SETP_MSG_FW_BEGIN:
        status = boot_fw_begin(boot, request);
        setp_put_u32(body, boot->next_offset); body_length = 4U;
        break;
    case SETP_MSG_FW_DATA:
        status = boot_fw_data(boot, request);
        setp_put_u32(body, boot->next_offset); body_length = 4U;
        break;
    case SETP_MSG_FW_END:
        status = boot_fw_end(boot, request);
        setp_put_u32(body, boot->next_offset); body_length = 4U;
        break;
    case SETP_MSG_FW_ABORT:
        if (request->payload_length != 0U) status = SETP_STATUS_INVALID_LENGTH;
        else setp_tms2812_boot_abort(boot);
        setp_put_u32(body, boot->next_offset); body_length = 4U;
        break;
    case SETP_MSG_FW_STATUS:
        if (request->payload_length != 0U) status = SETP_STATUS_INVALID_LENGTH;
        else {
            boot_status(boot, &fw_status);
            body_length = (uint16_t)setp_fw_status_encode(&fw_status, body, sizeof(body));
        }
        break;
    case SETP_MSG_FW_ACTIVATE:
        if (request->payload_length != 0U) status = SETP_STATUS_INVALID_LENGTH;
        else if (boot->state != SETP_FW_READY) status = SETP_STATUS_WRONG_STATE;
        else { boot->state = SETP_FW_ACTIVE; reboot = 1U; }
        setp_put_u32(body, boot->next_offset); body_length = 4U;
        break;
    default:
        status = SETP_STATUS_UNSUPPORTED;
        break;
    }
    boot->last_error = status == SETP_STATUS_OK ? 0U : status;
    if (!boot_send_response(boot, request, request_id, status, body, body_length)) return false;
    if (reboot != 0U) boot->port.reboot(boot->port_user);
    return true;
}

bool setp_tms2812_boot_init(setp_tms2812_boot_t *boot,
                            const setp_tms2812_boot_config_t *config,
                            const setp_tms2812_boot_port_t *port,
                            void *port_user)
{
    if ((boot == NULL) || (config == NULL) || (port == NULL)
            || (config->model == NULL) || (config->model_length > SETP_DEVICE_MODEL_MAX)
            || (config->max_image_size == 0U) || (config->max_block_size == 0U)
            || (config->max_block_size > SETP_TMS2812_MAX_BLOCK_SIZE)
            || (port->send_can == NULL) || (port->erase_image == NULL)
            || (port->write_image == NULL) || (port->read_image == NULL)
            || (port->reboot == NULL)) return false;
    (void)memset(boot, 0, sizeof(*boot));
    boot->config = *config;
    boot->port = *port;
    boot->port_user = port_user;
    boot->state = SETP_FW_IDLE;
    setp_can_rx_init(&boot->rx);
    return true;
}

void setp_tms2812_boot_abort(setp_tms2812_boot_t *boot)
{
    if (boot == NULL) return;
    boot->state = SETP_FW_IDLE;
    boot->next_offset = 0U;
    boot->last_error = 0U;
    (void)memset(&boot->manifest, 0, sizeof(boot->manifest));
}

bool setp_tms2812_boot_process(setp_tms2812_boot_t *boot,
                               const setp_can_frame_t *frame,
                               uint32_t now_ms)
{
    setp_can_packet_t packet;
    setp_can_rx_result_t result;
    setp_can_id_t can_id;
    setp_frame_t request;
    if ((boot == NULL) || (frame == NULL)) return false;
    result = setp_can_rx_feed(&boot->rx, frame, now_ms, &packet);
    if (result != SETP_CAN_RX_COMPLETE) return result == SETP_CAN_RX_NONE;
    if (!setp_can_id_unpack(packet.can_id, &can_id)
            || (can_id.destination != boot->config.node_id)
            || !setp_frame_decode_datagram(packet.data, packet.length, &request)) return false;
    return boot_dispatch(boot, &request, &can_id, now_ms);
}