#include "ai_wire.h"
#include <string.h>

#define AIW_VERSION 2u
#define AIW_TYPE_REQUEST 0x30u
#define AIW_TYPE_RESULT 0x31u
#define AIW_TYPE_UI_SESSION_OPEN 0x32u
#define AIW_TYPE_UI_COMMAND 0x33u
#define AIW_TYPE_UI_ACK 0x34u
#define AIW_MAGIC_0 0x4Eu
#define AIW_MAGIC_1 0x4Cu
#define AIW_MAGIC_2 0x41u
#define AIW_MAGIC_3 0x49u

uint16_t aiw_crc16_ccitt_false(const uint8_t *data, size_t length) {
    uint16_t crc = 0xFFFFu;
    if (!data && length != 0u) return 0u;
    for (size_t i = 0; i < length; ++i) {
        crc ^= (uint16_t)data[i] << 8;
        for (unsigned bit = 0; bit < 8u; ++bit) {
            crc = (crc & 0x8000u) ? (uint16_t)((crc << 1) ^ 0x1021u)
                                  : (uint16_t)(crc << 1);
        }
    }
    return crc;
}

static void put16(uint8_t *p, uint16_t value) {
    p[0] = (uint8_t)value;
    p[1] = (uint8_t)(value >> 8);
}

static void put32(uint8_t *p, uint32_t value) {
    for (unsigned i = 0; i < 4u; ++i) p[i] = (uint8_t)(value >> (8u * i));
}

static void put64(uint8_t *p, uint64_t value) {
    for (unsigned i = 0; i < 8u; ++i) p[i] = (uint8_t)(value >> (8u * i));
}

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

static uint32_t get32(const uint8_t *p) {
    uint32_t value = 0u;
    for (unsigned i = 0; i < 4u; ++i) value |= (uint32_t)p[i] << (8u * i);
    return value;
}

static uint64_t get64(const uint8_t *p) {
    uint64_t value = 0u;
    for (unsigned i = 0; i < 8u; ++i) value |= (uint64_t)p[i] << (8u * i);
    return value;
}

static void set_header(uint8_t *out, uint8_t type, uint16_t payload_len) {
    out[0] = AIW_MAGIC_0;
    out[1] = AIW_MAGIC_1;
    out[2] = AIW_MAGIC_2;
    out[3] = AIW_MAGIC_3;
    out[4] = AIW_VERSION;
    out[5] = type;
    put16(out + 6, payload_len);
}

static bool valid_ticket(const aiw_ticket *ticket) {
    return ticket && ticket->session_id != 0u && ticket->ticket_id != 0u &&
           ticket->calibration_id != 0u && ticket->channels != 0u &&
           ticket->channels <= AIW_MAX_CHANNELS && ticket->min_source_frame_seq != 0u;
}

static bool valid_result(const aiw_result *result) {
    if (!result || result->session_id == 0u || result->ticket_id == 0u ||
        result->calibration_id == 0u || result->source_frame_seq == 0u ||
        result->source_frame_seq <= result->capture_baseline_seq ||
        result->channels == 0u || result->channels > AIW_MAX_CHANNELS ||
        (result->status_flags & (uint8_t)~AIW_RESULT_ALLOWED_FLAGS) != 0u) return false;
    if ((result->status_flags & AIW_RESULT_FLAG_HAND_PRESENT) == 0u) {
        for (uint16_t i = 0; i < result->channels; ++i) {
            if (result->pwm[i] != 0u) return false;
        }
    }
    return true;
}

static bool valid_ui_ack(const aiw_ui_ack *ack) {
    if (ack == NULL || ack->session_id == 0u ||
        ack->selected_mode > 2u || (ack->flags & (uint8_t)~AIW_UI_ACK_ALLOWED_FLAGS) != 0u ||
        (unsigned)ack->status > (unsigned)AIW_UI_ACK_SESSION_REJECTED) return false;
    if (ack->sequence == 0u) {
        return (ack->status == AIW_UI_ACK_SESSION_OPENED ||
                ack->status == AIW_UI_ACK_SESSION_REJECTED) &&
               ack->selected_mode == 0u && ack->flags == 0u &&
               ack->remaining_ms == 0u;
    }
    return ack->status <= AIW_UI_ACK_SEQUENCE_COLLISION &&
           (((ack->flags & AIW_UI_ACK_FLAG_RUNNING) != 0u) || ack->remaining_ms == 0u);
}

