#include "mqtt_comms.h"
#include "mqtt_comms_internal.h"

#include "esp_check.h"
#include "esp_log.h"
#include "event_notifier.h"

#include "endian.h"

#include "ota_update.h"

#define MQTT_EVENT_QUEUE_TIMEOUT pdMS_TO_TICKS(500)

esp_err_t mqtt_comms_init(mqtt_comms_config_t* conf)
{
    is_connected = false;

    ESP_RETURN_ON_FALSE(conf->event.group != NULL, ESP_ERR_INVALID_ARG, TAG, "No event group provided");

    component_id = conf->component_id;
    events = conf->event;
    queues = conf->queues;
    
    mqtt_event_queue = xQueueCreate(5, sizeof(mqtt_request_t));
    ESP_RETURN_ON_FALSE(mqtt_event_queue != NULL, ESP_ERR_NO_MEM, TAG, "Could not create mqtt_event_queue");

    disconnect_message_sent = xSemaphoreCreateBinary();
    ESP_RETURN_ON_FALSE(disconnect_message_sent != NULL, ESP_ERR_NO_MEM, TAG, "Could not create disconnect semaphore");
    disconnect_message_id = -1;

    ESP_RETURN_ON_FALSE(conf->device_id != NULL, ESP_ERR_INVALID_ARG, TAG, "No device ID provided");
    ESP_RETURN_ON_FALSE(strlen(conf->device_id) <= MQTT_MAX_DEVICE_NAME_LENGTH, ESP_ERR_INVALID_SIZE, TAG, "Device ID name too long");
    strncpy(device_id, conf->device_id, sizeof(device_id));
    device_id_len = strlen(device_id);

    connection_payload[0] = '\1';
    memcpy(&connection_payload[1], conf->custom_conn_payload, conf->custom_conn_payload_len);
    connection_payload_len = conf->custom_conn_payload_len + 1;
    snprintf(connection_status_topic, sizeof(connection_status_topic),
            "%s/%s", device_id, CONNECTION_STATUS);

    // Store subscriptions list
    for (int i = 0; i < conf->subs.count; i++)
    {
        mqtt_create_subscriber(&conf->subs.topics[i]);
    }
    // Store publisher list
    for (int i = 0; i < conf->pubs.count; i++)
    {
        mqtt_create_publisher(&conf->pubs.topics[i]);
    }
    
    xTaskCreatePinnedToCore(
        handle_message_task,
        "mqtt_comms_task",
        8192,
        NULL,
        8,
        &handle,
        1
    );
    xTaskCreatePinnedToCore(
        mqtt_event_task,
        "mqtt_event_task",
        4096,
        NULL,
        7,
        NULL,
        1
    );

    return ESP_OK;
}

static void subscribe_to_ota_updates()
{
    mqtt_topic_subscription_t topic;
    topic.listen_to_broadcast = false;
    topic.topic.qos = 1;

    topic.topic.filter = OTA_TOPIC_START_REQUEST;
    subscribe_to_topic(&topic);
    topic.topic.filter = OTA_TOPIC_DATA;
    subscribe_to_topic(&topic);
    topic.topic.filter = OTA_TOPIC_ABORT;
    subscribe_to_topic(&topic);
    topic.topic.filter = OTA_TOPIC_FINISH_REQUEST;
    subscribe_to_topic(&topic);
    topic.topic.filter = OTA_TOPIC_PROGRESS_REQUEST;
    subscribe_to_topic(&topic);
    topic.topic.filter = OTA_TOPIC_PROGRESS_RESPONSE;
    subscribe_to_topic(&topic);
}

esp_err_t mqtt_create_subscriber(mqtt_topic_subscription_t* topic_info)
{
    // We can only store a limited number of subscriptions
    const char* topic_name = topic_info->topic.filter;
    ESP_RETURN_ON_FALSE(sub_count < MQTT_COMMS_MAX_SUBS, ESP_ERR_NO_MEM, TAG, "No memory to subscribe");
    ESP_RETURN_ON_FALSE(validate_topic_name(topic_name), ESP_ERR_INVALID_ARG, TAG, "Topic name %s is invalid", topic_name);

    // Store subscription info
    topic_info->topic_name_len = strlen(topic_name);
    subscriptions[sub_count] = *topic_info;
    sub_count++;

    // Subscribe to topic now if we're connected
    if (is_connected)
    {
        ESP_LOGI(TAG, "Attempting connect to %s", &topic_name);
        if (subscribe_to_topic(topic_info) < 0)
        {
            return ESP_FAIL;
        }
    }

    return ESP_OK;
}

