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

static const char *LOG_TAG = "mesh_ota_access";

esp_err_t errOTAMeshTransmit(mesh_addr_t* pMeshNodeAddr)
{
    esp_err_t err = ESP_OK;
    const esp_partition_t* pBootPartition; //pointer to boot partition (that will booted after reset)
    MESH_PACKET_t sMeshPacket; //packet for sending and receiving
    // uint32_t u32Index = 0U; //index for partition read offset
    bool bAbort = false; //abort the OTA process
    bool bNodeIsResponding = false; //remote node is still active
    uint32_t u32OTABytesWritten = 0U;
    uint32_t u32SegmentCounter = 0U;

    pBootPartition = esp_ota_get_boot_partition(); //get boot partition (that will booted after reset), not the running partition

    //loop through partition to read in segmensts until end or error or abort called
    while( ((OTA_MESH_SEGMENT_SIZE * u32SegmentCounter) < pBootPartition->size) && (err == ESP_OK) && (bAbort == false))
        {
            bNodeIsResponding = false; //reset to default

            // read partition with offset based in index
            ERROR_CHECK(esp_partition_read(pBootPartition, (OTA_MESH_SEGMENT_SIZE * u32SegmentCounter), sMeshPacket.au8Payload, OTA_MESH_SEGMENT_SIZE));
            u32OTABytesWritten = ((u32SegmentCounter+1) * OTA_MESH_SEGMENT_SIZE);
            vPrintOTAProgress(&(pBootPartition->size), &u32OTABytesWritten, Transmitter);

            if(err == ESP_OK)
                {
                    //no error while read --> send OTA_DATA packet
                    sMeshPacket.type = OTA_Data;

                    if((OTA_MESH_SEGMENT_SIZE * (u32SegmentCounter+1)) >= pBootPartition->size) //check if last segment
                        {
                            //last partition image segment --> send OTA_Complete
                            ESP_LOGI(LOG_TAG, "OTA-TX: last segment--> send Complete");
                            sMeshPacket.type = OTA_Complete;
                        }
                    //ESP_LOGI(LOG_TAG, "OTA-TX: send packet");
                    err = errSendMeshPacket(pMeshNodeAddr, &sMeshPacket);
                }
            else
                {
                    // error while read --> send OTA_ABORT and abort this OTA process
                    sMeshPacket.type = OTA_Abort;
                    bAbort = true;
                    ESP_LOGI(LOG_TAG, "OTA-TX: error while read --> send ABORT");
                    errSendMeshPacket(pMeshNodeAddr, &sMeshPacket);
                }

            // loop through all OTA messages or until abort is called or error
            for (uint32_t u32Index = 0; ((u32Index < QUEUE_MESSAGE_OTA_SIZE) && (bAbort == false) && (err == ESP_OK)); u32Index++) //loop through all OTA messages
                {
                    // if(uxQueueSpacesAvailable(queueMessageOTA) < QUEUE_MESSAGE_OTA_SIZE)
                    //   {
                    //queue not empty
                    if (xQueueReceive(queueMessageOTA, &sMeshPacket, ((OTA_MESH_TIMEOUT) / portTICK_PERIOD_MS)) != pdTRUE)
                        {
                            ESP_LOGE(LOG_TAG, "Unable to receive OTA Messages from queue");
                            err = ESP_FAIL;
                        }

                    if((err == ESP_OK) && (bCheckMACEquality(sMeshPacket.meshSenderAddr.addr, pMeshNodeAddr->addr))) //if OTA_Version_Request
                        {
                            //packet from node received
                            switch (sMeshPacket.type)
                                {
                                case OTA_ACK: //increase index for next round
                                    u32Index++;
                                    bNodeIsResponding = true;
                                    u32Index = QUEUE_MESSAGE_OTA_SIZE;//this will end the loop through all OTA messages
                                    break;
                                case OTA_Abort: //abort this OTA process
                                    bAbort = true;
                                    bNodeIsResponding = true;
                                    break;
                                default:
                                    //receives wrong OTA message type from node --> back to queue
                                    //vAddOtaMessageToQueue(&sMeshPacket);
                                    break;
                                }
                        }
                    else if (err == ESP_OK)
                        {
                            //received from wrong node --> back to queue
                            vAddOtaMessageToQueue(&sMeshPacket);
                        }
                    /*
                    }
                    else
                    {
                        // OTA Message queue is empty --> wait some time
                        ESP_LOGI(LOG_TAG, "OTA-TX: ota message queue empty --> wait");
                        vTaskDelay( (OTA_MESH_TIMEOUT) / portTICK_PERIOD_MS);
                    }
                    */

                }//end OTA message loop

            if(bNodeIsResponding == false)
                {
                    //no abort was called but node didn’t responded
                    ESP_LOGE(LOG_TAG, "OTA-TX: no abort was called but node didn’t responded --> error");
                    bAbort = true;
                    err = ESP_FAIL; //this OTA process failed with error
                }
            u32SegmentCounter++;
        }//end of partition segment loop
    vClearOtaMessageQueue(pMeshNodeAddr);
    return err;
}

