#include "RemoteManagement.h"
#include "OTAPacketBuilder.h"

#define TAG "RemoteManagement"

void RemoteManagement::mark_ota_success() {
    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) {
            ESP_LOGI(TAG, "Diagnostics completed successfully! Continuing execution ...");
            esp_ota_mark_app_valid_cancel_rollback();
            // ESP_LOGE(TAG, "Diagnostics failed! Start rollback to the previous version ...");
            // esp_ota_mark_app_invalid_rollback_and_reboot();
        }
    }
}

bool RemoteManagement::start_ota() {
    ESP_LOGI(TAG, "Start OTA");
    return ESP_OK == esp_ota_begin(update_partition, OTA_SIZE_UNKNOWN, &ota_handle);
}

bool RemoteManagement::write_ota(const uint8_t* ota_data, size_t size) {

    flatbuffers::Verifier verifier(ota_data, size);
    if (!Messaging::VerifyOTAPacketBuffer(verifier)) {
        return false;
    }

    const auto &ota_packet = Flatbuffers::OTAPacketBuilder::parse_ota_packet(ota_data);

    if (ota_packet->sequence_number() <= packet_number) { // got an old packet
        return true;
    } else if (ota_packet->sequence_number() == packet_number + 1) { // got the next packet
        ESP_LOGI(TAG, "Write %d bytes of OTA packet %d", ota_packet->length(), ota_packet->sequence_number());
        if (ESP_OK == esp_ota_write(ota_handle, ota_packet->payload()->data(), ota_packet->length())) {
            packet_number++;
            return true;
        } else {
           return false;
        }
    }
    return false; // got an incorrect packet
}

bool RemoteManagement::ota_end() {
    ESP_LOGI(TAG, "OTA end");
    if (ESP_OK != esp_ota_end(ota_handle)) {
        return false;
    }
    return ESP_OK == esp_ota_set_boot_partition(update_partition);
}

void RemoteManagement::restart() {
    ESP_LOGI(TAG, "restart");
    esp_restart();
}

void RemoteManagement::register_calls(MessagingInterface &messaging_interface) {
    messaging_interface.register_function(4, [](const uint8_t *parameters, size_t parameter_size, Writer& writer) {
        const auto out = start_ota();
        ESP_LOGI(TAG, "Start returned");
        std::vector<uint8_t> vec{(uint8_t)(out ? 1 : 0)};
        writer.write(vec);
    });
    messaging_interface.register_function(5, [](const uint8_t *parameters, size_t parameter_size, Writer& writer) {
        const auto out = write_ota(parameters, parameter_size);
        std::vector<uint8_t> vec{(uint8_t)(out ? 1 : 0)};
        writer.write(vec);
    });
    messaging_interface.register_function(6, [](const uint8_t *parameters, size_t parameter_size, Writer& writer) {
        const auto out = ota_end();
        std::vector<uint8_t> vec{(uint8_t)(out ? 1 : 0)};
        writer.write(vec);
    });
    messaging_interface.register_function(7, [](const uint8_t *parameters, size_t parameter_size, Writer& writer) {
        restart();
        std::vector<uint8_t> vec{1};
        writer.write(vec);
    });
}