esp_err_t mqtt_create_publisher(mqtt_topic_publisher_t* topic_info)
{
    const char* topic_name = topic_info->topic.filter;
    ESP_RETURN_ON_FALSE(pub_count < MQTT_COMMS_MAX_PUBS, ESP_ERR_NO_MEM, TAG, "No memory to create publisher");
    ESP_RETURN_ON_FALSE(validate_topic_name(topic_name), ESP_ERR_INVALID_ARG, TAG, "Topic name %s is invalid", topic_name);

    char* full_name = pub_full_names[pub_count];
    snprintf(full_name, MQTT_MAX_TOPIC_FULL_NAME_LENGTH + 1,
        "%s/%s", device_id, topic_info->topic.filter);
    topic_info->topic.filter = full_name;  // use the new full-name memory address

    publishers[pub_count] = *topic_info;
    pub_count++;

    return ESP_OK;
}

static void mqtt_event_handler(void* arg, esp_event_base_t event_base,
                        int32_t event_id, void* event_data)
{
    TickType_t start_ticks = xTaskGetTickCount();
    BaseType_t ret;
    mqtt_request_t event_req;
    ESP_LOGI(TAG, "event: base=%s, event_id=%d", event_base, event_id);
    esp_mqtt_event_handle_t event = event_data;
    switch((esp_mqtt_event_id_t)event_id)
    {
        case MQTT_EVENT_CONNECTED:
            ESP_LOGI(TAG, "MQTT_EVENT_CONNECTED");
            is_connected = true;
            // Tell system that we are connected
            set_bits_and_notify_once(events.group, events.mqtt_connected_bit);
            clear_bits_and_notify(events.group, events.mqtt_down_bit);
            event_req = MQTT_REQ_EVENT_CONNECTED;
            ret = xQueueSendToBack(mqtt_event_queue, &event_req, MQTT_EVENT_QUEUE_TIMEOUT);
            if (ret != pdPASS)
            {
                ESP_LOGE(TAG, "Failed to add MQTT_REQ_EVENT_CONNECTED to queue");
            }
            break;
        case MQTT_EVENT_DISCONNECTED:
            ESP_LOGW(TAG, "MQTT_EVENT_DISCONNECTED");
            is_connected = false;
            clear_bits_and_notify_once(events.group, events.mqtt_connected_bit);
            set_bits_and_notify(events.group, events.mqtt_down_bit);
            // Don't need to queue up a reconnect because the MQTT client does it automatically
            break;
        case MQTT_EVENT_DATA:
            ESP_LOGI(TAG, "MQTT_EVENT_DATA");
            ESP_LOGI(TAG, "TOPIC=%.*s\r\n", event->topic_len, event->topic);
            ESP_LOGI(TAG, "DATA=%d bytes: %.*s\r\n", event->data_len, event->data_len, event->data);
            handle_received_message(event);
            break;
        case MQTT_EVENT_ERROR:
            ESP_LOGI(TAG, "MQTT_EVENT_ERROR");
            if (event->error_handle->error_type == MQTT_ERROR_TYPE_TCP_TRANSPORT)
            {
                ESP_LOGE(TAG, "reported from esp-tls: %d", event->error_handle->esp_tls_last_esp_err);
                ESP_LOGE(TAG, "reported from tls stack: %d", event->error_handle->esp_tls_stack_err);
                ESP_LOGE(TAG, "captured as transport's socket errno: %d",  event->error_handle->esp_transport_sock_errno);
                ESP_LOGI(TAG, "Last errno string (%s)", strerror(event->error_handle->esp_transport_sock_errno));
            }
            ESP_LOGE(TAG, "error type %d", event->error_handle->error_type);
            break;
        case MQTT_EVENT_PUBLISHED:
            ESP_LOGI(TAG, "MQTT_EVENT_PUBLISHED");
            ESP_LOGI(TAG,
                "PUBLISHED msg=%d outbox=%d",
                event->msg_id,
                esp_mqtt_client_get_outbox_size(client));
            if (event->msg_id == disconnect_message_id)
            {
                // If we have sent a disconnect message, tell the deep sleep prep task
                // Necessary so we don't go to sleep before the message sends
                xSemaphoreGive(disconnect_message_sent);
            }
            break;
        default:
            ESP_LOGI(TAG, "UNEXPECTED EVENT");
            break;  // we don't care about this
    }
    
    TickType_t end_ticks = xTaskGetTickCount();
    ESP_LOGI(TAG, "EVENT END: %d ms", pdTICKS_TO_MS(end_ticks - start_ticks));
}

