#define OTA_c
#include "header.h"
#include "OTA.h"
#define ON 1
#define OFF 0

#define HASH_LEN 32

static const int RX_BUF_SIZE = 1024;

bool OTA_status= OFF;

void infinite_loop(void)
{
    static const char *TAG = "OTA Function\t";
    int i = 0;
    ESP_LOGI(TAG, "When a new firmware is available on the server, press the reset button to download it");
    while(1) {
        ESP_LOGI(TAG, "Waiting for a new firmware ... %d", ++i);
        vTaskDelay(2000 / portTICK_PERIOD_MS);
    }
}

void OTA()
{
    char ota_write_data[1025]={0};
    static const char *TAG = "OTA Function\t";
    esp_err_t err;
    /* update handle : set by esp_ota_begin(), must be freed via esp_ota_end() */
    esp_ota_handle_t update_handle = 0 ;
    const esp_partition_t *update_partition = NULL;

    ESP_LOGI(TAG, "Starting OTA Download ....");

    const esp_partition_t *configured = esp_ota_get_boot_partition();
    const esp_partition_t *running = esp_ota_get_running_partition();

    if (configured != running) {
        ESP_LOGW(TAG, "Configured OTA boot partition at offset 0x%08x, but running from offset 0x%08x",
                 configured->address, running->address);
        ESP_LOGW(TAG, "(This can happen if either the OTA boot data or preferred boot image become corrupted somehow.)");
    }
    ESP_LOGI(TAG, "Running partition type %d subtype %d (offset 0x%08x)",
             running->type, running->subtype, running->address);

    
    update_partition = esp_ota_get_next_update_partition(NULL);
    ESP_LOGI(TAG, "Writing to partition subtype %d at offset 0x%x",
             update_partition->subtype, update_partition->address);
    assert(update_partition != NULL);

    int binary_file_length = 0;
    /*deal with all receive packet*/
    bool image_header_was_checked = false;    
    
    printf("\nOTA Status %d\n",OTA_status);
    bluetooth_write("$OTA::NEXT;#\n");
    char *l_ptr=NULL;
    int count=0,count2=0;
    int data_read=0;
    while (1) {
        //while(bt_flag==0);
        while(bt_flag==0)
        {
            count++;
            if(count>1000)
            {
                vTaskDelay(10 / portTICK_PERIOD_MS);
                printf(".");
                count=0;
            }
        }
        if(bt_flag==1)
        {
            count2++;
            bt_flag=0;
            printf("\nPacket No \t# %d\n",count2);
            
            data_read = data_count;
            printf("Data_read %d\n",data_read);
            strcpy(ota_write_data,bt_buf0);     
  
            l_ptr=strstr((char*)bt_buf0,"OTA::END;");
            if(l_ptr !=0)
            {
                printf("OTA Finished\n");
                data_count=data_read=-1;
                strcpy((char *)bt_buf0,"");
                break;
            }            
            else if (data_read > 0) {
                if (image_header_was_checked == false) {
                    esp_app_desc_t new_app_info;
                    printf("%d\n%d\n%d\n%d",data_read, sizeof(esp_image_header_t), sizeof(esp_image_segment_header_t),sizeof(esp_app_desc_t));
                    if (data_read > sizeof(esp_image_header_t) + sizeof(esp_image_segment_header_t) + sizeof(esp_app_desc_t)) {
                        // check current version with downloading
                        memcpy(&new_app_info, &ota_write_data[sizeof(esp_image_header_t) + sizeof(esp_image_segment_header_t)], sizeof(esp_app_desc_t));
                        ESP_LOGI(TAG, "New firmware version: %s", new_app_info.version);

                        esp_app_desc_t running_app_info;
                        if (esp_ota_get_partition_description(running, &running_app_info) == ESP_OK) {
                            ESP_LOGI(TAG, "Running firmware version: %s", running_app_info.version);
                        }

                        const esp_partition_t* last_invalid_app = esp_ota_get_last_invalid_partition();
                        esp_app_desc_t invalid_app_info;
                        if (esp_ota_get_partition_description(last_invalid_app, &invalid_app_info) == ESP_OK) {
                            ESP_LOGI(TAG, "Last invalid firmware version: %s", invalid_app_info.version);
                        }

                        // check current version with last invalid partition
                        if (last_invalid_app != NULL) {
                            if (memcmp(invalid_app_info.version, new_app_info.version, sizeof(new_app_info.version)) == 0) {
                                ESP_LOGW(TAG, "New version is the same as invalid version.");
                                ESP_LOGW(TAG, "Previously, there was an attempt to launch the firmware with %s version, but it failed.", invalid_app_info.version);
                                ESP_LOGW(TAG, "The firmware has been rolled back to the previous version.");
                                // Leave OTA
                                OTA_status=OFF;
                                return;
                            }
                        }

                        if (memcmp(new_app_info.version, running_app_info.version, sizeof(new_app_info.version)) == 0) {
                            ESP_LOGW(TAG, "Current running version is the same as a new. We will not continue the update.");
                            bluetooth_write("OTA::UPDATED;\n");
                            OTA_status=OFF;
                            return;
                        }

                        image_header_was_checked = true;

                        err = esp_ota_begin(update_partition, OTA_SIZE_UNKNOWN, &update_handle);
                        vTaskDelay(100 / portTICK_PERIOD_MS);
                        if (err != ESP_OK) {
                            ESP_LOGE(TAG, "esp_ota_begin failed (%s)", esp_err_to_name(err));
                            // Return
                            bluetooth_write("OTA::FAILED;\n");
                            OTA_status=OFF;
                            return;
                        }
                        ESP_LOGI(TAG, "esp_ota_begin succeeded");
                    } 
                    else {
                        ESP_LOGE(TAG, "received package is not fit len");    
                        
                    }
                }
                bluetooth_write("$OTA::NEXT;#\n");
                err = esp_ota_write( update_handle, (const void *)ota_write_data, data_read);
                vTaskDelay(100 / portTICK_PERIOD_MS);
                if (err != ESP_OK) {
                    bluetooth_write("OTA::FAILED;\n");
                    OTA_status=OFF;
                    return;
                }
                binary_file_length += data_read;
                ESP_LOGD(TAG, "Written image length %d", binary_file_length);
            } else if (data_read == 0) {
                ESP_LOGI(TAG, "Connection closed,all data received");
                break;
            }
        }
    }
    ESP_LOGI(TAG, "Total Write binary data length : %d", binary_file_length);

    if (esp_ota_end(update_handle) != ESP_OK) {
        ESP_LOGE(TAG, "esp_ota_end failed!");
        OTA_status=OFF;
        return ;
    }

    err = esp_ota_set_boot_partition(update_partition);
    if (err != ESP_OK) {
        ESP_LOGE(TAG, "esp_ota_set_boot_partition failed (%s)!", esp_err_to_name(err));
        OTA_status=OFF;
        return ;
    }
    else{
	ESP_LOGI(TAG, "Prepare to restart system!");
        bluetooth_write("\rPrepare to restart system!\n");
        esp_restart();
    }

    OTA_status=OFF;
    return ;
}

