#include "config_store.h"
#include "provisioning.h"
#include "ble_scanner.h"
#include "mqtt_publisher.h"
#include "led_indicator.h"

#include "esp_wifi.h"
#include "esp_event.h"
#include "esp_netif.h"
#include "esp_log.h"
#include "mdns.h"
#include "freertos/FreeRTOS.h"
#include "freertos/event_groups.h"
#include "freertos/task.h"
#include <stdio.h>
#include <string.h>

#define TAG "main"

#define WIFI_CONNECTED_BIT BIT5   /* offset to avoid clashing with mqtt_publisher bits */

#define DEFAULT_MQTT_PORT 1883
#define MDNS_QUERY_TIMEOUT_MS 3000
#define MDNS_RETRY_INTERVAL_MS 30000

static EventGroupHandle_t s_evt;

static void wifi_event_handler(void *arg, esp_event_base_t base,
                                int32_t id, void *data)
{
    if (base == WIFI_EVENT && id == WIFI_EVENT_STA_DISCONNECTED) {
        ESP_LOGW(TAG, "WiFi disconnected, reconnecting...");
        led_indicator_set(LED_CONNECTING);
        xEventGroupClearBits(s_evt, WIFI_CONNECTED_BIT);
        esp_wifi_connect();
    } else if (base == IP_EVENT && id == IP_EVENT_STA_GOT_IP) {
        ip_event_got_ip_t *evt = (ip_event_got_ip_t *)data;
        ESP_LOGI(TAG, "Got IP: " IPSTR, IP2STR(&evt->ip_info.ip));
        xEventGroupSetBits(s_evt, WIFI_CONNECTED_BIT);
    }
}

static void on_ble_scan_result(const char *tag_id, int8_t rssi)
{
    mqtt_publisher_send_rssi(tag_id, rssi);
}

static void wifi_init_sta(void)
{
    esp_event_handler_register(WIFI_EVENT, ESP_EVENT_ANY_ID, wifi_event_handler, NULL);
    esp_event_handler_register(IP_EVENT, IP_EVENT_STA_GOT_IP, wifi_event_handler, NULL);
    ESP_ERROR_CHECK(esp_wifi_connect());
}

static bool resolve_mqtt_broker_mdns(char *host_out, size_t host_len, uint16_t *port_out)
{
    ESP_ERROR_CHECK(mdns_init());

    mdns_result_t *results = NULL;
    bool found = false;

    esp_err_t err = mdns_query_ptr("_mqtt", "_tcp", MDNS_QUERY_TIMEOUT_MS, 4, &results);
    if (err != ESP_OK || !results) {
        ESP_LOGW(TAG, "mDNS query for _mqtt._tcp failed");
        goto free_results;
    }    

    mdns_result_t *r = results;
    while (r) {
        if (r->port) {
            if (r->addr) {
                snprintf(host_out, host_len, IPSTR, IP2STR(&r->addr->addr.u_addr.ip4));
                *port_out = r->port;
                ESP_LOGI(TAG, "mDNS found broker: %s:%u", host_out, *port_out);
                found = true;
                goto free_results;
            } else if (r->hostname) {
                esp_ip4_addr_t ip4;
                if (mdns_query_a(r->hostname, MDNS_QUERY_TIMEOUT_MS, &ip4) == ESP_OK) {
                    snprintf(host_out, host_len, IPSTR, IP2STR(&ip4));
                    *port_out = r->port;
                    ESP_LOGI(TAG, "mDNS found broker: %s:%u", host_out, *port_out);
                    found = true;
                    goto free_results;
                }
                ESP_LOGW(TAG, "mDNS A query for %s failed", r->hostname);
            }
        }
        r = r->next;
    }

free_results:
    mdns_query_results_free(results);

    mdns_free();
    return found;
}

