// Copyright 2026 Shanling.
// HID固件升级实现文件 (adapted for app_usb_aud_phaten_golden_6ch)

#define DEBUG_PRINT_ENABLE 0
#include "dfu_upgrade.h"
#include "user_func.h"
#include "dfu_flash_interface.h"
#include <string.h>
#include <stdio.h>
#include "debug_print.h"
#include "swlock.h"

// Global flash lock defined in lfs_services.c — protects all flash hardware access
extern swlock_t flash_lock;

#if HID_DFU_EN
// 外部函数声明
extern void hidSetChangePending(unsigned int id);
extern unsigned char check_sum(unsigned char *pbuf, unsigned len);
extern void device_reboot(void);
extern unsigned hidSendData[];

// 固件升级通知MCU的标志（phaten无UART MCU，保留为静态变量供内部使用）
static unsigned g_firmware_upgrade_mcu_notify = 0;
unsigned g_in_fw_upgrade = 0;  // 升级进行中标志，供其他模块检查以阻止非升级HID上报

// 升级状态（静态全局变量）
static upgrade_status_t g_upgrade_status;

// 全局标志：第一页写命令是否已执行
static uint8_t g_first_page_write_executed = 0;
// 全局标志：是否已执行dfu_flash_cmd_erase_all
static uint8_t g_flash_erased = 0;

// 初始化固件升级模块
void firmware_upgrade_init(void)
{
    memset(&g_upgrade_status, 0, sizeof(upgrade_status_t));
    g_upgrade_status.state = UPGRADE_IDLE;
    g_upgrade_status.expected_block_num = 0;
    debug_printf("Firmware upgrade module initialized\n");
}

// 获取当前升级状态
upgrade_state_t firmware_upgrade_get_state(void)
{
    return g_upgrade_status.state;
}

// 重置升级状态
void firmware_upgrade_reset(void)
{
    memset(&g_upgrade_status, 0, sizeof(upgrade_status_t));
    g_upgrade_status.state = UPGRADE_IDLE;
    g_upgrade_status.expected_block_num = 0;
    g_upgrade_status.first_page_written = 0;
    g_first_page_write_executed = 0;
    g_flash_erased = 0;
}

// 发送固件升级响应（主动上报）
// 填充hidSendData，通过UserHIDGetData上报给host
void send_firmware_upgrade_response(uint8_t cmd, uint8_t *response_data, uint16_t data_len)
{
    unsigned char *ptr = (unsigned char *)hidSendData;
    ptr[0] = 1;
    ptr[1] = 0x77;  // 同步头
    ptr[2] = cmd;   // 命令码

    if (response_data != NULL && data_len > 0) {
        uint16_t copy_len = (data_len < 61) ? data_len : 61;
        memcpy(&ptr[3], response_data, copy_len);
    }

    for (int i = 3 + (int)data_len; i < 64; i++) {
        ptr[i] = 0x00;
    }

    hidSetChangePending(0x1);
}

