static const char *TAG = "ota";
#include <esp_log.h>

#include <esp_ota_ops.h>
#include <lwip/apps/tftp_server.h>

#include <cstring>

#include "ota.hpp"


class OTA_Session {
public:
    OTA_Session() :
                part(esp_ota_get_next_update_partition(NULL)),
                ota(0), failed(false), total_written(0)
            {
        if (!part) {
            ESP_LOGE(TAG, "Could not find update partition");
            failed = true;
            return;
        }

        auto err = esp_ota_begin(part, OTA_SIZE_UNKNOWN, &ota);
        if (err != ESP_OK) {
            ESP_LOGE(TAG, "esp_ota_begin: %d", err);
            failed = true;
            return;
        }
    }

    ~OTA_Session() {
        if (failed) {
            ESP_LOGE(TAG, "Aborting OTA due to previous failure");
            esp_ota_abort(ota);
            return;
        }

        ESP_LOGI(TAG, "Finalizing OTA (%d bytes)", total_written);
        auto err = esp_ota_end(ota);
        if (err != ESP_OK) {
            ESP_LOGE(TAG, "esp_ota_end: %d", err);
            esp_ota_abort(ota);
            return;
        }

        ESP_LOGI(TAG, "Setting new boot partition");
        err = esp_ota_set_boot_partition(part);
        if (err != ESP_OK) {
            ESP_LOGE(TAG, "esp_ota_set_boot_partition: %d", err);
            return;
        }

        ESP_LOGI(TAG, "OTA Successful");

        esp_restart();
    }

    bool write(const void *data, int len) {
        if (failed) return false;
        auto err = esp_ota_write(ota, data, len);
        if (err != ESP_OK) {
            ESP_LOGE(TAG, "esp_ota_write(%d): %d", len, err);
            failed = true;
            return false;
        }
        total_written += len;
        return true;
    }

private:
    const esp_partition_t* part;
    esp_ota_handle_t ota;
    bool failed;
    int total_written;
};


static void* ota_open(const char *filename, const char *mode, u8_t write) {
    if (strcmp(mode, "octet") != 0) {
        ESP_LOGE(TAG, "Unexpected mode: %s", mode);
        return NULL;
    }

    if (strcmp(filename, "blinky.bin") != 0) {
        ESP_LOGE(TAG, "Unexpected filename: %s", mode);
        return NULL;
    }

    if (!write) {
        ESP_LOGE(TAG, "Illegal read attempt");
        return NULL;
    }

    ESP_LOGI(TAG, "Starting OTA");

    return new OTA_Session();
}

static void ota_close(void* handle) {
    OTA_Session *session = (OTA_Session*)handle;
    if (session) delete session;
    else ESP_LOGE(TAG, "Attempt to close invalid session");
}

static int ota_read(void* handle, void* buf, int len) {
    ESP_LOGE(TAG, "Super illegal read of %d bytes", len);
    return -1;
}

static int ota_write(void* handle, struct pbuf* p) {
    OTA_Session *session = (OTA_Session*)handle;
    if (!session) {
        ESP_LOGE(TAG, "Attempt to write to invalid session");
        return -1;
    }

    for (; p; p = p->next) {
        if (!session->write(p->payload, p->len)) {
            ESP_LOGE(TAG, "Failed to write %d bytes", p->len);
            return -1;
        }
    }

    return 0;
}


static tftp_context ota_context = {
    .open  = ota_open,
    .close = ota_close,
    .read  = ota_read,
    .write = ota_write
};

int start_ota_serv() {
    ESP_LOGI(TAG, "Starting OTA server");

    ESP_ERROR_CHECK(tftp_init(&ota_context));

    return 0;
}


int stop_ota_serv() {
    tftp_cleanup();
    return 0;
}