#include "ota_frame_protocol.h"
#include "ota_binary_stream.h"
#include "ota_offer_protocol.h"
#include "esp_log.h"
#include <string.h>

static const char *tag_s = "OTA_FRAME";

static ota_frame_tx_fn s_tx_fn;
static void *s_tx_ctx;

typedef enum {
    PS_SYNC0,
    PS_SYNC1,
    PS_VER,
    PS_OP,
    PS_SEQ0,
    PS_SEQ1,
    PS_LEN0,
    PS_LEN1,
    PS_PAYLOAD,
    PS_CRC,
    PS_TAIL,
} parser_state_t;

static struct {
    parser_state_t state;
    uint8_t ver;
    uint8_t op;
    uint16_t seq;
    uint16_t payload_len;
    uint16_t payload_pos;
    uint8_t payload[OTA_FRAME_MAX_PAYLOAD];
    uint8_t crc_rx;
} s_parser;

static uint8_t ota_frame_crc8(const uint8_t *data, uint16_t len)
{
    uint8_t c = 0;
    for (uint16_t i = 0; i < len; i++) {
        c ^= data[i];
    }
    return c;
}

void ota_frame_protocol_init(ota_frame_tx_fn tx, void *ctx)
{
    s_tx_fn = tx;
    s_tx_ctx = ctx;
    ota_frame_protocol_reset();
}

void ota_frame_protocol_reset(void)
{
    memset(&s_parser, 0, sizeof(s_parser));
    s_parser.state = PS_SYNC0;
}

bool ota_frame_rx_looks_framed(const uint8_t *data, uint16_t len)
{
    return (data != NULL && len >= 1 && data[0] == OTA_FRAME_MAGIC0 &&
            (len == 1 || data[1] == OTA_FRAME_MAGIC1));
}

bool ota_frame_rx_is_active(void)
{
    return s_parser.state != PS_SYNC0;
}

int ota_frame_build(uint8_t op, uint16_t seq, const uint8_t *payload, uint16_t payload_len,
                    uint8_t *out, uint16_t out_cap, uint16_t *out_len)
{
    if (out == NULL || out_len == NULL || payload_len > OTA_FRAME_MAX_PAYLOAD) {
        return -1;
    }
    uint16_t need = (uint16_t)(8 + payload_len + 2);
    if (out_cap < need) {
        return -2;
    }

    out[0] = OTA_FRAME_MAGIC0;
    out[1] = OTA_FRAME_MAGIC1;
    out[2] = OTA_FRAME_VERSION;
    out[3] = op;
    out[4] = (uint8_t)(seq & 0xFF);
    out[5] = (uint8_t)((seq >> 8) & 0xFF);
    out[6] = (uint8_t)(payload_len & 0xFF);
    out[7] = (uint8_t)((payload_len >> 8) & 0xFF);
    if (payload_len > 0 && payload != NULL) {
        memcpy(out + 8, payload, payload_len);
    }
    out[8 + payload_len] = ota_frame_crc8(out + 2, (uint16_t)(6 + payload_len));
    out[8 + payload_len + 1] = OTA_FRAME_TAIL_BYTE;
    *out_len = need;
    return 0;
}

static void ota_frame_tx_bytes(const uint8_t *frame, uint16_t len)
{
    if (s_tx_fn != NULL && frame != NULL && len > 0) {
        s_tx_fn(frame, len, s_tx_ctx);
    }
}

static void ota_frame_send_simple(uint8_t op, uint16_t seq, const uint8_t *payload, uint16_t plen)
{
    uint8_t buf[OTA_FRAME_BUF_SIZE];
    uint16_t out_len = 0;
    if (ota_frame_build(op, seq, payload, plen, buf, sizeof(buf), &out_len) == 0) {
        ota_frame_tx_bytes(buf, out_len);
    }
}

static void ota_frame_send_ack(uint16_t seq, uint8_t data_byte, uint8_t status)
{
    uint8_t pl[2] = { data_byte, status };
    ota_frame_send_simple(OTA_FRAME_OP_ACK, seq, pl, 2);
}

static void ota_frame_send_status_ack(uint16_t seq, uint8_t status)
{
    ota_frame_send_simple(OTA_FRAME_OP_ACK, seq, &status, 1);
}

static void ota_frame_send_nak(uint16_t seq, uint8_t err_code)
{
    ota_frame_send_simple(OTA_FRAME_OP_NAK, seq, &err_code, 1);
}

static int map_stream_rc_to_ack(int rc)
{
    if (rc == 0) {
        return OTA_ACK_OK;
    }
    if (rc == -2) {
        return OTA_ACK_ERR_OVERFLOW;
    }
    if (rc == -1) {
        return OTA_ACK_ERR_NO_SESSION;
    }
    return OTA_ACK_ERR_FLASH;
}

