/*
 * OTA 状态管理器实现
 * 管理OTA更新状态，支持失败回退到上一版本
 */

#include "ota_manager.h"
#include "device_nvs.h"
#include "nvs.h"
#include "esp_log.h"
#include "esp_ota_ops.h"
#include "esp_app_desc.h"
#include "esp_partition.h"
#include "esp_system.h"
#include "esp_timer.h"
#include "freertos/FreeRTOS.h"
#include <string.h>
#include <stdio.h>

static const char *tag_s = "OTA_MGR";

#define OTA_MGR_NVS_NAMESPACE    "ota_mgr"
#define OTA_MGR_KEY_STATE        "state"
#define OTA_MGR_KEY_TYPE         "type"
#define OTA_MGR_KEY_TARGET_VER   "tgt_ver"
#define OTA_MGR_KEY_SOURCE_VER   "src_ver"
#define OTA_MGR_KEY_SOURCE_PART  "src_part"
#define OTA_MGR_KEY_TIMESTAMP    "timestamp"
#define OTA_MGR_KEY_BOOT_COUNT   "boot_cnt"
#define OTA_MGR_KEY_MAX_ATTEMPTS "max_retry"
#define OTA_MGR_KEY_LAST_SUCCESS_VER "last_ok_ver"

#define OTA_MAX_BOOT_ATTEMPTS    3        /* 最大启动尝试次数 */
#define OTA_SESSION_MAGIC        0x4F544100u /* "OTA\0" */

static ota_session_info_t current_session_s;
static bool session_loaded_s = false;

/**
 * @brief 从NVS读取会话信息
 */
static esp_err_t load_session_from_nvs(void)
{
    nvs_handle_t handle;
    esp_err_t err = nvs_open(OTA_MGR_NVS_NAMESPACE, NVS_READONLY, &handle);
    if (err != ESP_OK) {
        return err;
    }

    uint32_t magic = 0;
    err = nvs_get_u32(handle, "magic", &magic);
    if (err != ESP_OK || magic != OTA_SESSION_MAGIC) {
        nvs_close(handle);
        memset(&current_session_s, 0, sizeof(current_session_s));
        session_loaded_s = false;
        return ESP_ERR_NOT_FOUND;
    }

    uint8_t state = 0;
    nvs_get_u8(handle, OTA_MGR_KEY_STATE, &state);
    current_session_s.state = (ota_state_t)state;

    uint8_t type = 0;
    nvs_get_u8(handle, OTA_MGR_KEY_TYPE, &type);
    current_session_s.type = (ota_type_t)type;

    size_t len = sizeof(current_session_s.target_version);
    nvs_get_str(handle, OTA_MGR_KEY_TARGET_VER, current_session_s.target_version, &len);

    len = sizeof(current_session_s.source_version);
    nvs_get_str(handle, OTA_MGR_KEY_SOURCE_VER, current_session_s.source_version, &len);

    len = sizeof(current_session_s.source_partition);
    if (nvs_get_str(handle, OTA_MGR_KEY_SOURCE_PART, current_session_s.source_partition, &len) != ESP_OK) {
        current_session_s.source_partition[0] = '\0';
    }

    nvs_get_u32(handle, OTA_MGR_KEY_TIMESTAMP, &current_session_s.timestamp);
    nvs_get_u8(handle, OTA_MGR_KEY_BOOT_COUNT, &current_session_s.boot_count);
    nvs_get_u8(handle, OTA_MGR_KEY_MAX_ATTEMPTS, &current_session_s.max_boot_attempts);

    nvs_close(handle);
    session_loaded_s = true;

    ESP_LOGI(tag_s, "Loaded session: state=%d, type=%d, target=%s, source=%s, boot_cnt=%d",
             current_session_s.state, current_session_s.type,
             current_session_s.target_version, current_session_s.source_version,
             current_session_s.boot_count);

    return ESP_OK;
}

/**
 * @brief 保存会话信息到NVS
 */
