///
/// © MiroZ 2024
///
/// Created on: Apr 23, 2019
/// Modified: 2024
////

#include "App.h"

#include "esp_system.h"
#include "esp_wifi.h"
#include "esp_log.h"
#include "esp_ota_ops.h"
#include "esp_http_client.h"
#include "esp_flash_partitions.h"
#include "esp_partition.h"
#include "string.h"
#include "mbedtls/aes.h"
#include "utilities.h"

#include "Ota.h"

#include <HTTPClient.h>

#define OTA_URL "http://bigfoot-inc.com/fw/wellhub.enc.bin"
extern const uint8_t server_cert[] asm("_binary_bigfoot_inc_pem_start");

static const char *TAG = "OTA";

#define FW_HEADER	32

Ota::Ota(AppIF & app) : m_app(app)
{

}

void Ota::taskFatalError()
{
	ESP_LOGE(TAG, "Exiting task due to fatal error...");
	(void) vTaskDelete(NULL);

	while (1)
	{
		;
	}
}

void Ota::start()
{
	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_LOGW(TAG, "Starting OTA check...");

	HTTPClient client;

	//client.begin(OTA_URL, (const char *)server_cert);
	client.begin(OTA_URL);
	int http_code = client.GET();

	if(http_code != HTTP_CODE_OK)
	{
		ESP_LOGE(TAG, "Server error %d", http_code);
		return;
	}

	int remote_len = client.getSize();
	ESP_LOGI(TAG, "remote file len %d", remote_len);

	if(remote_len < 4096)
	{
		client.end();
		return;
	}

	WiFiClient* stream = client.getStreamPtr();

	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;
	bool image_header_was_checked = false;

	// read the header
	int data_read = stream->readBytes(otaWriteData, FW_HEADER);
	if (data_read != FW_HEADER)
	{
		ESP_LOGE(TAG, "Error: SSL data read error");
		return;
	}

	esp_aes_context m_dec_ctx;
	uint8_t m_dec_iv[16];

	memset(m_dec_iv, 0, sizeof(m_dec_iv));
	esp_aes_init(&m_dec_ctx);
	esp_aes_setkey(&m_dec_ctx,  (const uint8_t*)"thirty two bytes key sdf asdf 32", 256);

	esp_aes_crypt_cbc(&m_dec_ctx, ESP_AES_DECRYPT, 32, m_dec_iv, (const uint8_t*)otaWriteData, (uint8_t*)&otaWriteData[0]);

	if(strncmp((char*)&otaWriteData[0+24], "wellhub0", 8) != 0)
	{
		ESP_LOGE(TAG, "Invalid image!");
		esp_aes_free(&m_dec_ctx);
		client.end();
		return;
	}

	int left = remote_len-FW_HEADER; // encryption header

	while (1)
	{
		int chunk = left > BUFFSIZE ? BUFFSIZE : left;

		int data_read = stream->readBytes(otaWriteData, chunk);
		left -= data_read;

		esp_aes_crypt_cbc(&m_dec_ctx, ESP_AES_DECRYPT, data_read, m_dec_iv, (const uint8_t*)otaWriteData, (uint8_t*)&otaWriteData[0]);

		if (data_read < 0)
		{
			ESP_LOGE(TAG, "Error: SSL data read error");
			return;
		}
		else if (data_read > 0)
		{
			if (image_header_was_checked == false)
			{
				// this is the first 1k of received image
				esp_app_desc_t new_app_info;
				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,
					        &otaWriteData[sizeof(esp_image_header_t) + sizeof(esp_image_segment_header_t)],
					        sizeof(esp_app_desc_t));
					ESP_LOGI(TAG, "Remote firmware");
					ESP_LOGI(TAG, "version: %s", new_app_info.version);
					ESP_LOGI(TAG, "   name: %s", new_app_info.project_name);
					ESP_LOGI(TAG, " c time: %s", new_app_info.time);
					ESP_LOGI(TAG, " c date: %s", new_app_info.date);
					ESP_LOGI(TAG, "idf ver: %s", new_app_info.idf_ver);
					//ESP_LOGI(TAG, "    sha: %s", new_app_info.app_elf_sha256);

					esp_app_desc_t running_app_info;
					if (esp_ota_get_partition_description(running, &running_app_info) == ESP_OK)
					{
						ESP_LOGI(TAG, "");
						ESP_LOGI(TAG, "Running firmware");
						ESP_LOGI(TAG, "version: %s", running_app_info.version);
						ESP_LOGI(TAG, "   name: %s", running_app_info.project_name);
						ESP_LOGI(TAG, " c time: %s", running_app_info.time);
						ESP_LOGI(TAG, " c date: %s", running_app_info.date);
						ESP_LOGI(TAG, "idf ver: %s", running_app_info.idf_ver);
					//	ESP_LOGI(TAG, "    sha: %s", running_app_info.app_elf_sha256);
					}

					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) + sizeof(new_app_info.project_name) + sizeof(new_app_info.time)) == 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.");
							
							client.end();
							
							return;
						}
					}

					// TODO: add date/time check
					if (memcmp(new_app_info.version, running_app_info.version, sizeof(new_app_info.version) + sizeof(new_app_info.project_name) + sizeof(new_app_info.time) + sizeof(new_app_info.date)) == 0)
					{
						ESP_LOGI(TAG, "Current running version is the same as a remote. No need for update.");
						client.end();
						return;
					}

					m_app.getLed()->setPulse(255, 255, 0);

					image_header_was_checked = true;

					ESP_LOGI(TAG, "Starting...");
					err = esp_ota_begin(update_partition, OTA_SIZE_UNKNOWN, &update_handle);

					if (err != ESP_OK)
					{
						ESP_LOGE(TAG, "esp_ota_begin failed (%s)", esp_err_to_name(err));
						// httpCleanup(client);
						taskFatalError();
					}
				}
				else
				{
					ESP_LOGE(TAG, "received package does not fit len");
					return;
				}
			}

			err = esp_ota_write(update_handle, (const void *) otaWriteData, data_read);

			if (err != ESP_OK)
			{
				ESP_LOGE(TAG, "Error: esp_ota_write error");
				taskFatalError();
			}
			binary_file_length += data_read;

			int progress = binary_file_length*100/remote_len;
			static int last_reported = -1;

			if(last_reported != progress)
			{
				last_reported = progress;
				ESP_LOGI(TAG, "%d%%", last_reported);
				delay(1);
			}
		}
		else if (left == 0)
		{
			ESP_LOGI(TAG, "All data received");
			break;
		}
	}

	client.end();

	ESP_LOGI(TAG, "Total Write binary data length : %d", binary_file_length);

	esp_aes_free(&m_dec_ctx);

	if (esp_ota_end(update_handle) != ESP_OK)
	{
		ESP_LOGE(TAG, "esp_ota_end failed!");
		// httpCleanup(client);
		taskFatalError();
	}

	ESP_LOGI(TAG, "esp_ota_end");

	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));
		// httpCleanup(client);
		taskFatalError();
	}
	ESP_LOGI(TAG, "Prepare to restart system!");
	esp_restart();
	return;
}

Ota::~Ota()
{
}