#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) {
        if (op == OTA_FRAME_OP_ACK && plen >= 2) {
            ESP_LOGD(tag_s, "TX ACK seq=%u data=0x%02X status=%u",
                     (unsigned)seq, payload[0], (unsigned)payload[1]);
        } else if (op == OTA_FRAME_OP_ACK && plen == 1) {
            ESP_LOGI(tag_s, "TX ACK seq=%u status=%u", (unsigned)seq, (unsigned)payload[0]);
        } else if (op == OTA_FRAME_OP_NAK) {
            ESP_LOGW(tag_s, "TX NAK seq=%u err=%u", (unsigned)seq,
                     plen > 0 ? (unsigned)payload[0] : 0U);
        }
        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_chunk_ack(uint16_t seq, uint16_t chunk_len,
                                     uint8_t last_byte, uint8_t status)
{
    uint8_t pl[4] = {
        (uint8_t)(chunk_len & 0xFF),
        (uint8_t)((chunk_len >> 8) & 0xFF),
        last_byte,
        status,
    };
    ota_frame_send_simple(OTA_FRAME_OP_ACK, seq, pl, 4);
}

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 const char *ota_frame_op_name(uint8_t op)
{
    switch (op) {
    case OTA_FRAME_OP_BEGIN:
        return "BEGIN";
    case OTA_FRAME_OP_DATA:
        return "DATA";
    case OTA_FRAME_OP_END:
        return "END";
    case OTA_FRAME_OP_ACK:
        return "ACK";
    case OTA_FRAME_OP_NAK:
        return "NAK";
    default:
        return "UNKNOWN";
    }
}

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;

    if (s_parser.op == OTA_FRAME_OP_DATA) {
        ESP_LOGD(tag_s, "RX %s seq=%u plen=%u", ota_frame_op_name(s_parser.op),
                 (unsigned)s_parser.seq, (unsigned)plen);
    } else {
        ESP_LOGI(tag_s, "RX %s seq=%u plen=%u", ota_frame_op_name(s_parser.op),
                 (unsigned)s_parser.seq, (unsigned)plen);
    }

    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 || plen > OTA_FRAME_MAX_PAYLOAD) {
            ota_frame_send_nak(s_parser.seq, OTA_ACK_ERR_BAD_LEN);
            return;
        }
        uint8_t pkt[1 + OTA_FRAME_MAX_PAYLOAD];
        pkt[0] = OTA_OP_DATA;
        memcpy(pkt + 1, p, plen);
        int rc = ota_binary_stream_process_packet(pkt, (uint16_t)(1 + plen));
        uint8_t ack_status = (uint8_t)map_stream_rc_to_ack(rc);
        if (plen == 1) {
            ota_frame_send_ack(s_parser.seq, p[0], ack_status);
        } else {
            ota_frame_send_chunk_ack(s_parser.seq, plen, p[plen - 1], ack_status);
        }
        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]);
    }
}