void print_sha256 (const uint8_t *image_hash, const char *label)
{
    static const char *TAG = "OTA sha256\t";
    char hash_print[HASH_LEN * 2 + 1];
    hash_print[HASH_LEN * 2] = 0;
    for (int i = 0; i < HASH_LEN; ++i) {
        sprintf(&hash_print[i * 2], "%02x", image_hash[i]);
    }
    ESP_LOGI(TAG, "%s: %s", label, hash_print);
}

bool diagnostic(void)
{
    static const char *TAG = "OTA diagnostic\t";
    gpio_config_t io_conf;
    io_conf.intr_type    = GPIO_PIN_INTR_DISABLE;
    io_conf.mode         = GPIO_MODE_INPUT;
    io_conf.pin_bit_mask = (1ULL << CONFIG_GPIO_DIAGNOSTIC);
    io_conf.pull_down_en = GPIO_PULLDOWN_DISABLE;
    io_conf.pull_up_en   = GPIO_PULLUP_ENABLE;
    gpio_config(&io_conf);

    ESP_LOGI(TAG, "Diagnostics (5 sec)...");
    vTaskDelay(5000 / portTICK_PERIOD_MS);

    bool diagnostic_is_ok = gpio_get_level(CONFIG_GPIO_DIAGNOSTIC);

    gpio_reset_pin(CONFIG_GPIO_DIAGNOSTIC);
    return diagnostic_is_ok;
}

void OTA_init()
{
     static const char *TAG = "OTA Init\t";
    uint8_t sha_256[HASH_LEN] = { 0 };
    esp_partition_t partition;

    // get sha256 digest for the partition table
    partition.address   = ESP_PARTITION_TABLE_OFFSET;
    partition.size      = ESP_PARTITION_TABLE_MAX_LEN;
    partition.type      = ESP_PARTITION_TYPE_DATA;
    esp_partition_get_sha256(&partition, sha_256);
    print_sha256(sha_256, "SHA-256 for the partition table: ");

    // get sha256 digest for bootloader
    partition.address   = ESP_BOOTLOADER_OFFSET;
    partition.size      = ESP_PARTITION_TABLE_OFFSET;
    partition.type      = ESP_PARTITION_TYPE_APP;
    esp_partition_get_sha256(&partition, sha_256);
    print_sha256(sha_256, "SHA-256 for bootloader: ");

    // get sha256 digest for running partition
    esp_partition_get_sha256(esp_ota_get_running_partition(), sha_256);
    print_sha256(sha_256, "SHA-256 for current firmware: ");

    const esp_partition_t *running = esp_ota_get_running_partition();
    esp_ota_img_states_t ota_state;
    if (esp_ota_get_state_partition(running, &ota_state) == ESP_OK) {
        if (ota_state == ESP_OTA_IMG_PENDING_VERIFY) {
            // run diagnostic function ...
            bool diagnostic_is_ok = diagnostic();
            if (diagnostic_is_ok) {
                ESP_LOGI(TAG, "Diagnostics completed successfully! Continuing execution ...");
                esp_ota_mark_app_valid_cancel_rollback();
            } else {
                ESP_LOGE(TAG, "Diagnostics failed! Start rollback to the previous version ...");
                esp_ota_mark_app_invalid_rollback_and_reboot();
            }
        }
    }

    // Initialize NVS.
    esp_err_t err = nvs_flash_init();
    if (err == ESP_ERR_NVS_NO_FREE_PAGES || err == ESP_ERR_NVS_NEW_VERSION_FOUND) {
        // OTA app partition table has a smaller NVS partition size than the non-OTA
        // partition table. This size mismatch may cause NVS initialization to fail.
        // If this happens, we erase NVS partition and initialize NVS again.
        ESP_ERROR_CHECK(nvs_flash_erase());
        err = nvs_flash_init();
    }
    ESP_ERROR_CHECK( err );
}