#include <stdbool.h>
#include "sdkconfig.h"
#include "esp_log.h"
#include "bootloader_init.h"
#include "bootloader_utility.h"
#include "bootloader_common.h"
#include "esp_partition.h" 
#include "esp_app_desc.h" 
#include "string.h"


#define MAX_PARTITIONS  8
#define LABEL_LENGTH    32

/**
 * @brief Helper macro to test the equality of two partition positions
 */
#define PART_POS_EQUAL(p0, p1)      (((p0).offset == (p1).offset) && ((p1).size == (p1).size))

static const char* TAG = "boot";
/**
 * @brief Structure describing a partition, with its name label and its position, used as an identifier here.
 */
typedef struct {
    const esp_partition_pos_t* pos;
    int number;
    char label[LABEL_LENGTH];
} bootloader_part_info;



extern esp_err_t bootloader_common_get_partition_description(const esp_partition_pos_t *partition, esp_app_desc_t *app_desc);

static int bootloader_utility_list_partitions(const bootloader_state_t* st, bootloader_part_info *bp, int length);
static int find_ota1_partition(bootloader_part_info *partitions, int count);




/*
 * We arrive here after the ROM bootloader finished loading this second stage bootloader from flash.
 * The hardware is mostly uninitialized, flash cache is down and the app CPU is in reset.
 * We do have a stack, so we can do the initialization in C.
 */
void __attribute__((noreturn)) call_start_cpu0(void)
{
     bootloader_state_t bs = {0};
     bootloader_part_info partitions[MAX_PARTITIONS];
     int ota1_index = -1;
     
    // 1. Hardware initialization
    if (bootloader_init() != ESP_OK) {
        esp_rom_printf("[%s] Hardware initialization failed, resetting...\n", TAG);
        bootloader_reset();
    }
    esp_rom_printf("[%s] Hardware initialization successful\n", TAG);

#ifdef CONFIG_BOOTLOADER_SKIP_VALIDATE_IN_DEEP_SLEEP
    bootloader_utility_load_boot_image_from_deep_sleep();
    esp_rom_printf("[%s] Loaded boot image from deep sleep\n", TAG);
#endif

      // 2. Load partition table
    if (!bootloader_utility_load_partition_table(&bs)) {
        ESP_LOGE(TAG, "load partition table error!");
        bootloader_reset();
    }

    // 3. Show the available partitions menu
    int count = bootloader_utility_list_partitions(&bs, partitions, MAX_PARTITIONS);
    if (count == 0) {
        esp_rom_printf("No boot partition available, rebooting...\n");
        esp_rom_delay_us(1000000);
        bootloader_reset();
    }
    
    // 4. Find OTA_1 partition
    ota1_index = find_ota1_partition(partitions, count);
    if (ota1_index == -1) {
        esp_rom_printf("OTA_1 partition not found, trying to boot from available partitions...\n");
        // Fallback to first available partition
        ota1_index = 0;
    }

    // 5. Print a custom message
    esp_rom_printf("[%s] %s\n", TAG, CONFIG_EXAMPLE_BOOTLOADER_WELCOME_MESSAGE);

    esp_rom_printf("BOOTING PARTITION %s (index: %d)\n", partitions[ota1_index].label, ota1_index);

    // 6. Load the app image for booting
    esp_rom_printf("Loading app from partition at offset 0x%x\n", (unsigned int)partitions[ota1_index].pos->offset);
    
    // Direct approach: Check if we found OTA_1 and load it specifically
    if (partitions[ota1_index].pos->offset == 0x010000 ) {
        esp_rom_printf("Loading OTA_1 partition directly from bootloader state\n");
        // Load OTA_1 directly using the bootloader state
        // OTA_1 is bs.ota[1], so we use bootloader utility with proper OTA index
        bootloader_utility_load_boot_image(&bs, 1); // OTA_1 is at index 1 in ota array
    } else {
        esp_rom_printf("Loading partition using standard method\n");
        bootloader_utility_load_boot_image(&bs, partitions[ota1_index].number);
    }
}