// 处理START命令
unsigned char handle_firmware_upgrade_start(uint8_t data[], uint16_t len)
{
    uint8_t response[64] = {0};

    if (g_upgrade_status.state != UPGRADE_IDLE) {
        debug_printf("Firmware upgrade error: not in IDLE state (current: %d)\n", g_upgrade_status.state);
        response[0] = STATUS_FAIL;
        send_firmware_upgrade_response(FIRMWARE_UPGRADE_START, response, 1);
        return STATUS_FAIL;
    }

    uint32_t firmware_size = data[2] | (data[3] << 8) | (data[4] << 16) | (data[5] << 24);

    debug_printf("g_first_page_write_executed: %d\n", g_first_page_write_executed);

    if (firmware_size == 0 || firmware_size > FLASH_MAX_UPGRADE_SIZE) {
        debug_printf("Firmware upgrade error: invalid size %lu\n", firmware_size);
        response[0] = STATUS_SIZE_INVALID;
        send_firmware_upgrade_response(FIRMWARE_UPGRADE_START, response, 1);
        return STATUS_SIZE_INVALID;
    }

    uint32_t aligned_size = ((firmware_size + FLASH_PAGE_SIZE - 1) / FLASH_PAGE_SIZE) * FLASH_PAGE_SIZE;

    // Acquire global flash lock — held until END or ABORT releases it
    swlock_acquire(&flash_lock);

    if (dfu_flash_cmd_init() != 0) {
        debug_printf("Firmware upgrade error: flash init failed\n");
        swlock_release(&flash_lock);
        response[0] = STATUS_FAIL;
        send_firmware_upgrade_response(FIRMWARE_UPGRADE_START, response, 1);
        return STATUS_FAIL;
    }

    if (g_first_page_write_executed == 0 && g_flash_erased == 0) {
        debug_printf("Force upgrade: erasing existing upgrade image\n");
        dfu_flash_cmd_erase_all();
        g_flash_erased = 1;
    } else {
        debug_printf("Force upgrade: erase skipped (already erased or begin_write started)\n");
    }

    g_upgrade_status.state = UPGRADE_PREPARING;
    g_upgrade_status.firmware_size = aligned_size;
    g_upgrade_status.total_blocks = (aligned_size + MAX_DATA_BLOCK_SIZE - 1) / MAX_DATA_BLOCK_SIZE;
    g_upgrade_status.received_blocks = 0;
    g_upgrade_status.expected_block_num = 0;
    g_upgrade_status.error_code = 0;

    memset(g_upgrade_status.page_buf.buffer, 0, FLASH_PAGE_SIZE);
    g_upgrade_status.page_buf.offset = 0;
    g_upgrade_status.page_buf.total_pages = aligned_size / FLASH_PAGE_SIZE;
    g_upgrade_status.page_buf.written_pages = 0;

    if (g_first_page_write_executed == 0) {
        debug_printf("Firmware upgrade: executing first page write command (flag=0) in START...\n");
        unsigned char cmd_data[16];
        memset(cmd_data, 0, sizeof(cmd_data));
        cmd_data[0] = 0;
        if (dfu_flash_cmd_write_page(cmd_data) != 0) {
            debug_printf("Firmware upgrade error: first page write command failed (begin_write error)\n");
            dfu_flash_cmd_deinit();
            g_upgrade_status.state = UPGRADE_IDLE;
            swlock_release(&flash_lock);
            response[0] = STATUS_FAIL;
            send_firmware_upgrade_response(FIRMWARE_UPGRADE_START, response, 1);
            return STATUS_FAIL;
        }
        g_first_page_write_executed = 1;
        debug_printf("Firmware upgrade: first page write command completed\n");
    } else {
        debug_printf("Firmware upgrade: first page write command already executed in ERASE, skipping %d\n", g_first_page_write_executed);
    }
    g_upgrade_status.first_page_written = 1;

    response[0] = STATUS_SUCCESS;
    response[1] = (aligned_size >> 0) & 0xFF;
    response[2] = (aligned_size >> 8) & 0xFF;
    response[3] = (aligned_size >> 16) & 0xFF;
    response[4] = (aligned_size >> 24) & 0xFF;
    response[5] = (g_upgrade_status.total_blocks >> 0) & 0xFF;
    response[6] = (g_upgrade_status.total_blocks >> 8) & 0xFF;

    send_firmware_upgrade_response(FIRMWARE_UPGRADE_START, response, 7);
    g_firmware_upgrade_mcu_notify = 1;
    g_in_fw_upgrade = 1;
    return STATUS_SUCCESS;
}

