/**
* @file Mesh_OTA.c
* @brief Start and implement OTA updates via HTTPS from server and other mesh nodes (bidirectional)
* @author Hendrik Schutter
* @date 21.01.2021
*/

#include "Mesh_OTA.h"
#include "Mesh_OTA_Util.h"
#include "Mesh_OTA_Globals.h"
#include "Mesh_OTA_Partition_Access.h"

static const char *LOG_TAG = "mesh_ota";

/**
* @fn esp_err_t errMeshOTAInitialize(void)
* @brief Starts Mesh OTA functionality
* @param void
* @return ESP32 error code
* @author Hendrik Schutter
* @date 21.01.2021
*
* Initialize queues and tasks
* Set callbacks
*/
esp_err_t errMeshOTAInitialize(void)
{
    esp_err_t err = ESP_OK;
    BaseType_t xReturned;
    bWantReboot = false;

    //create queue to store nodes for ota worker task
    queueNodes = xQueueCreate(QUEUE_NODES_SIZE, sizeof(mesh_addr_t));
    if (queueNodes == 0)  // Queue not created
        {
            ESP_LOGE(LOG_TAG, "Unable to create queue for nodes");
            err = ESP_FAIL;
        }

    if(err == ESP_OK)
        {
            //create queue to store ota messages
            queueMessageOTA = xQueueCreate(QUEUE_MESSAGE_OTA_SIZE, sizeof(MESH_PACKET_t));
            if (queueMessageOTA == 0)  // Queue not created
                {
                    ESP_LOGE(LOG_TAG, "Unable to create queue for OTA messages");
                    err = ESP_FAIL;
                }
        }

    if(err == ESP_OK)
        {
            bsStartStopServerWorker = xSemaphoreCreateBinary();
            if( bsStartStopServerWorker == NULL )
                {
                    ESP_LOGE(LOG_TAG, "Unable to create mutex to represent state of server worker");
                    err = ESP_FAIL;
                }
        }

    if(err == ESP_OK)
        {
            bsOTAProcess = xSemaphoreCreateBinary();
            if( bsOTAProcess == NULL )
                {
                    ESP_LOGE(LOG_TAG, "Unable to create mutex to grant access to OTA process");
                    err = ESP_FAIL;
                }
        }

    if(err == ESP_OK)
        {
            xSemaphoreGive(bsOTAProcess);  //unlock binary semaphore
            if( bsOTAProcess == NULL )
                {
                    ESP_LOGE(LOG_TAG, "Unable to unlock mutex to grant access to OTA process");
                    err = ESP_FAIL;
                }
        }

    //register callbacks in network
    ERROR_CHECK(errMeshNetworkSetChildConnectedHandle(vMeshOtaUtilAddNodeToPossibleUpdatableQueue));
    ERROR_CHECK(errMeshNetworkSetOTAMessageHandleHandle(vMeshOtaUtilAddOtaMessageToQueue));
    ERROR_CHECK(errMeshNetworkSetChangeStateOfServerWorkerHandle(vMeshOtaUtilChangeStateOfServerWorker));

    if(err == ESP_OK)
        {
            pOTAPartition =  esp_ota_get_next_update_partition(NULL); //get ota partition

            if(pOTAPartition == NULL)
                {
                    err = ESP_FAIL;
                    ESP_LOGE(LOG_TAG, "unable to get next ota partition");
                }
        }

    if(err == ESP_OK)
        {
            xReturned = xTaskCreate(vMeshOtaTaskServerWorker, "vMeshOtaTaskServerWorker", 8192, NULL, 5, NULL);
            if(xReturned != pdPASS)
                {
                    ESP_LOGE(LOG_TAG, "Unable to create the server worker task");
                    err = ESP_FAIL;
                }
        }


    if(err == ESP_OK)
        {
            xReturned = xTaskCreate(vMeshOtaTaskOTAWorker, "vMeshOtaTaskOTAWorker", 8192, NULL, 5, NULL);
            if(xReturned != pdPASS)
                {
                    ESP_LOGE(LOG_TAG, "Unable to create the OTA worker task");
                    err = ESP_FAIL;
                }
        }
    return err;
}

