/* ESP drivers */
#include "esp_log.h"
#include "esp_partition.h"
#include "esp_system.h"
#include "esp_tls.h"

/* FreeRTOS */
#include "freertos/FreeRTOS.h"
#include "freertos/semphr.h"
#include "freertos/task.h"

/* Cert bundle */
#include "esp_crt_bundle.h"

/* MQTT client */
#include "mqtt_client.h"

#define LOG_TAG "APP_MQTT"

#define APP_MQTT_PART_TYPE                   ((esp_partition_type_t)0x40)
#define APP_MQTT_PART_CRT_NAME               "tls_crt"
#define APP_MQTT_PART_KEY_NAME               "tls_key"
#define APP_MQTT_PART_CRT_SUBTYPE            0x00
#define APP_MQTT_PART_KEY_SUBTYPE            0x01
#define APP_MQTT_PART_HEADER_VALID_SIGNATURE (0x66CCFFAA)

typedef struct {
    uint32_t signature;
    uint32_t length;
    char     data;
} app_mqtt_part_header_t;

static SemaphoreHandle_t s_mqtt_semaphore;

static esp_mqtt_client_handle_t    s_mqtt_client;
static esp_partition_mmap_handle_t s_mqtt_crt_mmap_handle;
static esp_partition_mmap_handle_t s_mqtt_key_mmap_handle;

static void app_mqtt_event_handler(void *arg, esp_event_base_t event_base, int32_t event_id, void *event_data);

esp_err_t app_mqtt_init(uint32_t timeout_ms) {
    esp_err_t ret = ESP_OK;

    s_mqtt_semaphore = xSemaphoreCreateBinary();
    if (s_mqtt_semaphore == NULL) {
        ESP_LOGE(LOG_TAG, "Failed to create semaphore");
        return ESP_ERR_NO_MEM;
    }

    const esp_partition_t *part_crt =
        esp_partition_find_first(APP_MQTT_PART_TYPE, APP_MQTT_PART_CRT_SUBTYPE, APP_MQTT_PART_CRT_NAME);

    const esp_partition_t *part_key =
        esp_partition_find_first(APP_MQTT_PART_TYPE, APP_MQTT_PART_KEY_SUBTYPE, APP_MQTT_PART_KEY_NAME);

    if (part_crt == NULL) {
        ESP_LOGE(LOG_TAG, "Failed to find certificate partition.");
        goto destroy_semaphore_exit;
    }

    ESP_LOGI(LOG_TAG, "Found certificate partition at 0x%08lx", part_crt->address);

    if (part_key == NULL) {
        ESP_LOGE(LOG_TAG, "Failed to find key partition.");
        goto destroy_semaphore_exit;
    }

    ESP_LOGI(LOG_TAG, "Found key partition at 0x%08lx", part_key->address);

    app_mqtt_part_header_t *ptr_crt;
    app_mqtt_part_header_t *ptr_key;

    ret = esp_partition_mmap(part_crt, 0, part_crt->size, ESP_PARTITION_MMAP_DATA, (const void **)(&ptr_crt),
                             &s_mqtt_crt_mmap_handle);
    if (ret != ESP_OK) {
        ESP_LOGE(LOG_TAG, "Failed to map certificate partition");
        goto destroy_semaphore_exit;
    }

    ret = esp_partition_mmap(part_key, 0, part_key->size, ESP_PARTITION_MMAP_DATA, (const void **)(&ptr_key),
                             &s_mqtt_key_mmap_handle);
    if (ret != ESP_OK) {
        ESP_LOGE(LOG_TAG, "Failed to map key partition");
        goto unmap_cert_exit;
    }

    if (ptr_crt->signature != APP_MQTT_PART_HEADER_VALID_SIGNATURE) {
        ESP_LOGE(LOG_TAG, "Certificate partition content is invalid");
        goto unmap_cert_exit;
    }

    ESP_LOGI(LOG_TAG, "Certificate length: %ld", ptr_crt->length);

    if (ptr_key->signature != APP_MQTT_PART_HEADER_VALID_SIGNATURE) {
        ESP_LOGE(LOG_TAG, "Key partition content is invalid");
        goto unmap_cert_exit;
    }

    ESP_LOGI(LOG_TAG, "Key length: %ld", ptr_key->length);

    const esp_mqtt_client_config_t mqtt_cfg = {
        .broker =
            {
                .address.uri                    = CONFIG_APP_MQTT_BROKER_ADDR,
                .verification.crt_bundle_attach = esp_crt_bundle_attach,
            },
        .credentials.authentication =
            {
                .certificate     = &ptr_crt->data,
                .key             = &ptr_key->data,
            },
    };

    s_mqtt_client = esp_mqtt_client_init(&mqtt_cfg);
    if (s_mqtt_client == NULL) {
        ret = ESP_ERR_NO_MEM;
        goto unmap_cert_exit;
    }

    ret = esp_mqtt_client_register_event(s_mqtt_client, ESP_EVENT_ANY_ID, app_mqtt_event_handler, NULL);
    if (ret != ESP_OK) {
        ESP_LOGE(LOG_TAG, "Failed to register MQTT event.");
        goto destroy_client_exit;
    }

    ret = esp_mqtt_client_start(s_mqtt_client);
    if (ret != ESP_OK) {
        ESP_LOGE(LOG_TAG, "Failed to start MQTT client.");
        goto destroy_client_exit;
    }

    if (xSemaphoreTake(s_mqtt_semaphore, pdMS_TO_TICKS(timeout_ms)) != pdPASS) {
        ESP_LOGE(LOG_TAG, "Failed to connect to broker in time.");
        goto destroy_client_exit;
    }

    return ESP_OK;

destroy_client_exit:
    esp_mqtt_client_destroy(s_mqtt_client);

unmap_cert_exit:
    esp_partition_munmap(s_mqtt_crt_mmap_handle);

destroy_semaphore_exit:
    vSemaphoreDelete(s_mqtt_semaphore);

    return ret;
}

esp_err_t app_mqtt_deinit(void) {
    esp_err_t ret;
    esp_mqtt_client_stop(s_mqtt_client);

    ret = esp_mqtt_client_destroy(s_mqtt_client);
    if (ret != ESP_OK) {
        return ret;
    }

    vSemaphoreDelete(s_mqtt_semaphore);

    esp_partition_munmap(s_mqtt_crt_mmap_handle);
    esp_partition_munmap(s_mqtt_key_mmap_handle);

    return ESP_OK;
}

esp_err_t app_mqtt_publish(char *topic, char *payload) {
    return esp_mqtt_client_publish(s_mqtt_client, topic, payload, (int)strlen(payload), 0, 0);
}

static void app_mqtt_event_handler(void *arg, esp_event_base_t event_base, int32_t event_id, void *event_data) {
    switch ((esp_mqtt_event_id_t)event_id) {
        case MQTT_EVENT_CONNECTED:
            ESP_LOGI(LOG_TAG, "Connected to broker.");
            xSemaphoreGive(s_mqtt_semaphore);
            break;
        default:
            break;
    }
}