/*******************************************************************************
 * @file    dfu_manager.c
 * @brief   DFU state, watchdog, and MCUboot image control
 ******************************************************************************/
#include <zephyr/kernel.h>
#include <zephyr/dfu/mcuboot.h>
#include <zephyr/mgmt/mcumgr/mgmt/callbacks.h>
#include <hal/nrf_power.h>

#include "main.h"
#include "dfu_manager.h"
#include "debug_print.h"
#include "led_control.h"
#include "ble_service.h"

#define DFU_STALL_TIMEOUT_SEC   30      // DFU upload stall disconnect threshold
#define IMG_MGMT_ID_STATE       0U
#define IMG_MGMT_ID_UPLOAD      1U

static bool dfu_led_active;                       // DFU upload green LED state
static struct k_work_delayable dfu_watchdog_work; // DFU upload stall watchdog

static void dfu_watchdog_handler(struct k_work *work);

static enum mgmt_cb_return dfu_status_cb(uint32_t event,
                                         enum mgmt_cb_return prev_status,
                                         int32_t *rc,
                                         uint16_t *group,
                                         bool *abort_more,
                                         void *data,
                                         size_t data_size);

static enum mgmt_cb_return smp_cmd_status_cb(uint32_t event,
                                             enum mgmt_cb_return prev_status,
                                             int32_t *rc,
                                             uint16_t *group,
                                             bool *abort_more,
                                             void *data,
                                             size_t data_size);

static struct mgmt_callback dfu_status_callback = {
    .callback = dfu_status_cb,
    .event_id = MGMT_EVT_OP_IMG_MGMT_DFU_STARTED |
                MGMT_EVT_OP_IMG_MGMT_DFU_CHUNK |
                MGMT_EVT_OP_IMG_MGMT_DFU_CHUNK_WRITE_COMPLETE |
                MGMT_EVT_OP_IMG_MGMT_DFU_PENDING |
                MGMT_EVT_OP_IMG_MGMT_DFU_STOPPED,
};

static struct mgmt_callback smp_cmd_status_callback = {
    .callback = smp_cmd_status_cb,
    .event_id = MGMT_EVT_OP_CMD_DONE,
};

/* DFU watchdog와 mcumgr 콜백 초기화 */
void dfu_manager_init(void)
{
    k_work_init_delayable(&dfu_watchdog_work, dfu_watchdog_handler);
    mgmt_callback_register(&dfu_status_callback);
    mgmt_callback_register(&smp_cmd_status_callback);
}

/* 현재 이미지가 confirm 대기 상태인지 확인 */
bool dfu_is_confirm_pending_boot(void)
{
#if defined(CONFIG_BOOTLOADER_MCUBOOT)
    return !boot_is_img_confirmed();
#else
    return false;
#endif
}

/* 실행 중 이미지 confirm 처리 */
void dfu_confirm_running_image(void)
{
#if defined(CONFIG_BOOTLOADER_MCUBOOT)
    if (!boot_is_img_confirmed())
    {
        int err = boot_write_img_confirmed();
        if (err == 0)
        {
            DBG_CORE("[DFU] Running image confirmed\r\n");
        }
        else
        {
            DBG_ERR("[DFU] Image confirm failed (err %d)\r\n", err);
        }
    }
    else
    {
        DBG_CORE("[DFU] Running image already confirmed\r\n");
    }
#endif
}

/* MCUboot swap 상태 문자열 변환 */
static const char *swap_type_to_str(int swap_type)
{
    switch (swap_type)
    {
        case BOOT_SWAP_TYPE_NONE:
            return "none";
        case BOOT_SWAP_TYPE_TEST:
            return "test";
        case BOOT_SWAP_TYPE_PERM:
            return "perm";
        case BOOT_SWAP_TYPE_REVERT:
            return "revert";
        case BOOT_SWAP_TYPE_FAIL:
            return "fail";
        default:
            return "unknown";
    }
}