static bool packet_matches(const uint8_t *packet, size_t length, uint8_t type,
                           uint16_t payload_length) {
    if (packet == NULL || length != AIW_HEADER_BYTES + payload_length + AIW_CRC_BYTES ||
        packet[0] != AIW_MAGIC_0 || packet[1] != AIW_MAGIC_1 ||
        packet[2] != AIW_MAGIC_2 || packet[3] != AIW_MAGIC_3 ||
        packet[4] != AIW_VERSION || packet[5] != type ||
        get16(packet + 6) != payload_length) return false;
    return aiw_crc16_ccitt_false(packet, length - AIW_CRC_BYTES) ==
           get16(packet + length - AIW_CRC_BYTES);
}

bool aiw_encode_ticket(const aiw_ticket *ticket, uint8_t *out, size_t capacity,
                       size_t *written) {
    if (written) *written = 0u;
    if (!written || !out || !valid_ticket(ticket)) return false;
    const size_t packet_len = AIW_HEADER_BYTES + AIW_TICKET_PAYLOAD_BYTES + AIW_CRC_BYTES;
    if (capacity < packet_len) return false;
    set_header(out, AIW_TYPE_REQUEST, AIW_TICKET_PAYLOAD_BYTES);
    uint8_t *p = out + AIW_HEADER_BYTES;
    put64(p, ticket->session_id); p += 8;
    put32(p, ticket->ticket_id); p += 4;
    put32(p, ticket->calibration_id); p += 4;
    put16(p, ticket->channels); p += 2;
    put64(p, ticket->min_source_frame_seq); p += 8;
    put16(p, aiw_crc16_ccitt_false(out, (size_t)(p - out)));
    *written = packet_len;
    return true;
}

bool aiw_decode_ticket(const uint8_t *packet, size_t length, aiw_ticket *out) {
    if (!out) return false;
    memset(out, 0, sizeof(*out));
    const size_t expected = AIW_HEADER_BYTES + AIW_TICKET_PAYLOAD_BYTES + AIW_CRC_BYTES;
    if (!packet || length != expected ||
        packet[0] != AIW_MAGIC_0 || packet[1] != AIW_MAGIC_1 ||
        packet[2] != AIW_MAGIC_2 || packet[3] != AIW_MAGIC_3 ||
        packet[4] != AIW_VERSION || packet[5] != AIW_TYPE_REQUEST ||
        get16(packet + 6) != AIW_TICKET_PAYLOAD_BYTES) return false;
    if (aiw_crc16_ccitt_false(packet, length - AIW_CRC_BYTES) !=
        get16(packet + length - AIW_CRC_BYTES)) return false;
    const uint8_t *p = packet + AIW_HEADER_BYTES;
    aiw_ticket decoded = {0};
    decoded.session_id = get64(p); p += 8;
    decoded.ticket_id = get32(p); p += 4;
    decoded.calibration_id = get32(p); p += 4;
    decoded.channels = get16(p);
    p += 2;
    decoded.min_source_frame_seq = get64(p);
    if (!valid_ticket(&decoded)) return false;
    *out = decoded;
    return true;
}