esp_err_t errOTAMeshReceive(bool* pbNewOTAImage, mesh_addr_t* pMeshNodeAddr)
{
    esp_err_t err = ESP_OK;
    MESH_PACKET_t sMeshPacket; //packet for sending and receiving
    bool bComplete = false; //complete the OTA process
    bool bAbort = false; //abort the OTA process
    bool bNodeIsResponding = false; //remote node is still active
    uint32_t u32OTABytesWritten = 0U; //counter unsed for progress log
    static esp_ota_handle_t otaHandle; //OTA process handle
    *pbNewOTAImage = false;
    uint32_t u32SegmentCounter = 0U;

    ERROR_CHECK(esp_ota_begin(pOTAPartition, OTA_SIZE_UNKNOWN, &otaHandle)); //start ota update process

    //partition segement loop through partition to read in segmensts until end or error or abort called
    while((bComplete == false) && (err == ESP_OK) && (bAbort == false) && (u32OTABytesWritten <= pOTAPartition->size))
        {
            bNodeIsResponding = false; //reset to default

            // loop through all OTA messages or until abort is called
            for (uint32_t u32Index = 0; ((u32Index < QUEUE_MESSAGE_OTA_SIZE) && (bAbort == false)); u32Index++) //loop through all OTA messages
                {
                    //if(uxQueueSpacesAvailable(queueMessageOTA) < QUEUE_MESSAGE_OTA_SIZE)
                    //   {
                    //queue not empty
                    if (xQueueReceive(queueMessageOTA, &sMeshPacket, ((OTA_MESH_TIMEOUT) / portTICK_PERIOD_MS)) != pdTRUE)
                        {
                            ESP_LOGE(LOG_TAG, "Unable to receive OTA Messages from queue");
                            err = ESP_FAIL;
                        }

                    if((err == ESP_OK) && (bCheckMACEquality(sMeshPacket.meshSenderAddr.addr, pMeshNodeAddr->addr))) //if OTA_Version_Request
                        {
                            //packet from node received
                            switch (sMeshPacket.type)
                                {
                                case OTA_Complete: //signal end of this OTA process, fall through because same behavior as OTA_Data
                                    bComplete = true;
                                    ESP_LOGI(LOG_TAG, "OTA-RX: rec Complete --> last segment");
                                //fall through
                                case OTA_Data: //data segement received
                                    bNodeIsResponding = true;
                                    ERROR_CHECK(esp_ota_write(otaHandle, sMeshPacket.au8Payload, OTA_MESH_SEGMENT_SIZE));
                                    u32OTABytesWritten = ((u32SegmentCounter+1) * OTA_MESH_SEGMENT_SIZE);
                                    vPrintOTAProgress(&(pOTAPartition->size), &u32OTABytesWritten, Receiver);
                                    u32Index = QUEUE_MESSAGE_OTA_SIZE; //this will end the loop through all OTA messages
                                    break;
                                case OTA_Abort: //abort this OTA process
                                    bAbort = true;
                                    bNodeIsResponding = true;
                                    ESP_LOGI(LOG_TAG, "OTA-RX: rec Abort");
                                    //this will end the loop through all OTA messages
                                    break;
                                default:
                                    //receives wrong OTA message type from node --> back to queue
                                    //vAddOtaMessageToQueue(&sMeshPacket);
                                    break;
                                }
                        }
                    else if (err == ESP_OK)
                        {
                            //received from wrong node --> back to queue
                            vAddOtaMessageToQueue(&sMeshPacket);
                        }

                    /*  }
                    else
                      {
                          ESP_LOGI(LOG_TAG, "OTA-RX: ota message queue empty --> wait");
                          // OTA Message queue is empty --> wait some time
                          vTaskDelay( (OTA_MESH_TIMEOUT) / portTICK_PERIOD_MS);
                      }
                      */

                }//end of OTA message loop

            if(bNodeIsResponding == false)
                {
                    //no abort was called but node didn’t responded --> error
                    ESP_LOGI(LOG_TAG, "OTA-RX: no abort was called but node didn’t responded --> error");
                    bAbort = true; //this will stop the partition segement loop
                    err = ESP_FAIL; //this OTA process failed with error
                }
            else
                {
                    //node has responded with OTA_DATA or OTA_Complete or OTA_ABORT
                    if(err == ESP_OK)
                        {

                            if(bAbort == false)
                                {
                                    //no error while ota write --> send OTA_ACK packet
                                    //ESP_LOGI(LOG_TAG, "OTA-RX: no error while ota write --> send OTA_ACK packet");
                                    sMeshPacket.type = OTA_ACK;
                                    err = errSendMeshPacket(pMeshNodeAddr, &sMeshPacket);
                                }
                        }
                    else
                        {
                            // error while read --> send OTA_ABORT and abort this OTA process
                            sMeshPacket.type = OTA_Abort;
                            bAbort = true;
                            ESP_LOGI(LOG_TAG, "OTA-RX: abort --> send ABORT");
                            errSendMeshPacket(pMeshNodeAddr, &sMeshPacket);
                        }
                }
            u32SegmentCounter++;
        }//end of partition segement loop

    if(bComplete == true)
        {
            //all OTA segments received --> validate
            ESP_LOGI(LOG_TAG, "OTA-RX: validate image ");
            ERROR_CHECK(esp_ota_end(otaHandle));
            ERROR_CHECK(esp_ota_set_boot_partition(pOTAPartition));
            if(err == ESP_OK)
                {
                    //successfully updated OTA partition
                    *pbNewOTAImage = true;
                }
        }
    else
        {
            //not all OTA segments received --> abort this OTA process
            ERROR_CHECK(esp_ota_abort(otaHandle));
        }
    vClearOtaMessageQueue(pMeshNodeAddr);
    return err;
}