static void handle_complete_frame(void)
{
    if (!ota_offer_is_armed() && !ota_binary_stream_is_active()) {
        ESP_LOGW(tag_s, "忽略 OTA 帧：未 armed");
        ota_frame_send_nak(s_parser.seq, OTA_ACK_ERR_NO_SESSION);
        return;
    }

    const uint8_t *p = s_parser.payload;
    uint16_t plen = s_parser.payload_len;

    switch (s_parser.op) {
    case OTA_FRAME_OP_BEGIN: {
        if (plen != 4) {
            ota_frame_send_nak(s_parser.seq, OTA_ACK_ERR_BAD_LEN);
            return;
        }
        uint8_t pkt[5] = { OTA_OP_BEGIN, p[0], p[1], p[2], p[3] };
        int rc = ota_binary_stream_process_packet(pkt, 5);
        ota_frame_send_status_ack(s_parser.seq, map_stream_rc_to_ack(rc));
        break;
    }
    case OTA_FRAME_OP_DATA: {
        if (plen != 1) {
            ota_frame_send_nak(s_parser.seq, OTA_ACK_ERR_BAD_LEN);
            return;
        }
        uint8_t b = p[0];
        uint8_t pkt[2] = { OTA_OP_DATA, b };
        int rc = ota_binary_stream_process_packet(pkt, 2);
        ota_frame_send_ack(s_parser.seq, b, (uint8_t)map_stream_rc_to_ack(rc));
        break;
    }
    case OTA_FRAME_OP_END: {
        if (plen != 0) {
            ota_frame_send_nak(s_parser.seq, OTA_ACK_ERR_BAD_LEN);
            return;
        }
        uint8_t pkt[1] = { OTA_OP_END };
        int rc = ota_binary_stream_process_packet(pkt, 1);
        ota_frame_send_status_ack(s_parser.seq, map_stream_rc_to_ack(rc));
        break;
    }
    default:
        ota_frame_send_nak(s_parser.seq, OTA_ACK_ERR_BAD_OP);
        break;
    }
}

static void parser_push_byte(uint8_t b)
{
    switch (s_parser.state) {
    case PS_SYNC0:
        if (b == OTA_FRAME_MAGIC0) {
            s_parser.state = PS_SYNC1;
        }
        break;
    case PS_SYNC1:
        if (b == OTA_FRAME_MAGIC1) {
            s_parser.state = PS_VER;
        } else {
            s_parser.state = PS_SYNC0;
            if (b == OTA_FRAME_MAGIC0) {
                s_parser.state = PS_SYNC1;
            }
        }
        break;
    case PS_VER:
        s_parser.ver = b;
        s_parser.state = PS_OP;
        break;
    case PS_OP:
        s_parser.op = b;
        s_parser.state = PS_SEQ0;
        break;
    case PS_SEQ0:
        s_parser.seq = b;
        s_parser.state = PS_SEQ1;
        break;
    case PS_SEQ1:
        s_parser.seq |= (uint16_t)b << 8;
        s_parser.state = PS_LEN0;
        break;
    case PS_LEN0:
        s_parser.payload_len = b;
        s_parser.state = PS_LEN1;
        break;
    case PS_LEN1:
        s_parser.payload_len |= (uint16_t)b << 8;
        s_parser.payload_pos = 0;
        if (s_parser.payload_len > OTA_FRAME_MAX_PAYLOAD) {
            s_parser.state = PS_SYNC0;
            ota_frame_send_nak(0, OTA_ACK_ERR_BAD_LEN);
            return;
        }
        if (s_parser.payload_len == 0) {
            s_parser.state = PS_CRC;
        } else {
            s_parser.state = PS_PAYLOAD;
        }
        break;
    case PS_PAYLOAD:
        s_parser.payload[s_parser.payload_pos++] = b;
        if (s_parser.payload_pos >= s_parser.payload_len) {
            s_parser.state = PS_CRC;
        }
        break;
    case PS_CRC:
        s_parser.crc_rx = b;
        s_parser.state = PS_TAIL;
        break;
    case PS_TAIL:
        if (b != OTA_FRAME_TAIL_BYTE) {
            s_parser.state = PS_SYNC0;
            ota_frame_send_nak(s_parser.seq, OTA_ACK_ERR_CRC);
            return;
        }
        if (s_parser.ver != OTA_FRAME_VERSION) {
            s_parser.state = PS_SYNC0;
            ota_frame_send_nak(s_parser.seq, OTA_ACK_ERR_BAD_OP);
            return;
        }
        {
            uint8_t hdr[6 + OTA_FRAME_MAX_PAYLOAD];
            hdr[0] = s_parser.ver;
            hdr[1] = s_parser.op;
            hdr[2] = (uint8_t)(s_parser.seq & 0xFF);
            hdr[3] = (uint8_t)((s_parser.seq >> 8) & 0xFF);
            hdr[4] = (uint8_t)(s_parser.payload_len & 0xFF);
            hdr[5] = (uint8_t)((s_parser.payload_len >> 8) & 0xFF);
            if (s_parser.payload_len > 0) {
                memcpy(hdr + 6, s_parser.payload, s_parser.payload_len);
            }
            uint8_t crc_calc = ota_frame_crc8(hdr, (uint16_t)(6 + s_parser.payload_len));
            if (crc_calc != s_parser.crc_rx) {
                s_parser.state = PS_SYNC0;
                ota_frame_send_nak(s_parser.seq, OTA_ACK_ERR_CRC);
                return;
            }
            handle_complete_frame();
        }
        s_parser.state = PS_SYNC0;
        break;
    default:
        s_parser.state = PS_SYNC0;
        break;
    }
}

void ota_frame_rx_feed(const uint8_t *data, uint16_t len)
{
    if (data == NULL || len == 0) {
        return;
    }
    for (uint16_t i = 0; i < len; i++) {
        parser_push_byte(data[i]);
    }
}