/**
* @fn void vMeshOtaTaskServerWorker(void *arg)
* @brief Task for updating from server via HTTPS
* @param arg
* @return void
* @author Hendrik Schutter
* @date 21.01.2021
*/
void vMeshOtaTaskServerWorker(void *arg)
{
    esp_err_t err = ESP_OK;
    bool bNewOTAImage; //true if a new ota image was downloaded and validated
    bool bFirstRun = true;

    while(true)
        {
            err = ESP_OK;
            bNewOTAImage = false;
            xSemaphoreTake(bsStartStopServerWorker, portMAX_DELAY); //wait for binary semaphore that allows to start the worker
            xSemaphoreGive(bsStartStopServerWorker); //free binary semaphore, this allows the handler to change is to taken

            if (esp_mesh_is_root()) //check again that this node is the root node
                {
                    ESP_LOGI(LOG_TAG, "Checking firmware image on server");

                    if(bFirstRun == true)
                        {
                            //init on first run
                            ERROR_CHECK(errHTTPSClientInitialize());
                            bFirstRun = false;
                        }

                    ERROR_CHECK(errHTTPSClientConnectToServer());
                    ERROR_CHECK(errHTTPSClientValidateServer());
                    ERROR_CHECK(errHTTPSClientSendRequest());

                    ERROR_CHECK(errMeshOtaPartitionAccessHttps(&bNewOTAImage));
                    errHTTPSClientReset();

                    if(bNewOTAImage == true)
                        {
                            //set want reboot
                            ESP_LOGI(LOG_TAG, "Updated successfully via HTTPS, set pending reboot");
                            bWantReboot = true;
                        }
                    vTaskDelay( (SERVER_CHECK_INTERVAL*1000) / portTICK_PERIOD_MS); //sleep till next server checks
                }
        }
}

/**
* @fn void vMeshOtaTaskServerWorker(void *arg)
* @brief Task for updating from nodes in mesh network
* @param arg
* @return void
* @author Hendrik Schutter
* @date 21.01.2021
*/
void vMeshOtaTaskOTAWorker(void *arg)
{
    esp_err_t err = ESP_OK;
    bool bNewOTAImage; //true if a new ota image was downloaded and validated
    mesh_addr_t meshNodeAddr; //node that should be checked for ota update
    BaseType_t xReturned;

    while(true)
        {
            err = ESP_OK;
            bNewOTAImage = false;

            if((uxQueueSpacesAvailable(queueNodes) - QUEUE_NODES_SIZE) == 0)
                {
                    //nodes queue is empty

                    xReturned =  xSemaphoreTake(bsOTAProcess, portMAX_DELAY); //wait for binary semaphore that allows to start the OTA process
                    if((xReturned == pdTRUE) && (bWantReboot == true) && (OTA_ALLOW_REBOOT == 1))
                        {
                            ESP_LOGE(LOG_TAG, "ESP32 Reboot ...");
                            vTaskDelay( (1000) / portTICK_PERIOD_MS);
                            esp_restart();
                        }

                    xSemaphoreGive(bsOTAProcess); //free binary semaphore, this allows other tasks to start the OTA process
                    ERROR_CHECK(errMeshOtaSlaveEndpoint(&bNewOTAImage, &meshNodeAddr));
                }
            else
                {
                    //queue not empty
                    if (xQueueReceive(queueNodes, &meshNodeAddr, ((100) / portTICK_PERIOD_MS)) != pdTRUE)
                        {
                            ESP_LOGE(LOG_TAG, "Unable to receive OTA Messages from Queue");
                            err = ESP_FAIL;
                        }

                    ERROR_CHECK(errMeshOtaMasterEndpoint(&bNewOTAImage, &meshNodeAddr));

                    if (err != ESP_OK)
                        {
                            //OTA process faild --> add back to queue
                            vMeshOtaUtilAddNodeToPossibleUpdatableQueue(meshNodeAddr.addr);
                        }
                    else
                        {
                            vMeshOtaUtilClearNeighboursQueue(&meshNodeAddr); //remove this node from queue
                        }
                }

            if(bNewOTAImage == true)
                {
                    //set want reboot
                    ESP_LOGI(LOG_TAG, "Updated successfully via Mesh, set pending reboot");
                    bWantReboot = true;

                    vMeshOtaUtilAddAllNeighboursToQueue(&meshNodeAddr); //add all existing neighbours to queue
                }
            vTaskDelay( (1000) / portTICK_PERIOD_MS);
        }
}