static void handle_received_message(const esp_mqtt_event_handle_t event)
{
    event_t message;
    mqtt_tx_message_event_t *message_data;
    message_data = (mqtt_tx_message_event_t*)message.data;
    // Copy topic name
    if (unlikely(event->topic_len > MQTT_MAX_TOPIC_FULL_NAME_LENGTH))
    {
        ESP_LOGE(TAG, "MQTT message topic name is too long");
        return;
    }

    message_data->topic = get_sub_topic_id(event->topic, event->topic_len);
    message.component_id = component_id;
    ESP_LOGI(TAG, "Received topic %.*s, %d %d", event->topic_len, event->topic, message_data->topic, event->data_len);

    // Put message on queue to be handled elsewhere
    message_data->type = MQTT_MSG_EVENT_MESSAGE;
    message_data->data_len = event->data_len;
    memcpy(message_data->data, event->data, message_data->data_len);
    xQueueSendToBack(queues.outbound, (void*)&message, MQTT_MAX_OUTGOING_QUEUE_DELAY);
}

static void mqtt_event_task(void* args)
{
    mqtt_request_t req;

    ESP_LOGI(TAG, "Waiting for WiFi");
    xEventGroupWaitBits(events.group, events.network_connected_bit, pdFALSE, pdFALSE, portMAX_DELAY);
    ESP_LOGI(TAG, "Connected to WiFi");

    char uri_buffer[23];
    sprintf(uri_buffer, "mqtt://%s", MQTT_BROKER_IP_ADDRESS);
    esp_mqtt_client_config_t mqtt_cfg = {
        .broker.address.uri = uri_buffer,
        .buffer.size = MQTT_MAX_BUFFER_LENGTH,

        .session.last_will.topic = connection_status_topic,
        .session.last_will.msg = "\0",  // 0 to indicate disconnect
        .session.last_will.msg_len = 1,
        .session.last_will.qos = 2,
        .session.last_will.retain = 1,
        // .session.keepalive = 15,  // ping timeout of ~15 seconds

        // .network.reconnect_timeout_ms = 500,
        // .network.timeout_ms = 5*1000,
        // .network.refresh_connection_after_ms = 0,
        // .network.disable_auto_reconnect = false,
        .outbox.limit = 1000,
        .task.priority = 15,
    };

    client = esp_mqtt_client_init(&mqtt_cfg);
    ESP_ERROR_CHECK(esp_mqtt_client_register_event(client, ESP_EVENT_ANY_ID, mqtt_event_handler, NULL));  // TODO limit the events
    esp_mqtt_client_start(client);

    while (true)
    {
        if (xQueueReceive(mqtt_event_queue, &req, portMAX_DELAY))
        {
            switch (req)
            {
                case MQTT_REQ_EVENT_CONNECTED:
                    // Tell subscribers that we are connected
                    esp_mqtt_client_publish(
                        client, connection_status_topic, connection_payload,
                        connection_payload_len, 2, true);
                    // Subscribe to configured topics
                    subscribe_to_all();
                    break;
                case MQTT_REQ_EVENT_DISCONNECTED:
                    // Try reconnecting after a delay
                    // ! Only needed if auto-reconnect isn't enabled
                    vTaskDelay(pdMS_TO_TICKS(500));
                    esp_mqtt_client_start(client);
                    break;
                default:
                    break;
            }
        }
    }
}

// Test function to print out the state of tasks - won't be in the full release
#include "freertos/task.h"
static void dump_tasks(void)
{
    static char *buf = NULL;
    if (buf == NULL) buf = malloc(2000);
    if (!buf) return;

    vTaskList(buf);
    ESP_LOGI(TAG, "%s", buf);

    TaskHandle_t handle;
    TaskStatus_t info;
    handle = xTaskGetHandle("mqtt_task");
    if (handle)
    {
        vTaskGetInfo(handle, &info, pdTRUE, eInvalid);
        ESP_LOGI(TAG,
            "mqtt_task: state=%d priority=%u base_priority=%u stack_high_water=%u task_num=%u",
            info.eCurrentState,
            (unsigned)info.uxCurrentPriority,
                (unsigned)info.uxBasePriority,
            (unsigned)info.usStackHighWaterMark,
            (unsigned)info.xTaskNumber);
    }
    else
    {
        ESP_LOGW(TAG, "couldn't get mqtt_task handle");
    }
}

static void handle_message_task(void* args)
{
    mqtt_rx_message_event_t message;
    int response;
    int counter = 0;
    while (true)
    {
        xQueueReceive(queues.inbound, (void*)&message, portMAX_DELAY);

        // Skip message if MQTT isn't connected
        if (!is_connected)
        {
            ESP_LOGW(TAG, "Skipping message due to no MQTT connection");
            continue;
        }

        switch (message.type)
        {
            case MQTT_MSG_EVENT_MESSAGE:
                response = publish_to_topic_id(message.topic, message.data, message.data_len);
                ESP_LOGI(TAG, "publish_to_topic_id=%d, outbox size=%d, queue size=%d", response, esp_mqtt_client_get_outbox_size(client), uxQueueMessagesWaiting(queues.inbound));
                if (counter++ == 10)
                {
                    dump_tasks();
                    counter = 0;
                }
                break;
            default:
                ESP_LOGW(TAG, "Unexpected MQTT message event");
                break;
        }
    }
}

