/*
 * OTA 二进制流接收模块实现
 * 通用 OTA 实现，支持 BLE 和 UART 两种传输方式
 */

#include "sdkconfig.h"
#include "ota_binary_stream.h"
#include "ota_offer_protocol.h"
#include "ota_manager.h"
#include "esp_ota_ops.h"
#include "esp_partition.h"
#include "esp_log.h"
#include "esp_system.h"
#include "freertos/FreeRTOS.h"
#include <string.h>
#include <stdio.h>
#include <stdarg.h>

static const char *tag_s = "OTA_STREAM";

/* OTA 会话状态 */
typedef struct {
    bool active;
    esp_ota_handle_t handle;
    const esp_partition_t *partition;
    size_t expected_size;
    size_t written_size;
    size_t last_reported_size;
    ota_transport_type_t transport;
    ota_binary_stream_status_callback_t status_callback;
} ota_stream_session_t;

static ota_stream_session_t session_s = {0};

/* 格式化并发送状态回调 */
static void ota_stream_send_status_fmt(const char *fmt, ...)
{
    if (session_s.status_callback == NULL) {
        return;
    }

    char buf[192];
    va_list ap;
    va_start(ap, fmt);
    vsnprintf(buf, sizeof(buf), fmt, ap);
    va_end(ap);

    session_s.status_callback(buf);
}

/* 发送错误状态并结束 OTA 会话（仅用于致命错误，须已 armed） */
static void ota_stream_send_error(const char *step, const char *err)
{
    ESP_LOGE(tag_s, "OTA error at %s: %s", step, err);
    if (ota_offer_is_armed() || session_s.active) {
        ota_stream_send_status_fmt("{\"ota\":\"%s\",\"ok\":0,\"err\":\"%s\"}", step, err);
        ota_offer_notify_done(false, (uint32_t)session_s.written_size, step);
    }
}

/**
 * @brief 擦除 OTA 目标分区（失败清理，防止下次启动误用损坏固件）
 */
static void ota_stream_erase_partition(void)
{
    if (session_s.partition == NULL) {
        return;
    }

    ESP_LOGW(tag_s, "OTA 失败：擦除分区 %s 防止误启动", session_s.partition->label);

    const esp_partition_t *part = session_s.partition;
    esp_err_t err = esp_partition_erase_range(part, 0, part->size);
    if (err != ESP_OK) {
        ESP_LOGE(tag_s, "擦除分区失败: %s", esp_err_to_name(err));
    } else {
        ESP_LOGI(tag_s, "分区 %s 已擦除 (%u bytes)", part->label, (unsigned)part->size);
    }
}

/* 中止 OTA 会话 */
static void ota_stream_abort_internal(void)
{
    if (session_s.active && session_s.handle != 0) {
        esp_ota_abort(session_s.handle);
        ESP_LOGW(tag_s, "OTA session aborted");
    }

    /* 如果已写入部分数据（>4KB），擦除分区避免启动误用 */
    if (session_s.written_size > 4096 && session_s.partition != NULL) {
        ota_stream_erase_partition();
    }

    session_s.active = false;
    session_s.handle = 0;
    session_s.partition = NULL;
    session_s.expected_size = 0;
    session_s.written_size = 0;
    session_s.last_reported_size = 0;
}

void ota_binary_stream_abort(void)
{
    ota_stream_abort_internal();
}

bool ota_binary_stream_is_active(void)
{
    return session_s.active;
}

uint32_t ota_binary_stream_get_written_bytes(void)
{
    return (uint32_t)session_s.written_size;
}

esp_err_t ota_binary_stream_init(ota_transport_type_t transport,
                                  ota_binary_stream_status_callback_t status_callback,
                                  const char *target_version)
{
    if (session_s.active) {
        ESP_LOGW(tag_s, "OTA already active");
        return ESP_ERR_INVALID_STATE;
    }

    session_s.transport = transport;
    session_s.status_callback = status_callback;
    (void)target_version;

    ESP_LOGI(tag_s, "OTA stream init: transport=%d (mgr session starts at BEGIN/0x01)", transport);
    return ESP_OK;
}

