#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";

esp_err_t errMeshOTAInitialize()
{
    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;
                }
        }

    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;
}

void vMeshOtaTaskServerWorker(void *arg)
{
    esp_err_t err;
    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)
                        {
                            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;
                            vMeshOtaUtilAddAllNeighboursToQueue(); //add all existing neighbours to queue (aparent will not be added because this node is the root)
                        }
                    vTaskDelay( (SERVER_CHECK_INTERVAL*1000) / portTICK_PERIOD_MS); //sleep till next server checks
                }
        }
}

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

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

            if((uxQueueSpacesAvailable(queueNodes) - QUEUE_NODES_SIZE) == 0)
                {
                    //nodes queue is empty
                    ESP_LOGI(LOG_TAG, "nodes queue is empty");

                    if(bWantReboot == true)
                        {
                            //ESP_LOGI(LOG_TAG, "ESP32 Reboot ...");
                            //vTaskDelay( (1000) / portTICK_PERIOD_MS);
                            //esp_restart();
                        }

                    ERROR_CHECK(errMeshOtaSlaveEndpoint(&bNewOTAImage));
                }
            else
                {
                    //queue not empty
                    ESP_LOGI(LOG_TAG, "nodes queue not empty: %i", (QUEUE_NODES_SIZE - uxQueueSpacesAvailable(queueNodes)));

                    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);
                        }
                }

            if(bNewOTAImage == true)
                {
                    //set want reboot
                    ESP_LOGI(LOG_TAG, "Updated successfully via Mesh, set pending reboot");
                    bWantReboot = true;
                    vMeshOtaUtilAddAllNeighboursToQueue(); //add all existing neighbours to queue
                }
            vTaskDelay( (1000) / portTICK_PERIOD_MS);
        }
}

esp_err_t errMeshOtaSlaveEndpoint(bool* const cpbNewOTAImage)
{
    esp_err_t err = ESP_OK;
    MESH_PACKET_t sOTAMessage;
    const esp_partition_t* pBootPartition; //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)
                {
                    ESP_LOGE(LOG_TAG, "Unable to receive OTA Messages from Queue");
                    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

                    pBootPartition = esp_ota_get_boot_partition(); //get boot partition (that will booted after reset), not the running partition
                    ERROR_CHECK(esp_ota_get_partition_description(pBootPartition, &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");
                            //  --> 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");
                            //  --> 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;
}

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* pBootPartition; //pointer to boot partition (that will booted after reset)
    esp_app_desc_t bootPartitionDesc; //Metadata from boot partition
    bool bNodeIsConnected = false;
    bool bNodeIsResponding = 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

            ESP_LOGI(LOG_TAG, "Mesh-Master: send Version_Request to 0x%x", cpcMeshNodeAddr->addr[5]);
            ERROR_CHECK(errMeshOtaUtilSendOTAVersionRequest(cpcMeshNodeAddr)); //send OTA_VERSION_REQUEST with local version in payload

            for (uint32_t u32Index = 0; u32Index < QUEUE_MESSAGE_OTA_SIZE; 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
                                {
                                    bNodeIsResponding = true;
                                    pBootPartition = esp_ota_get_boot_partition(); //get boot partition (that will booted after reset), not the running partition
                                    ERROR_CHECK(esp_ota_get_partition_description(pBootPartition, &bootPartitionDesc)); //get metadata of partition

                                    if((bMeshOtaUtilNewerVersion((bootPartitionDesc).version, (char*) sOTAMessage.au8Payload)) && (err == ESP_OK)) //compare local and remote version
                                        {
                                            //remote newer as local
                                            ESP_LOGI(LOG_TAG, "Mesh: remote image on node is newer --> OTA update required");
                                            //  --> 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, "Mesh: remote image on node is older --> OTA send required");
                                            //  --> this version newer --> start OTA_Tx
                                            ERROR_CHECK(errMeshOtaPartitionAccessMeshTransmit(&sOTAMessage.meshSenderAddr));
                                        }
                                }
                            else if (err == ESP_OK)
                                {
                                    //received from wrong node or type --> back to queue
                                    vMeshOtaUtilAddOtaMessageToQueue(&sOTAMessage);
                                }
                        }
                    else
                        {
                            // OTA Message queue is empty --> wait some time
                            ESP_LOGI(LOG_TAG, "OTA-Master: 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;
}