// 处理DATA命令
unsigned char handle_firmware_upgrade_data(uint8_t data[], uint16_t len)
{
    uint8_t response[64] = {0};

    if (g_upgrade_status.state != UPGRADE_PREPARING && g_upgrade_status.state != UPGRADE_TRANSFERRING) {
        debug_printf("Firmware upgrade error: not in PREPARING or TRANSFERRING state (current: %d)\n", g_upgrade_status.state);
        response[0] = STATUS_FAIL;
        send_firmware_upgrade_response(FIRMWARE_UPGRADE_DATA, response, 1);
        return STATUS_FAIL;
    }

    uint16_t block_num = data[2] | (data[3] << 8);
    uint8_t data_len = data[4];

    if (data_len != MAX_DATA_BLOCK_SIZE) {
        debug_printf("Firmware upgrade error: invalid data length %d (must be %d)\n", data_len, MAX_DATA_BLOCK_SIZE);
        response[0] = STATUS_SIZE_INVALID;
        send_firmware_upgrade_response(FIRMWARE_UPGRADE_DATA, response, 1);
        return STATUS_SIZE_INVALID;
    }

    if (block_num != g_upgrade_status.expected_block_num) {
        debug_printf("Firmware upgrade error: block number mismatch (expected: %d, got: %d)\n",
               g_upgrade_status.expected_block_num, block_num);
        response[0] = STATUS_BLOCK_NUM_ERROR;
        response[1] = (g_upgrade_status.expected_block_num >> 0) & 0xFF;
        response[2] = (g_upgrade_status.expected_block_num >> 8) & 0xFF;
        send_firmware_upgrade_response(FIRMWARE_UPGRADE_DATA, response, 3);
        return STATUS_BLOCK_NUM_ERROR;
    }

    uint8_t received_checksum = data[5 + data_len];
    uint8_t calculated_checksum = check_sum(data, 5 + data_len);

    if (received_checksum != calculated_checksum) {
        debug_printf("Firmware upgrade error: checksum mismatch (expected: 0x%02x, got: 0x%02x)\n",
               calculated_checksum, received_checksum);
        response[0] = STATUS_CHECKSUM_ERROR;
        send_firmware_upgrade_response(FIRMWARE_UPGRADE_DATA, response, 1);
        return STATUS_CHECKSUM_ERROR;
    }

    if (g_upgrade_status.state == UPGRADE_PREPARING) {
        g_upgrade_status.state = UPGRADE_TRANSFERRING;

        if (block_num == 0) {
            if (g_first_page_write_executed == 0 && g_flash_erased == 0) {
                debug_printf("Firmware upgrade: erasing existing upgrade image\n");
                dfu_flash_cmd_erase_all();
                g_flash_erased = 1;
            }
        }
    }

    uint16_t buf_offset = g_upgrade_status.page_buf.offset;
    uint16_t remaining = FLASH_PAGE_SIZE - buf_offset;
    uint16_t to_copy = (data_len < remaining) ? data_len : remaining;

    memcpy(&g_upgrade_status.page_buf.buffer[buf_offset], &data[5], to_copy);
    g_upgrade_status.page_buf.offset += to_copy;

    if (g_upgrade_status.page_buf.offset >= FLASH_PAGE_SIZE) {
        unsigned char cmd_data[16];

        if (g_upgrade_status.page_buf.written_pages == 0) {
            debug_printf("Firmware upgrade: writing first page data (flag=0 already executed in START)\n");
        } else {
            cmd_data[0] = 1;
            dfu_flash_cmd_write_page(cmd_data);
        }

        for (int i = 0; i < 4; i++) {
            dfu_flash_cmd_write_page_data(&g_upgrade_status.page_buf.buffer[i * 64]);
        }

        g_upgrade_status.page_buf.written_pages++;
        g_upgrade_status.page_buf.offset = 0;
    }

    if (to_copy < data_len) {
        uint16_t remaining_data = data_len - to_copy;
        memcpy(&g_upgrade_status.page_buf.buffer[0], &data[5 + to_copy], remaining_data);
        g_upgrade_status.page_buf.offset = remaining_data;
    }

    g_upgrade_status.received_blocks++;
    g_upgrade_status.expected_block_num++;

    response[0] = STATUS_SUCCESS;
    response[1] = (block_num >> 0) & 0xFF;
    response[2] = (block_num >> 8) & 0xFF;

    send_firmware_upgrade_response(FIRMWARE_UPGRADE_DATA, response, 3);
    return STATUS_SUCCESS;
}