/**
* @fn esp_err_t errMeshOtaSlaveEndpoint(bool* const cpbNewOTAImage)
* @brief Endpoint for OTA process that is called from remote node
* @param cpbNewOTAImage pointer to boolean to signal if a new image was successfully received
* @param cpcMeshNodeAddr pointer to mesh node that transmitted the update
* @return ESP32 error code
* @author Hendrik Schutter
* @date 21.01.2021
*
* Answers the OTA_Version_Request with OTA_Version_Response
* calls errMeshOtaPartitionAccessMeshReceive OR errMeshOtaPartitionAccessMeshTransmit based on version number
*/
esp_err_t errMeshOtaSlaveEndpoint(bool* const cpbNewOTAImage, mesh_addr_t* const cpcMeshNodeAddr)
{
    esp_err_t err = ESP_OK;
    MESH_PACKET_t sOTAMessage;
    const esp_partition_t* cpBootPartition = NULL; //pointer to boot partition (that will booted after reset)
    esp_app_desc_t bootPartitionDesc; //Metadate from boot partition
    *cpbNewOTAImage = false; //set default false

    //read OTAMessages queue
    if(uxQueueSpacesAvailable(queueMessageOTA) < QUEUE_MESSAGE_OTA_SIZE)
        {
            //queue not empty
            if (xQueueReceive(queueMessageOTA, &sOTAMessage, ((100) / portTICK_PERIOD_MS)) != pdTRUE)
                {
                    err = ESP_FAIL;
                }

            if((err == ESP_OK) && (sOTAMessage.type == OTA_Version_Request)) //if OTA_Version_Request
                {
                    xSemaphoreTake(bsOTAProcess, portMAX_DELAY); //wait for binary semaphore that allows to start the OTA process

                    memcpy(cpcMeshNodeAddr, &sOTAMessage.meshSenderAddr, sizeof(mesh_addr_t));

                    cpBootPartition = esp_ota_get_boot_partition(); //get boot partition (that will booted after reset), not the running partition
                    ERROR_CHECK(esp_ota_get_partition_description(cpBootPartition, &bootPartitionDesc)); //get metadata of partition

                    //send OTA_Version_Response to sender of OTA_Version_Request packet wirh version in payload
                    ERROR_CHECK(errMeshOtaUtilSendOtaVersionResponse(&sOTAMessage.meshSenderAddr));

                    if((bMeshOtaUtilNewerVersion((bootPartitionDesc).version, (char*) sOTAMessage.au8Payload)) && (err == ESP_OK)) //compare local and remote version
                        {
                            //remote newer as local
                            ESP_LOGI(LOG_TAG, "remote image on node is newer --> OTA update required from node \"%x:%x:%x:%x:%x:%x\"", sOTAMessage.meshSenderAddr.addr[0], sOTAMessage.meshSenderAddr.addr[1], sOTAMessage.meshSenderAddr.addr[2], sOTAMessage.meshSenderAddr.addr[3], sOTAMessage.meshSenderAddr.addr[4], sOTAMessage.meshSenderAddr.addr[5]);
                            //  --> this version older --> start OTA_Rx --> set cpbNewOTAImage true
                            ERROR_CHECK(errMeshOtaPartitionAccessMeshReceive(cpbNewOTAImage, &sOTAMessage.meshSenderAddr));
                        }

                    if((bMeshOtaUtilNewerVersion((char*) sOTAMessage.au8Payload, (bootPartitionDesc).version)) && (err == ESP_OK)) //compare remote and local version
                        {
                            //local newer as remote
                            ESP_LOGI(LOG_TAG, "remote image on node is older --> OTA send required to node \"%x:%x:%x:%x:%x:%x\"", sOTAMessage.meshSenderAddr.addr[0], sOTAMessage.meshSenderAddr.addr[1], sOTAMessage.meshSenderAddr.addr[2], sOTAMessage.meshSenderAddr.addr[3], sOTAMessage.meshSenderAddr.addr[4], sOTAMessage.meshSenderAddr.addr[5]);
                            //  --> this version newer --> start OTA_Tx
                            ERROR_CHECK(errMeshOtaPartitionAccessMeshTransmit(&sOTAMessage.meshSenderAddr));
                        }
                    xSemaphoreGive(bsOTAProcess); //free binary semaphore, this allows other tasks to start the OTA process
                }
        }
    return err;
}