void ota_binary_stream_deinit(void)
{
    if (session_s.active) {
        ota_stream_abort_internal();
    }
    session_s.status_callback = NULL;
}

/* 处理 OTA 开始命令 */
static int ota_stream_handle_begin(const uint8_t *data, uint16_t len)
{
    if (session_s.active) {
        ESP_LOGW(tag_s, "OTA begin while active, abort previous session");
        ota_stream_abort_internal();
    }

    if (len < 5) {
        /* 需要至少 1字节opcode + 4字节长度 */
        ota_stream_send_error("begin", "size");
        return -1;
    }

    /* 解析固件长度（小端序 uint32） */
    uint32_t image_size = ((uint32_t)data[1]) |
                          ((uint32_t)data[2] << 8) |
                          ((uint32_t)data[3] << 16) |
                          ((uint32_t)data[4] << 24);

    if (image_size == 0 || image_size > (6 * 1024 * 1024)) {
        /* 固件大小不合理（最大6MB） */
        ota_stream_send_error("begin", "size_invalid");
        return -2;
    }

    ESP_LOGI(tag_s, "OTA begin: size=%lu bytes", (unsigned long)image_size);

    /* 查找 OTA 更新分区 */
    session_s.partition = esp_ota_get_next_update_partition(NULL);
    if (session_s.partition == NULL) {
        ota_stream_send_error("begin", "no_partition");
        return -3;
    }

    ESP_LOGI(tag_s, "OTA partition: %s at 0x%08x",
             session_s.partition->label,
             (unsigned)session_s.partition->address);

    /* 开始 OTA */
    esp_err_t err = esp_ota_begin(session_s.partition, image_size, &session_s.handle);
    if (err != ESP_OK) {
        ESP_LOGE(tag_s, "esp_ota_begin failed: %s", esp_err_to_name(err));
        ota_stream_send_error("begin", "esp_ota_begin");
        session_s.handle = 0;
        session_s.partition = NULL;
        return -4;
    }

    ota_type_t type = (session_s.transport == OTA_TRANSPORT_BLE) ? OTA_TYPE_BLE : OTA_TYPE_UART;
    esp_err_t mgr_err = ota_manager_begin_session(type, CONFIG_MY_APP_VERSION);
    if (mgr_err != ESP_OK) {
        ESP_LOGE(tag_s, "ota_manager_begin_session failed: %s", esp_err_to_name(mgr_err));
        esp_ota_abort(session_s.handle);
        session_s.handle = 0;
        session_s.partition = NULL;
        ota_stream_send_error("begin", "mgr_session");
        return -5;
    }

    session_s.expected_size = image_size;
    session_s.written_size = 0;
    session_s.last_reported_size = 0;
    session_s.active = true;

    ota_stream_send_status_fmt("{\"ota\":\"begin\",\"ok\":1,\"expect\":%lu}",
                                 (unsigned long)image_size);

    return 0;
}

/* 处理 OTA 数据包 */
static int ota_stream_handle_data(const uint8_t *data, uint16_t len)
{
    if (!session_s.active) {
        /* 仅 NAK，不结束 OTA：可能 BEGIN 未到或 App 重传；避免 done 风暴 */
        ESP_LOGW(tag_s, "OTA DATA without active session (armed=%d)",
                 (int)ota_offer_is_armed());
        return -1;
    }

    if (len < 1) {
        /* 空数据包 */
        return 0;
    }

    /* data[0] 是 opcode，后面是实际数据 */
    const uint8_t *payload = data + 1;
    uint16_t payload_len = len - 1;

    if (payload_len == 0) {
        return 0;
    }

    /* 检查是否超出预期大小 */
    if (session_s.written_size + payload_len > session_s.expected_size) {
        ESP_LOGE(tag_s, "OTA overflow: %u + %u > %u",
                 (unsigned)session_s.written_size,
                 (unsigned)payload_len,
                 (unsigned)session_s.expected_size);
        ota_stream_send_error("chunk", "overflow");
        ota_stream_abort_internal();
        return -2;
    }

    /* 写入数据 */
    esp_err_t err = esp_ota_write(session_s.handle, payload, payload_len);
    if (err != ESP_OK) {
        ESP_LOGE(tag_s, "esp_ota_write failed: %s", esp_err_to_name(err));
        ota_stream_send_error("chunk", "write");
        ota_stream_abort_internal();
        return -3;
    }

    session_s.written_size += payload_len;

    /* BLE 调试 JSON；UART OTA 期间 UART1 只走 OT 帧，不发 chunk 状态 JSON */
    if (session_s.transport != OTA_TRANSPORT_UART) {
        if (session_s.written_size - session_s.last_reported_size >= 8192 ||
            session_s.written_size >= session_s.expected_size) {
            ota_stream_send_status_fmt("{\"ota\":\"chunk\",\"ok\":1,\"written\":%u}",
                                        (unsigned)session_s.written_size);
            session_s.last_reported_size = session_s.written_size;
        }
    }

    return 0;
}