static esp_err_t save_session_to_nvs(void)
{
    nvs_handle_t handle;
    esp_err_t err = nvs_open(OTA_MGR_NVS_NAMESPACE, NVS_READWRITE, &handle);
    if (err != ESP_OK) {
        ESP_LOGE(tag_s, "Failed to open NVS: %s", esp_err_to_name(err));
        return err;
    }

    nvs_set_u32(handle, "magic", OTA_SESSION_MAGIC);
    nvs_set_u8(handle, OTA_MGR_KEY_STATE, (uint8_t)current_session_s.state);
    nvs_set_u8(handle, OTA_MGR_KEY_TYPE, (uint8_t)current_session_s.type);
    nvs_set_str(handle, OTA_MGR_KEY_TARGET_VER, current_session_s.target_version);
    nvs_set_str(handle, OTA_MGR_KEY_SOURCE_VER, current_session_s.source_version);
    nvs_set_str(handle, OTA_MGR_KEY_SOURCE_PART, current_session_s.source_partition);
    nvs_set_u32(handle, OTA_MGR_KEY_TIMESTAMP, current_session_s.timestamp);
    nvs_set_u8(handle, OTA_MGR_KEY_BOOT_COUNT, current_session_s.boot_count);
    nvs_set_u8(handle, OTA_MGR_KEY_MAX_ATTEMPTS, current_session_s.max_boot_attempts);

    nvs_commit(handle);
    nvs_close(handle);

    return ESP_OK;
}

/**
 * @brief 清除会话信息
 */
static esp_err_t clear_session_nvs(void)
{
    nvs_handle_t handle;
    esp_err_t err = nvs_open(OTA_MGR_NVS_NAMESPACE, NVS_READWRITE, &handle);
    if (err != ESP_OK) {
        return err;
    }

    nvs_erase_all(handle);
    nvs_commit(handle);
    nvs_close(handle);

    memset(&current_session_s, 0, sizeof(current_session_s));
    session_loaded_s = false;

    ESP_LOGI(tag_s, "Session cleared");
    return ESP_OK;
}

/**
 * @brief 检查当前运行的分区是否是通过OTA更新的分区
 */
static bool is_running_ota_partition(void)
{
    const esp_partition_t *running = esp_ota_get_running_partition();
    if (running == NULL) {
        return false;
    }
    return (strcmp(running->label, "ota_0") == 0 ||
            strcmp(running->label, "ota_1") == 0);
}

static bool partition_is_same(const esp_partition_t *a, const esp_partition_t *b)
{
    if (a == NULL || b == NULL) {
        return false;
    }
    return (a->address == b->address) && (a->size == b->size);
}

/** 传输失败或中断后仍停留在升级前分区（无需切 boot） */
static bool still_on_source_partition(void)
{
    const esp_partition_t *running = esp_ota_get_running_partition();
    if (running == NULL || current_session_s.source_partition[0] == '\0') {
        return false;
    }
    return strcmp(running->label, current_session_s.source_partition) == 0;
}


