/*
 * Custom bootloader for ESP32-S3 that always boots from ota_1 partition
 * with OTA ping-pong update behavior
 * 
 * Based on ESP-IDF v5.4.1 bootloader
 */

#include <stdint.h>
#include <stdbool.h>
#include <stddef.h>
#include "sdkconfig.h"
#include "esp_log.h"
#include "bootloader_init.h"
#include "bootloader_utility.h"
#include "bootloader_common.h"
#include "esp_flash_partitions.h"
#include "esp_image_format.h"

static const char *TAG = "boot";

// OTA_1 partition offset - this is the partition we want to boot from
#define OTA_1_OFFSET 0x210000

/*
 * We arrive here after the ROM bootloader finished loading this second stage bootloader from flash.
 */
void __attribute__((noreturn)) call_start_cpu0(void)
{
    // Initialize hardware
    if (bootloader_init() != ESP_OK) {
        bootloader_reset();
    }

    // Print custom bootloader message
    ESP_LOGI(TAG, "Custom bootloader for ESP32-S3 - Always boot from ota_1");

    // Load partition table
    bootloader_state_t bs = {0};
    if (!bootloader_utility_load_partition_table(&bs)) {
        ESP_LOGE(TAG, "Failed to load partition table");
        bootloader_reset();
    }

    // Find the boot partition index
    int boot_index = -1;
    
    // Always try to boot from ota_1 first by finding the partition at offset 0x210000
    ESP_LOGI(TAG, "Searching for ota_1 partition at offset 0x%x...", OTA_1_OFFSET);
    
    // Iterate through OTA partitions to find ota_1 by offset
    for (int i = 0; i < bs.app_count; i++) {
        if (bs.ota[i].offset == OTA_1_OFFSET) {
            ESP_LOGI(TAG, "Found ota_1 partition at index %d (offset 0x%x)", i, OTA_1_OFFSET);
            boot_index = i;
            break;
        }
    }
    
    // If ota_1 not found by offset, try to find it by subtype
    if (boot_index < 0) {
        ESP_LOGI(TAG, "Searching for ota_1 partition by subtype...");
        // Try to find it in the partition table by checking each OTA slot
        for (int i = 0; i < bs.app_count; i++) {
            // Load the partition info to check its subtype
            esp_partition_pos_t part_pos = bs.ota[i];
            if (part_pos.offset != 0) {
                // Check if this is the OTA_1 partition (subtype 0x11)
                if ((PART_TYPE_APP << 8 | (PART_SUBTYPE_OTA_FLAG | 1)) == 
                    ((uint16_t)PART_TYPE_APP << 8 | (PART_SUBTYPE_OTA_FLAG | 1))) {
                    ESP_LOGI(TAG, "Found ota_1 partition at index %d", i);
                    boot_index = i;
                    break;
                }
            }
        }
    }
    
    // If ota_1 still not found, check if it's directly in the partition table
    if (boot_index < 0 && bs.ota[1].offset != 0) {
        ESP_LOGI(TAG, "Using second OTA slot as ota_1");
        boot_index = 1;  // Assume the second OTA slot is ota_1
    }
    
    // If ota_1 still not found, fall back to default selection
    if (boot_index < 0) {
        ESP_LOGW(TAG, "ota_1 partition not found, falling back to factory app");
        boot_index = FACTORY_INDEX;  // Boot from factory app
    }
    
    // Boot the selected partition
    ESP_LOGI(TAG, "Booting from partition slot %d", boot_index);
    bootloader_utility_load_boot_image(&bs, boot_index);
    
    // Should never reach here
    ESP_LOGE(TAG, "Failed to boot from any partition");
    bootloader_reset();
}