/* 处理 OTA 结束命令 */
static int ota_stream_handle_end(void)
{
    if (!session_s.active) {
        ESP_LOGW(tag_s, "OTA END without active session");
        return -1;
    }

    /* 检查写入大小是否匹配 */
    if (session_s.written_size != session_s.expected_size) {
        ESP_LOGE(tag_s, "OTA incomplete: %u/%u",
                 (unsigned)session_s.written_size,
                 (unsigned)session_s.expected_size);
        ota_stream_send_error("end", "incomplete");
        ota_stream_abort_internal();
        return -2;
    }

    /* 结束 OTA 写入 */
    esp_err_t err = esp_ota_end(session_s.handle);
    if (err != ESP_OK) {
        ESP_LOGE(tag_s, "esp_ota_end failed: %s", esp_err_to_name(err));
        ota_stream_send_error("end", "esp_ota_end");
        ota_stream_abort_internal();
        return -3;
    }

    session_s.handle = 0;
    session_s.active = false;

    /* 设置启动分区 */
    err = esp_ota_set_boot_partition(session_s.partition);
    if (err != ESP_OK) {
        ESP_LOGE(tag_s, "esp_ota_set_boot_partition failed: %s", esp_err_to_name(err));
        ota_stream_send_error("end", "set_boot");
        return -4;
    }

    ESP_LOGI(tag_s, "OTA successful, written %u bytes, next boot from %s",
             (unsigned)session_s.written_size, session_s.partition->label);

    /* 通过OTA管理器结束会话 */
    ota_manager_end_session();

    ota_stream_send_status_fmt("{\"ota\":\"end\",\"ok\":1,\"written\":%u}",
                                (unsigned)session_s.written_size);

    ota_offer_notify_done(true, (uint32_t)session_s.written_size, NULL);

    /* notify_done 内已连发 OTA_DONE_REPORT_COUNT 次 1002/done，再等 DMA 发完后重启切分区 */
    vTaskDelay(pdMS_TO_TICKS(OTA_DONE_REBOOT_DELAY_MS));
    ESP_LOGI(tag_s, "Rebooting to %s...", session_s.partition->label);
    esp_restart();

    return 0;
}

int ota_binary_stream_process_packet(const uint8_t *data, uint16_t len)
{
    if (data == NULL || len == 0) {
        return -1;
    }

    uint8_t opcode = data[0];

    switch (opcode) {
        case OTA_OP_BEGIN:
            return ota_stream_handle_begin(data, len);

        case OTA_OP_DATA:
            return ota_stream_handle_data(data, len);

        case OTA_OP_END:
            return ota_stream_handle_end();

        default:
            ESP_LOGW(tag_s, "Unknown OTA opcode: 0x%02X", opcode);
            ota_stream_send_error("unknown", "invalid_opcode");
            return -10;
    }
}

esp_err_t ota_binary_stream_get_status_json(char *buf, size_t len)
{
    if (buf == NULL || len == 0) {
        return ESP_ERR_INVALID_ARG;
    }

    if (!session_s.active) {
        strncpy(buf, "{\"ota\":\"idle\"}", len);
        return ESP_OK;
    }

    snprintf(buf, len,
             "{\"ota\":\"active\",\"written\":%u,\"expect\":%u,\"transport\":%d}",
             (unsigned)session_s.written_size,
             (unsigned)session_s.expected_size,
             session_s.transport);

    return ESP_OK;
}