static mqtt_topic_id_t get_sub_topic_id(const char* topic, int topic_len)
{
    const char* topic_suffix;
    int suffix_len;
    if (memcmp(topic, device_id, device_id_len) == 0)
    {
        topic_suffix = topic + device_id_len + 1;  // +1 for the slash
        suffix_len = topic_len - (device_id_len + 1);
    }
    else if (memcmp(topic, BROADCAST_ID, broadcast_id_len) == 0)
    {
        topic_suffix = topic + broadcast_id_len + 1;  // +1 for the slash
        suffix_len = topic_len - (broadcast_id_len + 1);
    }
    else
    {
        return MQTT_TOPIC_UNKNOWN;
    }

    for (int i = 0; i < sub_count; i++)
    {
        const mqtt_topic_subscription_t* sub = &subscriptions[i];
        if (sub->topic_name_len == suffix_len
            && memcmp(sub->topic.filter, topic_suffix, suffix_len) == 0)
        {
            return sub->id;
        }
    }
    return MQTT_TOPIC_UNKNOWN;
}

static int get_pub_topic_index(mqtt_topic_id_t id)
{
    for (int i = 0; i < pub_count; i++)
    {
        if (publishers[i].id == id)
        {
            return i;
        }
    }
    return -1;
}

static int subscribe_to_topic(const mqtt_topic_subscription_t* topic_info)
{
    int result;
    static char full_topic_name[MQTT_MAX_TOPIC_FULL_NAME_LENGTH + 1];
    snprintf(full_topic_name, MQTT_MAX_TOPIC_FULL_NAME_LENGTH + 1,
        "%s/%s", device_id, topic_info->topic.filter);
    result = esp_mqtt_client_subscribe(client, full_topic_name, topic_info->topic.qos);

    if (result >= 0 && topic_info->listen_to_broadcast)
    {
        snprintf(full_topic_name, MQTT_MAX_TOPIC_FULL_NAME_LENGTH + 1,
            "%s/%s", BROADCAST_ID, topic_info->topic.filter);
        result = esp_mqtt_client_subscribe(client, full_topic_name, topic_info->topic.qos);
    }

    return result;
}

static int publish_to_topic(const char* topic_name, int qos, const uint8_t* data, size_t data_len)
{
    char full_name[MQTT_MAX_TOPIC_FULL_NAME_LENGTH + 1];
    snprintf(full_name, MQTT_MAX_TOPIC_FULL_NAME_LENGTH + 1,
        "%s/%s", device_id, topic_name);
    return esp_mqtt_client_publish(client, full_name, (const char*)data, data_len, qos, false);
}

static int publish_to_topic_id(mqtt_topic_id_t id, const uint8_t* data, size_t data_len)
{
    int index = get_pub_topic_index(id);
    if (index < 0) return -1;
    const esp_mqtt_topic_t* topic = &publishers[index].topic;
    return esp_mqtt_client_publish(client, topic->filter, (const char*)data, data_len, topic->qos, false);
}

static bool validate_topic_name(const char* name)
{
    // If the topic name starts or ends with a slash, it's bad
    return name[0] != '/' && name[strlen(name)-1] != '/';
}

static void subscribe_to_all()
{
    ESP_LOGI(TAG, "subscribe_to_all");
    // Subscribe to all the topics we're tracking
    for (int i = 0; i < sub_count; i++)
    {
        if (subscribe_to_topic(&subscriptions[i]) < 0)
        {
            ESP_LOGE(TAG, "Unable to subscribe to topic %s", subscriptions[i].topic.filter);
        }
        else
        {
            ESP_LOGI(TAG, "Subscribed to topic %s", subscriptions[i].topic.filter);
        }
    }

    subscribe_to_ota_updates();
}


uint64_t mqtt_prepare_deep_sleep()
{
    // Tell subscribers that we have disconnected
    xSemaphoreTake(disconnect_message_sent, 0);  // clear semaphore if it has accidentally been set
    disconnect_message_id = esp_mqtt_client_publish(client, connection_status_topic, "\0", 1, 2, true);
    if (disconnect_message_id < 0)
    {
        ESP_LOGE(TAG, "Disconnect message couldn't be sent: %d", disconnect_message_id);
    }
    else if (xSemaphoreTake(disconnect_message_sent, pdMS_TO_TICKS(5000)) != pdTRUE)
    {
        // Log that the disconnect message wasn't sent, but there is nothing else we can do about it
        ESP_LOGE(TAG, "Didn't send disconnect message in time");
    }
    esp_mqtt_client_stop(client);
    esp_mqtt_client_destroy(client);

    return 0;
}