/**
* @fn esp_err_t errMeshOtaMasterEndpoint(bool* const cpbNewOTAImage, const mesh_addr_t* const cpcMeshNodeAddr)
* @brief Endpoint for OTA process that calls remote node
* @param cpbNewOTAImage pointer to boolean to signal if a new image was successfully received
* @param cpcMeshNodeAddr pointer to remote node addr
* @return ESP32 error code
* @author Hendrik Schutter
* @date 21.01.2021
*
* Sends the OTA_Version_Request to remote node
* calls errMeshOtaPartitionAccessMeshReceive OR errMeshOtaPartitionAccessMeshTransmit based on version number received
*/
esp_err_t errMeshOtaMasterEndpoint(bool* const cpbNewOTAImage, const mesh_addr_t* const cpcMeshNodeAddr)
{
    esp_err_t err = ESP_OK;
    MESH_PACKET_t sOTAMessage;
    const esp_partition_t* cpBootPartition = NULL; //pointer to boot partition (that will booted after reset)
    esp_app_desc_t bootPartitionDesc; //Metadata from boot partition
    bool bNodeIsConnected = false;
    bool bNodeIsResponding = false;
    bool bSameVersion = false;

    *cpbNewOTAImage = false; //set default false

    if(bMeshNetworkIsNodeNeighbour(cpcMeshNodeAddr) == true) //check if node is still connected
        {
            bNodeIsConnected = true; //node is one of the neighbours
            xSemaphoreTake(bsOTAProcess, portMAX_DELAY); //wait for binary semaphore that allows to start the OTA process

            ERROR_CHECK(errMeshOtaUtilSendOtaVersionRequest(cpcMeshNodeAddr)); //send OTA_VERSION_REQUEST with local version in payload

            for (uint32_t u32Index = 0; ((u32Index < QUEUE_MESSAGE_OTA_SIZE) && (bSameVersion == false)); u32Index++) //loop through all OTA messages
                {
                    //if(uxQueueSpacesAvailable(queueMessageOTA) < QUEUE_MESSAGE_OTA_SIZE)
                    //  {
                    //queue not empty
                    if (xQueueReceive(queueMessageOTA, &sOTAMessage, ((3000) / portTICK_PERIOD_MS)) != pdTRUE)
                        {
                            ESP_LOGE(LOG_TAG, "Unable to receive OTA Messages from queue");
                            err = ESP_FAIL;
                        }

                    if((err == ESP_OK) && (sOTAMessage.type == OTA_Version_Response) && (bMeshNetworkCheckMacEquality(sOTAMessage.meshSenderAddr.addr, cpcMeshNodeAddr->addr))) //if OTA_Version_Request
                        {
                            bSameVersion = true;
                            bNodeIsResponding = true;
                            cpBootPartition = esp_ota_get_boot_partition(); //get boot partition (that will booted after reset), not the running partition
                            ERROR_CHECK(esp_ota_get_partition_description(cpBootPartition, &bootPartitionDesc)); //get metadata of partition

                            if((bMeshOtaUtilNewerVersion((bootPartitionDesc).version, (char*) sOTAMessage.au8Payload)) && (err == ESP_OK)) //compare local and remote version
                                {
                                    //remote newer as local
                                    bSameVersion = false;
                                    ESP_LOGI(LOG_TAG, "remote image on node is newer --> OTA update required from node \"%x:%x:%x:%x:%x:%x\"", sOTAMessage.meshSenderAddr.addr[0], sOTAMessage.meshSenderAddr.addr[1], sOTAMessage.meshSenderAddr.addr[2], sOTAMessage.meshSenderAddr.addr[3], sOTAMessage.meshSenderAddr.addr[4], sOTAMessage.meshSenderAddr.addr[5]);
                                    //  --> this version older --> start OTA_Rx --> set cpbNewOTAImage true
                                    ERROR_CHECK(errMeshOtaPartitionAccessMeshReceive(cpbNewOTAImage, &sOTAMessage.meshSenderAddr));
                                    u32Index = QUEUE_MESSAGE_OTA_SIZE;
                                }

                            if((bMeshOtaUtilNewerVersion((char*) sOTAMessage.au8Payload, (bootPartitionDesc).version)) && (err == ESP_OK)) //compare remote and local version
                                {
                                    //local newer as remote
                                    bSameVersion = false;
                                    ESP_LOGI(LOG_TAG, "remote image on node is older --> OTA send required to node \"%x:%x:%x:%x:%x:%x\"", sOTAMessage.meshSenderAddr.addr[0], sOTAMessage.meshSenderAddr.addr[1], sOTAMessage.meshSenderAddr.addr[2], sOTAMessage.meshSenderAddr.addr[3], sOTAMessage.meshSenderAddr.addr[4], sOTAMessage.meshSenderAddr.addr[5]);
                                    //  --> this version newer --> start OTA_Tx
                                    ERROR_CHECK(errMeshOtaPartitionAccessMeshTransmit(&sOTAMessage.meshSenderAddr));
                                    u32Index = QUEUE_MESSAGE_OTA_SIZE;
                                }
                        }
                    else
                        {
                            // OTA Message queue is empty --> wait some time
                            vTaskDelay( (1000/QUEUE_MESSAGE_OTA_SIZE) / portTICK_PERIOD_MS);
                        }
                }//end loop
            xSemaphoreGive(bsOTAProcess); //free binary semaphore, this allows other tasks to start the OTA process
        }

    if((bNodeIsResponding == false) && (bNodeIsConnected == true))
        {
            //add node back to queue if connected and NOT responding
            ESP_LOGI(LOG_TAG, "OTA-Master: connected and NOT responding --> add node back to queue ");
            vMeshOtaUtilAddNodeToPossibleUpdatableQueue(cpcMeshNodeAddr->addr);
        }
    return err;
}