/* MCUboot 이미지 상태 로그 출력 */
void dfu_log_mcuboot_state(void)
{
#if defined(CONFIG_BOOTLOADER_MCUBOOT)
    struct mcuboot_img_header header;
    uint8_t active_slot = boot_fetch_active_slot();
    int swap_type = mcuboot_swap_type();
    int err = boot_read_bank_header(active_slot, &header, sizeof(header));

    DBG_CORE("[DFU] active_slot=%u swap_type=%s(%d) confirmed=%d\r\n",
             active_slot,
             swap_type_to_str(swap_type),
             swap_type,
             boot_is_img_confirmed());

    if (err == 0)
    {
        DBG_CORE("[DFU] image version %u.%u.%u+%u size=0x%x\r\n",
                 header.h.v1.sem_ver.major,
                 header.h.v1.sem_ver.minor,
                 header.h.v1.sem_ver.revision,
                 header.h.v1.sem_ver.build_num,
                 header.h.v1.image_size);
    }
    else
    {
        DBG_ERR("[DFU] image header read failed (slot %u, err %d)\r\n",
                active_slot,
                err);
    }
#endif
}

/* DFU 업로드 무진행 시 BLE 연결 강제 해제 */
static void dfu_watchdog_handler(struct k_work *work)
{
    ARG_UNUSED(work);

    if (ble_is_connected())
    {
        DBG_CORE("[DFU] upload stalled (%ds), forcing disconnect\r\n", DFU_STALL_TIMEOUT_SEC);
        (void)ble_disconnect_active();
    }
}

/* DFU 업로드 상태에 따라 LED와 watchdog 제어 */
static enum mgmt_cb_return dfu_status_cb(uint32_t event, enum mgmt_cb_return prev_status, int32_t *rc, uint16_t *group, bool *abort_more, void *data, size_t data_size)
{
    ARG_UNUSED(prev_status);
    ARG_UNUSED(rc);
    ARG_UNUSED(group);
    ARG_UNUSED(abort_more);
    ARG_UNUSED(data);
    ARG_UNUSED(data_size);

    switch (event)
    {
        case MGMT_EVT_OP_IMG_MGMT_DFU_STARTED:
        case MGMT_EVT_OP_IMG_MGMT_DFU_CHUNK:
        case MGMT_EVT_OP_IMG_MGMT_DFU_CHUNK_WRITE_COMPLETE:
            dfu_led_active = true;
            led_ble_solid();
            k_work_reschedule(&dfu_watchdog_work, K_SECONDS(DFU_STALL_TIMEOUT_SEC));
            break;

        case MGMT_EVT_OP_IMG_MGMT_DFU_PENDING:
            dfu_led_active = true;
            led_ble_solid();
            k_work_cancel_delayable(&dfu_watchdog_work);

#if NRF_POWER_HAS_GPREGRET
            nrf_power_gpregret_set(NRF_POWER, DFU_RESUME_GPREGRET_REG, DFU_RESUME_MAGIC);
#endif
            break;

        case MGMT_EVT_OP_IMG_MGMT_DFU_STOPPED:
            k_work_cancel_delayable(&dfu_watchdog_work);
            if (dfu_led_active)
            {
                dfu_led_active = false;
                led_set_state(LED_STATE_ADVERTISING);
            }
            break;

        default:
            break;
    }

    return MGMT_CB_OK;
}

/* DFU 명령 실패 시 LED와 resume 플래그 복구 */
static enum mgmt_cb_return smp_cmd_status_cb(uint32_t event, enum mgmt_cb_return prev_status, int32_t *rc, uint16_t *group, bool *abort_more, void *data, size_t data_size)
{
    ARG_UNUSED(event);
    ARG_UNUSED(prev_status);
    ARG_UNUSED(rc);
    ARG_UNUSED(group);
    ARG_UNUSED(abort_more);

    if (data == NULL || data_size != sizeof(struct mgmt_evt_op_cmd_arg))
    {
        return MGMT_CB_OK;
    }

    const struct mgmt_evt_op_cmd_arg *cmd = data;

    if (cmd->group == MGMT_GROUP_ID_IMAGE &&
        (cmd->id == IMG_MGMT_ID_UPLOAD || cmd->id == IMG_MGMT_ID_STATE) &&
        cmd->err != MGMT_ERR_EOK)
    {
        DBG_ERR("[DFU] command failed id=%u err=%d\r\n", cmd->id, cmd->err);

        dfu_led_active = false;
        led_set_state(LED_STATE_ADVERTISING);

#if NRF_POWER_HAS_GPREGRET
        nrf_power_gpregret_set(NRF_POWER, DFU_RESUME_GPREGRET_REG, 0U);
#endif
    }

    return MGMT_CB_OK;
}