bool aiw_encode_result(const aiw_result *result, uint8_t *out, size_t capacity,
                       size_t *written) {
    if (written) *written = 0u;
    if (!written || !out || !valid_result(result)) return false;
    const size_t payload_len = AIW_RESULT_FIXED_PAYLOAD_BYTES + result->channels;
    const size_t packet_len = AIW_HEADER_BYTES + payload_len + AIW_CRC_BYTES;
    if (capacity < packet_len) return false;

    set_header(out, AIW_TYPE_RESULT, (uint16_t)payload_len);
    uint8_t *p = out + AIW_HEADER_BYTES;
    put64(p, result->session_id); p += 8;
    put32(p, result->ticket_id); p += 4;
    put32(p, result->calibration_id); p += 4;
    put16(p, result->channels); p += 2;
    put64(p, result->capture_baseline_seq); p += 8;
    put64(p, result->source_frame_seq); p += 8;
    *p++ = result->status_flags;
    memcpy(p, result->pwm, result->channels); p += result->channels;
    const uint16_t crc = aiw_crc16_ccitt_false(out, (size_t)(p - out));
    put16(p, crc);
    *written = packet_len;
    return true;
}

bool aiw_decode_result(const uint8_t *packet, size_t length, aiw_result *out) {
    if (!out) return false;
    memset(out, 0, sizeof(*out));
    if (!packet || length < AIW_HEADER_BYTES + AIW_RESULT_FIXED_PAYLOAD_BYTES + AIW_CRC_BYTES ||
        length > AIW_MAX_PACKET_BYTES) return false;
    if (packet[0] != AIW_MAGIC_0 || packet[1] != AIW_MAGIC_1 ||
        packet[2] != AIW_MAGIC_2 || packet[3] != AIW_MAGIC_3 ||
        packet[4] != AIW_VERSION || packet[5] != AIW_TYPE_RESULT) return false;
    const uint16_t payload_len = get16(packet + 6);
    if ((size_t)payload_len + AIW_HEADER_BYTES + AIW_CRC_BYTES != length ||
        payload_len < AIW_RESULT_FIXED_PAYLOAD_BYTES) return false;
    const uint8_t *p = packet + AIW_HEADER_BYTES;
    aiw_result decoded = {0};
    decoded.session_id = get64(p); p += 8;
    decoded.ticket_id = get32(p); p += 4;
    decoded.calibration_id = get32(p); p += 4;
    decoded.channels = get16(p); p += 2;
    decoded.capture_baseline_seq = get64(p); p += 8;
    decoded.source_frame_seq = get64(p); p += 8;
    decoded.status_flags = *p++;
    if (decoded.channels == 0u || decoded.channels > AIW_MAX_CHANNELS ||
        payload_len != AIW_RESULT_FIXED_PAYLOAD_BYTES + decoded.channels) return false;
    const uint16_t expected_crc = aiw_crc16_ccitt_false(packet, length - AIW_CRC_BYTES);
    const uint16_t received_crc = get16(packet + length - AIW_CRC_BYTES);
    if (expected_crc != received_crc) return false;
    memcpy(decoded.pwm, p, decoded.channels);
    if (!valid_result(&decoded)) return false;
    *out = decoded;
    return true;
}

bool aiw_encode_ui_session_open(const aiw_ui_session_open *open, uint8_t *out,
                                size_t capacity, size_t *written) {
    if (written) *written = 0u;
    const size_t packet_length = AIW_HEADER_BYTES + 8u + AIW_CRC_BYTES;
    if (!written || !out || !open || open->session_id == 0u || capacity < packet_length)
        return false;
    set_header(out, AIW_TYPE_UI_SESSION_OPEN, 8u);
    put64(out + AIW_HEADER_BYTES, open->session_id);
    put16(out + packet_length - AIW_CRC_BYTES,
          aiw_crc16_ccitt_false(out, packet_length - AIW_CRC_BYTES));
    *written = packet_length;
    return true;
}

bool aiw_decode_ui_session_open(const uint8_t *packet, size_t length,
                                aiw_ui_session_open *out) {
    if (!out) return false;
    memset(out, 0, sizeof(*out));
    if (!packet_matches(packet, length, AIW_TYPE_UI_SESSION_OPEN, 8u)) return false;
    out->session_id = get64(packet + AIW_HEADER_BYTES);
    return out->session_id != 0u;
}