// 处理END命令
unsigned char handle_firmware_upgrade_end(uint8_t data[], uint16_t len)
{
    uint8_t response[64] = {0};

    if (g_upgrade_status.state != UPGRADE_TRANSFERRING) {
        debug_printf("Firmware upgrade error: not in TRANSFERRING state (current: %d)\n", g_upgrade_status.state);
        response[0] = STATUS_FAIL;
        send_firmware_upgrade_response(FIRMWARE_UPGRADE_END, response, 1);
        return STATUS_FAIL;
    }

    g_upgrade_status.state = UPGRADE_COMPLETING;

    uint32_t firmware_size = data[2] | (data[3] << 8) | (data[4] << 16) | (data[5] << 24);
    debug_printf("Firmware upgrade END: firmware_size=%lu\n", firmware_size);

    if (g_upgrade_status.page_buf.offset > 0) {
        for (uint16_t i = g_upgrade_status.page_buf.offset; i < FLASH_PAGE_SIZE; i++) {
            g_upgrade_status.page_buf.buffer[i] = 0x00;
        }

        unsigned char cmd_data[16];
        cmd_data[0] = 1;
        dfu_flash_cmd_write_page(cmd_data);

        for (int i = 0; i < 4; i++) {
            dfu_flash_cmd_write_page_data(&g_upgrade_status.page_buf.buffer[i * 64]);
        }

        g_upgrade_status.page_buf.written_pages++;
        debug_printf("Firmware upgrade: last page written (total pages: %d)\n",
               g_upgrade_status.page_buf.written_pages);
    }

    unsigned char cmd_data[16];
    cmd_data[0] = 2;
    int result = dfu_flash_cmd_write_page(cmd_data);

    if (result != 0) {
        debug_printf("Firmware upgrade error: flash write termination failed\n");
        g_upgrade_status.state = UPGRADE_ERROR;
        g_upgrade_status.error_code = STATUS_FAIL;
        swlock_release(&flash_lock);
        response[0] = STATUS_FAIL;
        send_firmware_upgrade_response(FIRMWARE_UPGRADE_END, response, 1);
        return STATUS_FAIL;
    }

    debug_printf("Firmware upgrade: image verification passed\n");
    g_upgrade_status.state = UPGRADE_COMPLETED;

    response[0] = STATUS_SUCCESS;
    send_firmware_upgrade_response(FIRMWARE_UPGRADE_END, response, 1);
    debug_printf("Firmware upgrade completed successfully\n");
    g_firmware_upgrade_mcu_notify = 2;
    g_in_fw_upgrade = 0;
    swlock_release(&flash_lock);
    g_upgrade_status.state = UPGRADE_IDLE;
    return STATUS_SUCCESS;
}

// 处理STATUS命令
unsigned char handle_firmware_upgrade_status(uint8_t data[], uint16_t len)
{
    uint8_t response[64] = {0};

    response[0] = (uint8_t)g_upgrade_status.state;
    response[1] = (g_upgrade_status.received_blocks >> 0) & 0xFF;
    response[2] = (g_upgrade_status.received_blocks >> 8) & 0xFF;
    response[3] = (g_upgrade_status.total_blocks >> 0) & 0xFF;
    response[4] = (g_upgrade_status.total_blocks >> 8) & 0xFF;
    response[5] = (g_upgrade_status.page_buf.written_pages >> 0) & 0xFF;
    response[6] = (g_upgrade_status.page_buf.written_pages >> 8) & 0xFF;
    response[7] = (g_upgrade_status.page_buf.total_pages >> 0) & 0xFF;
    response[8] = (g_upgrade_status.page_buf.total_pages >> 8) & 0xFF;

    uint32_t bytes_transferred = g_upgrade_status.received_blocks * MAX_DATA_BLOCK_SIZE;
    response[9]  = (bytes_transferred >> 0) & 0xFF;
    response[10] = (bytes_transferred >> 8) & 0xFF;
    response[11] = (bytes_transferred >> 16) & 0xFF;
    response[12] = (bytes_transferred >> 24) & 0xFF;

    response[13] = (g_upgrade_status.firmware_size >> 0) & 0xFF;
    response[14] = (g_upgrade_status.firmware_size >> 8) & 0xFF;
    response[15] = (g_upgrade_status.firmware_size >> 16) & 0xFF;
    response[16] = (g_upgrade_status.firmware_size >> 24) & 0xFF;

    response[17] = g_upgrade_status.error_code;

    send_firmware_upgrade_response(FIRMWARE_UPGRADE_STATUS, response, 18);
    return STATUS_SUCCESS;
}