void app_main(void)
{
    ESP_ERROR_CHECK(config_store_init());

    led_indicator_init();
    led_indicator_set(LED_CONNECTING);

    ESP_ERROR_CHECK(esp_netif_init());
    ESP_ERROR_CHECK(esp_event_loop_create_default());

    s_evt = xEventGroupCreate();

    /* WiFi hardware must be up before the provisioning manager runs */
    esp_netif_create_default_wifi_sta();
    wifi_init_config_t wcfg = WIFI_INIT_CONFIG_DEFAULT();
    ESP_ERROR_CHECK(esp_wifi_init(&wcfg));
    ESP_ERROR_CHECK(esp_wifi_set_mode(WIFI_MODE_STA));
    ESP_ERROR_CHECK(esp_wifi_start());

    if (!config_store_is_provisioned()) {
        ESP_LOGI(TAG, "Not provisioned - starting BLE provisioning");
        provisioning_run();  /* blocks until complete */
        /* BT memory released by NETWORK_PROV_SCHEME_BLE_EVENT_HANDLER_FREE_BTDM */
    }

    // connect to wifi
    ESP_LOGI(TAG, "Connecting to WiFi...");
    led_indicator_set(LED_CONNECTING);

    // wifi_init_sta registers events and calls connect; skip if already inited
    wifi_init_sta();

    xEventGroupWaitBits(s_evt, WIFI_CONNECTED_BIT, pdFALSE, pdTRUE, portMAX_DELAY);
    ESP_LOGI(TAG, "WiFi connected");

    // derive sensor_id from WiFi MAC
    uint8_t mac[6];
    esp_wifi_get_mac(WIFI_IF_STA, mac);
    char sensor_id[32];
    snprintf(sensor_id, sizeof(sensor_id),
             "anchor_%02x%02x%02x", mac[3], mac[4], mac[5]);
    ESP_LOGI(TAG, "Sensor ID: %s", sensor_id);

    // resolve MQTT broker
    char broker_host[MQTT_HOST_MAX_LEN] = {0};
    uint16_t broker_port = DEFAULT_MQTT_PORT;
    char broker_uri[160];

    bool resolved = false;
    while (!resolved) {
        // favour config store override over mdns
        resolved = (config_store_get_mqtt_override(broker_host, &broker_port) == ESP_OK);
        if (resolved) {
            ESP_LOGI(TAG, "Using NVS MQTT override: %s:%u", broker_host, broker_port);
        } else {
           resolved = resolve_mqtt_broker_mdns(broker_host, sizeof(broker_host), &broker_port);        
        }

        if (!resolved) {
            ESP_LOGW(TAG, "No broker found, retrying in %ds...", MDNS_RETRY_INTERVAL_MS / 1000);
            led_indicator_set(LED_ERROR);
            vTaskDelay(pdMS_TO_TICKS(MDNS_RETRY_INTERVAL_MS));
            led_indicator_set(LED_CONNECTING);
        }
    }
    snprintf(broker_uri, sizeof(broker_uri), "mqtt://%s:%u", broker_host, broker_port);

    // init MQTT
    ESP_ERROR_CHECK(mqtt_publisher_init(sensor_id, broker_uri, s_evt));
    xEventGroupWaitBits(s_evt, MQTT_CONNECTED_BIT, pdFALSE, pdTRUE, portMAX_DELAY);

    // init BLE scanner
    ble_scanner_init(on_ble_scan_result);
    ble_scanner_start();

    led_indicator_set(LED_SCANNING);
    ESP_LOGI(TAG, "Anchor running - scanning BLE and publishing RSSI");

    // main event loop
    const EventBits_t cmd_bits =
        MQTT_CALIBRATE_START | MQTT_CALIBRATE_STOP |
        MQTT_SELECTED_BIT    | MQTT_DESELECTED_BIT;

    bool calibrating = false;
    bool selected    = false;

    for (;;) {
        EventBits_t bits = xEventGroupWaitBits(
            s_evt, cmd_bits, pdTRUE, pdFALSE, pdMS_TO_TICKS(100));

        if (bits & MQTT_CALIBRATE_START) {
            ESP_LOGI(TAG, "Calibration started");
            calibrating = true;
            led_indicator_set(LED_CALIBRATING);
        }
        if (bits & MQTT_CALIBRATE_STOP) {
            ESP_LOGI(TAG, "Calibration stopped");
            calibrating = false;
            led_indicator_set(selected ? LED_SELECTED : LED_SCANNING);
        }
        if (bits & MQTT_SELECTED_BIT) {
            ESP_LOGI(TAG, "Anchor selected");
            selected = true;
            if (!calibrating) led_indicator_set(LED_SELECTED);
        }
        if (bits & MQTT_DESELECTED_BIT) {
            ESP_LOGI(TAG, "Anchor deselected");
            selected = false;
            if (!calibrating) led_indicator_set(LED_SCANNING);
        }
    }
}