bool aiw_encode_ui_command(const aiw_ui_command *command, uint8_t *out,
                           size_t capacity, size_t *written) {
    if (written) *written = 0u;
    const size_t packet_length = AIW_HEADER_BYTES + 13u + AIW_CRC_BYTES;
    if (!written || !out || !command || command->session_id == 0u ||
        command->sequence == 0u || command->button > AIW_UI_COMMAND_MAX_BUTTON ||
        capacity < packet_length) return false;
    set_header(out, AIW_TYPE_UI_COMMAND, 13u);
    uint8_t *p = out + AIW_HEADER_BYTES;
    put64(p, command->session_id); p += 8;
    put32(p, command->sequence); p += 4;
    *p++ = command->button;
    put16(p, aiw_crc16_ccitt_false(out, (size_t)(p - out)));
    *written = packet_length;
    return true;
}

bool aiw_decode_ui_command(const uint8_t *packet, size_t length,
                           aiw_ui_command *out) {
    if (!out) return false;
    memset(out, 0, sizeof(*out));
    if (!packet_matches(packet, length, AIW_TYPE_UI_COMMAND, 13u)) return false;
    const uint8_t *p = packet + AIW_HEADER_BYTES;
    out->session_id = get64(p); p += 8;
    out->sequence = get32(p); p += 4;
    out->button = *p;
    return out->session_id != 0u && out->sequence != 0u &&
           out->button <= AIW_UI_COMMAND_MAX_BUTTON;
}

bool aiw_encode_ui_ack(const aiw_ui_ack *ack, uint8_t *out, size_t capacity,
                       size_t *written) {
    if (written) *written = 0u;
    const size_t packet_length = AIW_HEADER_BYTES + 19u + AIW_CRC_BYTES;
    if (!written || !out || !valid_ui_ack(ack) || capacity < packet_length) return false;
    set_header(out, AIW_TYPE_UI_ACK, 19u);
    uint8_t *p = out + AIW_HEADER_BYTES;
    put64(p, ack->session_id); p += 8;
    put32(p, ack->sequence); p += 4;
    *p++ = (uint8_t)ack->status;
    *p++ = ack->selected_mode;
    *p++ = ack->flags;
    put32(p, ack->remaining_ms); p += 4;
    put16(p, aiw_crc16_ccitt_false(out, (size_t)(p - out)));
    *written = packet_length;
    return true;
}

bool aiw_decode_ui_ack(const uint8_t *packet, size_t length, aiw_ui_ack *out) {
    if (!out) return false;
    memset(out, 0, sizeof(*out));
    if (!packet_matches(packet, length, AIW_TYPE_UI_ACK, 19u)) return false;
    const uint8_t *p = packet + AIW_HEADER_BYTES;
    out->session_id = get64(p); p += 8;
    out->sequence = get32(p); p += 4;
    const uint8_t status = *p++;
    if (status > AIW_UI_ACK_SESSION_REJECTED) return false;
    out->status = (aiw_ui_ack_status)status;
    out->selected_mode = *p++;
    out->flags = *p++;
    out->remaining_ms = get32(p);
    return valid_ui_ack(out);
}

bool aiw_decode_message(const uint8_t *packet, size_t length, aiw_message *out) {
    if (!out || !packet || length < AIW_HEADER_BYTES) return false;
    memset(out, 0, sizeof(*out));
    switch (packet[5]) {
        case AIW_TYPE_RESULT:
            out->type = AIW_MESSAGE_RESULT;
            return aiw_decode_result(packet, length, &out->value.result);
        case AIW_TYPE_UI_SESSION_OPEN:
            out->type = AIW_MESSAGE_UI_SESSION_OPEN;
            return aiw_decode_ui_session_open(packet, length, &out->value.ui_session_open);
        case AIW_TYPE_UI_COMMAND:
            out->type = AIW_MESSAGE_UI_COMMAND;
            return aiw_decode_ui_command(packet, length, &out->value.ui_command);
        case AIW_TYPE_UI_ACK:
            out->type = AIW_MESSAGE_UI_ACK;
            return aiw_decode_ui_ack(packet, length, &out->value.ui_ack);
        default:
            return false;
    }
}