esp_err_t ota_manager_init(void)
{
    ESP_LOGI(tag_s, "OTA Manager initializing...");

    /* 获取当前固件信息 */
    const esp_partition_t *running = esp_ota_get_running_partition();
    const esp_app_desc_t *app_desc = esp_app_get_description();

    if (running) {
        ESP_LOGI(tag_s, "Running partition: %s", running->label);
    }
    if (app_desc) {
        ESP_LOGI(tag_s, "Running version: %s", app_desc->version);
    }

    /* 从NVS加载会话信息 */
    esp_err_t err = load_session_from_nvs();

    if (err != ESP_OK) {
        /* 没有OTA会话，正常启动 */
        ESP_LOGI(tag_s, "No pending OTA session, normal boot");

        /* 如果运行的是OTA分区，说明是首次烧录或已稳定运行，保存版本信息 */
        if (is_running_ota_partition() && app_desc) {
            ota_manager_update_stored_version(app_desc->version);
        }
        return ESP_OK;
    }

    /* 有OTA会话，检查状态 */
    if (current_session_s.state == OTA_STATE_PENDING_VERIFY) {
        /*
         * esp_ota_end 已校验镜像；能从此分区启动即包有效。
         * 成功条件：运行分区与升级前不同（已从另一 OTA 槽启动）。
         */
        ESP_LOGI(tag_s, "New firmware boot detected, verifying partition switch...");

        current_session_s.boot_count++;
        save_session_to_nvs();

        bool switched = false;
        if (running && current_session_s.source_partition[0] != '\0') {
            switched = (strcmp(running->label, current_session_s.source_partition) != 0);
        } else {
            /* 旧会话无 source_partition：bootloader 能加载即视为成功 */
            switched = (running != NULL);
        }

        if (switched) {
            ESP_LOGI(tag_s, "OTA OK: boot from %s (was %s), version %s",
                     running ? running->label : "?",
                     current_session_s.source_partition[0] ? current_session_s.source_partition : "?",
                     app_desc ? app_desc->version : "?");

            current_session_s.state = OTA_STATE_SUCCESS;
            save_session_to_nvs();
            if (app_desc) {
                ota_manager_update_stored_version(app_desc->version);
            }
        } else {
            ESP_LOGW(tag_s, "OTA verify failed: still on %s",
                     running ? running->label : "unknown");
            clear_session_nvs();
        }

    } else if (current_session_s.state == OTA_STATE_IN_PROGRESS) {
        /*
         * 旧版固件在 UART/BLE 链路 init 时误调 begin_session("unknown") 会留下假会话；
         * 无实际写入时 target=unknown 且 boot_cnt=0，直接清除，避免每次上电 WARN。
         */
        if (current_session_s.boot_count == 0 &&
            strcmp(current_session_s.target_version, "unknown") == 0) {
            ESP_LOGI(tag_s, "Cleared stale OTA session (no transfer; re-flash or erase NVS if WARN persists)");
            clear_session_nvs();
            return ESP_OK;
        }

        /* OTA进行中但设备重启了，可能是断电或崩溃 */
        ESP_LOGW(tag_s, "OTA was in progress but interrupted");

        if (still_on_source_partition()) {
            ESP_LOGI(tag_s, "Still on %s after interrupt, clear session",
                     current_session_s.source_partition);
            clear_session_nvs();
            return ESP_OK;
        }

        current_session_s.state = OTA_STATE_FAILED;
        save_session_to_nvs();

        if (current_session_s.boot_count >= 1) {
            ESP_LOGW(tag_s, "OTA interrupted, triggering rollback");
            ota_manager_rollback();
            return ESP_FAIL;
        }

    } else if (current_session_s.state == OTA_STATE_SUCCESS) {
        /* 新固件已启动但未收到应用层确认，重复重启则回退 */
        current_session_s.boot_count++;
        save_session_to_nvs();

        ESP_LOGI(tag_s, "OTA success pending confirm, boot %u/%u",
                 current_session_s.boot_count,
                 current_session_s.max_boot_attempts);

        if (current_session_s.boot_count >= current_session_s.max_boot_attempts) {
            ESP_LOGW(tag_s, "OTA not confirmed, rolling back");
            ota_manager_rollback();
            return ESP_FAIL;
        }

    } else if (current_session_s.state == OTA_STATE_FAILED) {
        ESP_LOGW(tag_s, "Previous OTA marked as failed");

        if (still_on_source_partition()) {
            ESP_LOGI(tag_s, "Still on pre-OTA partition %s, clear failed session",
                     current_session_s.source_partition);
            clear_session_nvs();
            return ESP_OK;
        }

        const esp_partition_t *boot = esp_ota_get_boot_partition();
        if (running && boot && !partition_is_same(boot, running)) {
            ESP_LOGW(tag_s, "Boot partition %s != running %s, try rollback",
                     boot->label, running->label);
            ota_manager_rollback();
            return ESP_FAIL;
        }

        ESP_LOGI(tag_s, "No rollback needed, clear failed session");
        clear_session_nvs();
    }

    return ESP_OK;
}