esp_err_t errOTAHTTPS(bool* pbNewOTAImage)
{
    esp_err_t err = ESP_OK;
    char u8OTABuffer[OTA_HTTPS_SEGMENT_SIZE]; //store image segment from server before ota write
    uint32_t u32BufferLenght = OTA_HTTPS_SEGMENT_SIZE; //size of buffer
    uint32_t u32BytesRead = 0; //number of bytes that are read from server, <= u32BufferLenght
    char pcRemoteVersionNumber[12]; //string for version number in server image
    const esp_partition_t* pBootPartition; //pointer to boot partition (that will booted after reset)
    static esp_ota_handle_t otaHandle; //OTA process handle
    uint32_t u32StartOffset = 0U; //start offset for image (exclude the http response data)
    esp_app_desc_t bootPartitionDesc; //Metadate from boot partition
    uint32_t u32OTABytesWritten = 0U; //counter unsed for progress log

    ERROR_CHECK(errHTTPSClientRetrieveData(u8OTABuffer, &u32BufferLenght, &u32BytesRead)); //read first bytes if image, including the version

    ERROR_CHECK(errExtractVersionNumber(u8OTABuffer, &u32BytesRead, pcRemoteVersionNumber)); //extract version numbers

    if(err == ESP_OK) //check if version number is found
        {
            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

            if(bNewerVersion((bootPartitionDesc).version, pcRemoteVersionNumber)) //compare local and remote version
                {
                    // server image is newer --> OTA update required
                    ESP_LOGI(LOG_TAG, "Server: image is newer --> OTA update required");

                    ERROR_CHECK(errFindImageStart(u8OTABuffer, &u32BufferLenght, &u32StartOffset)); //get image start offset

                    ERROR_CHECK(esp_ota_begin(pOTAPartition, OTA_SIZE_UNKNOWN, &otaHandle)); //start ota update process
                    if(err == ESP_OK)
                        {
                            //image download and ota partition write
                            ESP_LOGI(LOG_TAG, "start OTA download via HTTPS");
                            do
                                {
                                    vPrintOTAProgress(&(pOTAPartition->size), &u32OTABytesWritten, Receiver);
                                    ERROR_CHECK(esp_ota_write(otaHandle, (const void*) u8OTABuffer+u32StartOffset, (u32BytesRead-u32StartOffset)));

                                    if(err == ESP_OK)
                                        {
                                            //write was succsesfull
                                            u32StartOffset = 0U; //reset the offset for next download
                                            ERROR_CHECK(errHTTPSClientRetrieveData(u8OTABuffer, &u32BufferLenght, &u32BytesRead)); //download next data segment
                                            u32OTABytesWritten = u32OTABytesWritten + u32BytesRead; //update counter
                                        }
                                }
                            while ((u32BytesRead > 0) && (err == ESP_OK) && (u32OTABytesWritten <= pOTAPartition->size)); //loop until error or complete image downloaded
                        }

                    if(err == ESP_OK)
                        {
                            //no error occurred --> finish ota update process
                            ERROR_CHECK(esp_ota_end(otaHandle)); //finish process
                            ERROR_CHECK(esp_ota_set_boot_partition(pOTAPartition)); //set new image as boot
                            if(err == ESP_OK)
                                {
                                    *pbNewOTAImage = true; //image validated
                                }
                        }
                    else
                        {
                            //error occurred --> abort ota update process
                            ESP_LOGE(LOG_TAG, "abort ota process due to error 0x%x -> %s", err, esp_err_to_name(err));
                            ERROR_CHECK(esp_ota_abort(otaHandle));
                            *pbNewOTAImage = false; //ota update failed
                        }
                }
            else
                {
                    ESP_LOGI(LOG_TAG, "server image is NOT newer --> OTA update NOT required");
                }
            xSemaphoreGive(bsOTAProcess); //free binary semaphore, this allows other tasks to start the OTA process
        } //end version number extracted
    return err;
}