// 处理ABORT命令
unsigned char handle_firmware_upgrade_abort(uint8_t data[], uint16_t len)
{
    uint8_t response[64] = {0};

    uint8_t erase_flag = data[2];
    debug_printf("Firmware upgrade ABORT: erase_flag=0x%02x\n", erase_flag);

    if (erase_flag == 0x01) {
        debug_printf("Firmware upgrade: erasing written data\n");
        dfu_flash_cmd_erase_all();
    }

    dfu_flash_cmd_deinit();
    firmware_upgrade_reset();
    g_upgrade_status.state = UPGRADE_ABORTED;

    response[0] = STATUS_SUCCESS;
    send_firmware_upgrade_response(FIRMWARE_UPGRADE_ABORT, response, 1);
    g_firmware_upgrade_mcu_notify = 3;
    g_in_fw_upgrade = 0;
    swlock_release(&flash_lock);
    g_upgrade_status.state = UPGRADE_IDLE;
    return STATUS_SUCCESS;
}

// 处理ERASE命令
unsigned char handle_firmware_upgrade_erase(uint8_t data[], uint16_t len)
{
    uint8_t response[64] = {0};

    debug_printf("Firmware upgrade ERASE: erasing all upgrade images\n");

    swlock_acquire(&flash_lock);

    if (dfu_flash_cmd_init() != 0) {
        debug_printf("Firmware upgrade error: flash init failed\n");
        swlock_release(&flash_lock);
        response[0] = STATUS_FAIL;
        send_firmware_upgrade_response(FIRMWARE_UPGRADE_ERASE, response, 1);
        return STATUS_FAIL;
    }

    if (g_first_page_write_executed == 0 && g_flash_erased == 0) {
        dfu_flash_cmd_erase_all();
        g_flash_erased = 1;
    } else {
        debug_printf("Firmware upgrade ERASE: erase skipped (already erased or begin_write started)\n");
    }

    dfu_flash_cmd_deinit();
    swlock_release(&flash_lock);

    response[0] = STATUS_SUCCESS;
    send_firmware_upgrade_response(FIRMWARE_UPGRADE_ERASE, response, 1);
    return STATUS_SUCCESS;
}

// 固件升级主处理函数（分发0xA7-0xAC，0xAE由eq.c直接处理）
unsigned char process_firmware_upgrade_cmd(uint8_t data[], uint16_t len)
{
    if (data[0] != 0x77) {
        debug_printf("Firmware upgrade error: invalid sync header 0x%02x\n", data[0]);
        return STATUS_FAIL;
    }

    uint8_t cmd = data[1];
    debug_printf("Firmware upgrade command: 0x%02x\n", cmd);

    switch (cmd) {
        case FIRMWARE_UPGRADE_START:
            return handle_firmware_upgrade_start(data, len);
        case FIRMWARE_UPGRADE_DATA:
            return handle_firmware_upgrade_data(data, len);
        case FIRMWARE_UPGRADE_END:
            return handle_firmware_upgrade_end(data, len);
        case FIRMWARE_UPGRADE_STATUS:
            return handle_firmware_upgrade_status(data, len);
        case FIRMWARE_UPGRADE_ABORT:
            return handle_firmware_upgrade_abort(data, len);
        case FIRMWARE_UPGRADE_ERASE:
            return handle_firmware_upgrade_erase(data, len);
        default:
            debug_printf("Firmware upgrade error: unknown command 0x%02x\n", cmd);
            return STATUS_FAIL;
    }
}

#endif // HID_DFU_EN