esp_err_t ota_manager_begin_session(ota_type_t type, const char *target_version)
{
    if (target_version == NULL || strlen(target_version) == 0) {
        return ESP_ERR_INVALID_ARG;
    }

    const esp_app_desc_t *app_desc = esp_app_get_description();
    const esp_partition_t *running = esp_ota_get_running_partition();

    memset(&current_session_s, 0, sizeof(current_session_s));
    current_session_s.state = OTA_STATE_IN_PROGRESS;
    current_session_s.type = type;
    strncpy(current_session_s.target_version, target_version, sizeof(current_session_s.target_version) - 1);
    if (app_desc) {
        strncpy(current_session_s.source_version, app_desc->version, sizeof(current_session_s.source_version) - 1);
    }
    if (running) {
        strncpy(current_session_s.source_partition, running->label,
                sizeof(current_session_s.source_partition) - 1);
    }
    current_session_s.timestamp = (uint32_t)(esp_timer_get_time() / 1000000ULL);
    current_session_s.boot_count = 0;
    current_session_s.max_boot_attempts = OTA_MAX_BOOT_ATTEMPTS;

    session_loaded_s = true;

    ESP_LOGI(tag_s, "OTA session started: type=%d, %s@%s -> %s",
             type, current_session_s.source_version,
             current_session_s.source_partition, target_version);

    return save_session_to_nvs();
}

esp_err_t ota_manager_end_session(void)
{
    if (!session_loaded_s || current_session_s.state != OTA_STATE_IN_PROGRESS) {
        return ESP_ERR_INVALID_STATE;
    }

    current_session_s.state = OTA_STATE_PENDING_VERIFY;
    ESP_LOGI(tag_s, "OTA session ended, waiting for verification");

    return save_session_to_nvs();
}

esp_err_t ota_manager_mark_success(void)
{
    if (!session_loaded_s) {
        return ESP_OK; /* 没有会话，不需要标记 */
    }

    /* 清除会话，OTA完全成功 */
    const esp_app_desc_t *app_desc = esp_app_get_description();
    if (app_desc) {
        ota_manager_update_stored_version(app_desc->version);
    }

    clear_session_nvs();
    ESP_LOGI(tag_s, "OTA marked as successful");

    return ESP_OK;
}

esp_err_t ota_manager_mark_failed_and_rollback(const char *reason)
{
    if (!session_loaded_s) {
        /* 没有会话，直接回退 */
        return ota_manager_rollback();
    }

    current_session_s.state = OTA_STATE_FAILED;
    save_session_to_nvs();

    ESP_LOGE(tag_s, "OTA marked as failed: %s", reason ? reason : "unknown");

    return ota_manager_rollback();
}

bool ota_manager_should_rollback(void)
{
    if (!session_loaded_s) {
        return false;
    }

    if (current_session_s.state != OTA_STATE_PENDING_VERIFY) {
        return false;
    }

    /* 增加启动计数 */
    current_session_s.boot_count++;
    save_session_to_nvs();

    ESP_LOGW(tag_s, "Boot count: %d/%d",
             current_session_s.boot_count,
             current_session_s.max_boot_attempts);

    /* 如果启动次数超过阈值，应该回退 */
    if (current_session_s.boot_count >= current_session_s.max_boot_attempts) {
        ESP_LOGE(tag_s, "Max boot attempts reached, rollback required");
        return true;
    }

    return false;
}