/**
 * @brief Find OTA_1 partition in the partition list
 * This function looks for OTA_1 partition either by checking if it's the second OTA partition (ota[1])
 * or by matching partition offset (0x010000  based on your partition table)
 */
static int find_ota1_partition(bootloader_part_info *partitions, int count)
{
    for (int i = 0; i < count; i++) {
        // Method 1: Check if this is OTA partition number 1 (second OTA partition)
        if (partitions[i].number == 1 && partitions[i].pos->offset >= 0x200000) {
            esp_rom_printf("Found OTA_1 partition by OTA number at index %d\n", i);
            return i;
        }
        
        // Method 2: Check by partition offset (0x010000  for ota_1 based on your table)
        if (partitions[i].pos->offset == 0x010000 ) {
            esp_rom_printf("Found OTA_1 partition by offset (0x010000 ) at index %d\n", i);
            return i;
        }
        
        // Method 3: Check by label if available (some apps might have descriptive names)
        // This is less reliable but can be used as additional check
        if (strstr(partitions[i].label, "ota") || strstr(partitions[i].label, "OTA")) {
            if (partitions[i].pos->offset == 0x010000 ) {
                esp_rom_printf("Found OTA_1 partition by label and offset at index %d\n", i);
                return i;
            }
        }
    }
    
    esp_rom_printf("OTA_1 partition not found in partition list\n");
    return -1;
}

static esp_err_t bootloader_utility_fill_info(const esp_partition_pos_t* partition, bootloader_part_info* info)
{
    esp_image_metadata_t image_data = {0};
    esp_app_desc_t app_desc;
    esp_err_t ret;

    ret = esp_image_verify(ESP_IMAGE_VERIFY_SILENT, partition, &image_data);

    if (ret == ESP_OK) {
        ret = bootloader_common_get_partition_description(partition, &app_desc);
    }

    if (ret == ESP_OK) {
        info->pos = partition;
        /* Make sure to always have a NULL-byte */
        strncpy(info->label, app_desc.project_name, LABEL_LENGTH - 1);
        info->label[LABEL_LENGTH - 1] = 0;
    }

    return ret;
}

/**
 * @brief Get bootable partitions information.
 * We could override the function `bootloader_utility_load_partition_table` to return these information at the same time
 * when loading and populating `bootloader_state_t` structure. But let's keep it simple and avoid copy-paste.
 */
static int bootloader_utility_list_partitions(const bootloader_state_t* st, bootloader_part_info *bp, int length)
{
    int index = 0;

    // Add factory partition if available and if `bp` is big enough
    if (st->factory.offset && index < length) {
        if (bootloader_utility_fill_info(&st->factory, bp + index) == ESP_OK) {
            bp[index++].number = FACTORY_INDEX;
            esp_rom_printf("Added factory partition at index %d\n", index-1);
        } else {
            esp_rom_printf("Error getting the description for the factory partition\n");
        }
    }

    // Add test partition if available
    if (st->test.offset && index < length) {
        if (bootloader_utility_fill_info(&st->test, bp + index) == ESP_OK) {
            bp[index++].number = TEST_APP_INDEX;
            esp_rom_printf("Added test partition at index %d\n", index-1);
        } else {
            esp_rom_printf("Error getting the description for the test partition\n");
        }
    }

    // Add OTA partitions - this is where ota_0 and ota_1 will be added
    for (int i = 0; i < st->app_count && index < length; i++) {
        if (bootloader_utility_fill_info(&st->ota[i], bp + index) == ESP_OK) {
            bp[index].number = i;  // This will be 0 for ota_0, 1 for ota_1, etc.
            esp_rom_printf("Added OTA_%d partition at index %d (offset: 0x%x)\n", 
                          i, index, (unsigned int)st->ota[i].offset);
            index++;
        } else {
            esp_rom_printf("Error getting the description for the ota_%d partition\n", i);
        }
    }

    return index;
}

// Return global reent struct if any newlib functions are linked to bootloader
struct _reent *__getreent(void)
{
    return _GLOBAL_REENT;
}