esp_err_t ota_manager_rollback(void)
{
    ESP_LOGW(tag_s, "Performing rollback...");

    /*
     * 优先使用 ESP-IDF 原生回退流程，避免手工选错分区导致不可启动。
     * 该 API 会在可回退场景下自动标记当前 app 无效并回退到上一个可用镜像。
     */
    esp_err_t rb_err = esp_ota_mark_app_invalid_rollback_and_reboot();
    if (rb_err == ESP_OK) {
        return ESP_OK; /* 不会执行到这里 */
    }
    ESP_LOGW(tag_s, "native rollback unavailable: %s, fallback to manual partition selection",
             esp_err_to_name(rb_err));

    const esp_partition_t *running = esp_ota_get_running_partition();
    if (running == NULL) {
        ESP_LOGE(tag_s, "Cannot determine running partition");
        return ESP_ERR_INVALID_STATE;
    }

    /*
     * 仅作为兜底：若 boot 分区已不是 running，优先回到该分区；
     * 否则在 ota_0/ota_1 之间切换，且明确不能回到当前 running 分区。
     */
    const esp_partition_t *prev_partition = NULL;
    const esp_partition_t *boot = esp_ota_get_boot_partition();
    if (boot && !partition_is_same(boot, running) && boot->type == ESP_PARTITION_TYPE_APP) {
        prev_partition = boot;
    } else if (strcmp(running->label, "ota_0") == 0) {
        prev_partition = esp_partition_find_first(ESP_PARTITION_TYPE_APP, ESP_PARTITION_SUBTYPE_APP_OTA_1, NULL);
    } else if (strcmp(running->label, "ota_1") == 0) {
        prev_partition = esp_partition_find_first(ESP_PARTITION_TYPE_APP, ESP_PARTITION_SUBTYPE_APP_OTA_0, NULL);
    }

    if (prev_partition == NULL) {
        ESP_LOGE(tag_s, "No previous partition found, cannot rollback");
        clear_session_nvs();
        return ESP_ERR_NOT_FOUND;
    }
    if (partition_is_same(prev_partition, running)) {
        ESP_LOGE(tag_s, "Rollback target equals running partition (%s)", running->label);
        return ESP_ERR_INVALID_STATE;
    }

    ESP_LOGI(tag_s, "Rollback to partition: %s", prev_partition->label);

    /* 设置启动分区为上一版本 */
    esp_err_t err = esp_ota_set_boot_partition(prev_partition);
    if (err != ESP_OK) {
        ESP_LOGE(tag_s, "Failed to set boot partition: %s, clear session", esp_err_to_name(err));
        clear_session_nvs();
        return err;
    }

    /* 更新会话状态 */
    if (session_loaded_s) {
        current_session_s.state = OTA_STATE_ROLLBACK;
        save_session_to_nvs();
    }

    ESP_LOGI(tag_s, "Rollback set, rebooting...");
    vTaskDelay(pdMS_TO_TICKS(1000));
    esp_restart();

    return ESP_OK; /* 实际上不会执行到这里 */
}

ota_state_t ota_manager_get_state(void)
{
    if (!session_loaded_s) {
        return OTA_STATE_IDLE;
    }
    return current_session_s.state;
}

esp_err_t ota_manager_get_session_info(ota_session_info_t *info)
{
    if (!session_loaded_s || info == NULL) {
        return ESP_ERR_INVALID_STATE;
    }
    memcpy(info, &current_session_s, sizeof(ota_session_info_t));
    return ESP_OK;
}

const char* ota_manager_get_current_partition_label(void)
{
    const esp_partition_t *running = esp_ota_get_running_partition();
    if (running == NULL) {
        return NULL;
    }
    return running->label;
}

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

    nvs_handle_t handle;
    esp_err_t err = nvs_open(OTA_MGR_NVS_NAMESPACE, NVS_READONLY, &handle);
    if (err != ESP_OK) {
        return err;
    }

    size_t required_len = len;
    err = nvs_get_str(handle, OTA_MGR_KEY_LAST_SUCCESS_VER, buf, &required_len);
    nvs_close(handle);

    if (err == ESP_ERR_NVS_NOT_FOUND) {
        buf[0] = '\0';
        return ESP_OK;
    }

    return err;
}

esp_err_t ota_manager_update_stored_version(const char *version)
{
    if (version == NULL) {
        return ESP_ERR_INVALID_ARG;
    }

    nvs_handle_t handle;
    esp_err_t err = nvs_open(OTA_MGR_NVS_NAMESPACE, NVS_READWRITE, &handle);
    if (err != ESP_OK) {
        return err;
    }

    err = nvs_set_str(handle, OTA_MGR_KEY_LAST_SUCCESS_VER, version);
    if (err == ESP_OK) {
        nvs_commit(handle);
    }
    nvs_close(handle);

    ESP_LOGI(tag_s, "Stored version updated: %s", version);
    return err;
}
