#include <string.h>
#include "esp_log.h"
#include "dhcp_wd.h"
#include "esp_wifi.h"
//#include "esp_eth.h"
#include "ping.h"

#define xstr(s) str(s)
#define str(s) #s

#define PERIOD_GW_PING_S  CONFIG_DHCPWD_PERIOD_GW_PING_S
#define GETIP_TIMEOUT_S   CONFIG_DHCPWD_GETIP_TIMEOUT_S
#define TASK_STACK_SIZE   CONFIG_DHCPWD_TASK_STACK_SIZE
#define TASK_PRIO         CONFIG_DHCPWD_TASK_PRIORITY

static const char *TAG = "dhcp_wd";

static void dhcp_watchdog_task(void *parm);

struct dhcp_wd_instance {
    TaskHandle_t task;
    esp_netif_t * iface;
    bool is_wifi;
    bool running;
};

#define STATES_ENUM \
    X(DISCONECTED) \
    X(CONECTED_WAIT_IP) \
    X(CONECTED_WAIT_IP2) \
    X(CONECTED)

enum dhcp_wd_state {
#undef X
#define X(s) STATE_##s,
    STATES_ENUM
};

const char *state_names[] = {
#undef X
#define X(s) xstr(s),
    STATES_ENUM
};

enum dhcp_wd_notify {
    NOTIFY_CONNECTED = BIT0,
    NOTIFY_DISCONNECTED = BIT1,
    NOTIFY_GOT_IP = BIT2,
    NOTIFY_SHUTDOWN = BIT3, // kills the task
};

/** Send a notification to the watchdog task */
esp_err_t dhcp_watchdog_notify(dhcp_wd_handle_t handle, const enum dhcp_wd_event event)
{
    assert(handle != NULL);
    assert(handle->task != NULL);

    uint32_t flag = 0;

    switch (event) {
        case DHCP_WD_NOTIFY_CONNECTED:
            flag = NOTIFY_CONNECTED;
            break;

        case DHCP_WD_NOTIFY_DISCONNECTED:
            flag = NOTIFY_DISCONNECTED;
            break;

        case DHCP_WD_NOTIFY_GOT_IP:
            flag = NOTIFY_GOT_IP;
            break;

        default:
            break;
    }

    BaseType_t ret = pdPASS;
    if (flag != 0) {
        ret = xTaskNotify(handle->task, flag, eSetBits);
    }

    return (pdPASS == ret) ? ESP_OK : ESP_FAIL;
}

/**
 * Start the watchdog
 */
esp_err_t dhcp_watchdog_start(esp_netif_t * netif, bool is_wifi, dhcp_wd_handle_t *pHandle)
{
    assert(pHandle != NULL);

    dhcp_wd_handle_t handle = calloc(1, sizeof(struct dhcp_wd_instance));
    if (!handle) return ESP_ERR_NO_MEM;
    *pHandle = handle;

    handle->iface = netif;
    handle->is_wifi = is_wifi;

    BaseType_t ret = xTaskCreate(dhcp_watchdog_task, "dhcp-wd", TASK_STACK_SIZE, (void *)handle, TASK_PRIO, &handle->task);
    handle->running = true;

    return (pdPASS == ret) ? ESP_OK : ESP_FAIL;
}

/**
 * Check if a watchdog is still running
 *
 * @param handle
 * @return is running
 */
bool dhcp_watchdog_is_running(dhcp_wd_handle_t handle)
{
    return handle->running;
}

/**
 * Stop the watchdog and free resources
 */
esp_err_t dhcp_watchdog_stop(dhcp_wd_handle_t *pHandle)
{
    assert(pHandle != NULL);
    assert(*pHandle != NULL);
    xTaskNotify((*pHandle)->task, NOTIFY_SHUTDOWN, eSetBits);
    *pHandle = NULL;
    return ESP_OK;
}

/**
 * @param parm - tcpip_adapter_if_t iface (cast to void *) - typically TCPIP_ADAPTER_IF_STA
 */
static void dhcp_watchdog_task(void *parm)
{
    enum dhcp_wd_state state = STATE_DISCONECTED;

    dhcp_wd_handle_t handle = parm;
    assert(handle != NULL);
    assert(handle->iface != NULL);

    ESP_LOGI(TAG, "Watchdog started");

    while (1) {
        uint32_t flags = 0;
        uint32_t wait_s;
        TickType_t waittime;

        switch (state) {
            case STATE_DISCONECTED:
                wait_s = waittime = portMAX_DELAY;
                break;

            case STATE_CONECTED_WAIT_IP:
            case STATE_CONECTED_WAIT_IP2:
                wait_s = GETIP_TIMEOUT_S;
                waittime = (GETIP_TIMEOUT_S * 1000) / portTICK_PERIOD_MS;
                break;

            case STATE_CONECTED:
                wait_s = PERIOD_GW_PING_S;
                waittime = (PERIOD_GW_PING_S * 1000) / portTICK_PERIOD_MS;
                break;

            default:
                assert(0);
        }

        ESP_LOGD(TAG, "State %s, wait %d s", state_names[state], wait_s);
        BaseType_t rv = xTaskNotifyWait(
            /* no clear on entry */ pdFALSE,
            /* clear all on exit */ ULONG_MAX,
                                    &flags, waittime);

        if (rv == pdPASS) {
            // the order here is important in case we get multiple events at once

            if (flags & NOTIFY_DISCONNECTED) {
                state = STATE_DISCONECTED;
            }

            if (flags & NOTIFY_CONNECTED) {
                state = STATE_CONECTED_WAIT_IP;
            }

            if (flags & NOTIFY_GOT_IP) {
                state = STATE_CONECTED;
            }

            if (flags & NOTIFY_SHUTDOWN) {
                // kill self
                handle->running = false;
                free(handle);
                vTaskDelete(NULL);
                return;
            }
        } else {
            // a timeout occurred

            switch (state) {
                case STATE_DISCONECTED:
                    // this shouldn't happen, we have infinite delay waiting for disconnected
                    ESP_LOGW(TAG, "dhcp_wd double discon evt");
                    break;

                case STATE_CONECTED_WAIT_IP:
                    ESP_LOGW(TAG, "Get IP timeout, restarting DHCP client");
                    // this is a bit suspicious
                    // try to restart the DHCPC client
                    ESP_ERROR_CHECK(esp_netif_dhcpc_stop(handle->iface));
                    ESP_ERROR_CHECK(esp_netif_dhcpc_start(handle->iface));
                    state = STATE_CONECTED_WAIT_IP2;
                    break;

                case STATE_CONECTED_WAIT_IP2:
                    ESP_LOGW(TAG, "Get IP timeout 2, restarting network stack");
                    // well now this is weird. try flipping the whole WiFi/Eth stack
                    if (handle->is_wifi) {
                        ESP_ERROR_CHECK(esp_wifi_disconnect());
                    }
                    // this will trigger the disconnected event and loop back into Disconnected
                    // the disconnect event handler calls connect again
                    state = STATE_DISCONECTED;
                    break;

                case STATE_CONECTED: {
                    // Ping gateway to check if we're still connected
                    enum dhcp_wd_test_result result = dhcp_wd_test_connection(handle->iface);

                    if (result == DHCP_WD_RESULT_PING_LOST) {
                        // looks like the gateway silently dropped us
                        // try kicking the DHCP client, if it helps
                        ESP_ERROR_CHECK(esp_netif_dhcpc_stop(handle->iface));
                        ESP_ERROR_CHECK(esp_netif_dhcpc_start(handle->iface));
                        state = STATE_CONECTED_WAIT_IP2;
                        // if not, it'll flip the whole wifi stack
                    } else {
                        ESP_LOGD(TAG, "Gateway ping OK");
                    }
                    break;
                }
            }
        }
    }
}

enum dhcp_wd_test_result dhcp_wd_test_connection(esp_netif_t *iface)
{
    ESP_LOGD(TAG, "Ping Gateway to check if IP is valid");

    ping_opts_t opts = PING_CONFIG_DEFAULT();
    opts.count = 3;
    opts.interval_ms = 0;
    opts.timeout_ms = 1000;

    esp_netif_ip_info_t ip_info = {};
    ESP_ERROR_CHECK(esp_netif_get_ip_info(iface, &ip_info));

    opts.ip_addr.addr = ip_info.gw.addr;

    ping_result_t result = {};

    if (ip_info.gw.addr != 0) {
        esp_err_t ret = ping(&opts, &result);
        if (ret != ESP_OK) {
            ESP_LOGW(TAG, "Ping error");
            return DHCP_WD_RESULT_PING_LOST;
        }
        ESP_LOGD(TAG, "Ping result: %d tx, %d rx", result.sent, result.received);
        if (result.received == 0) {
            ESP_LOGW(TAG, "Failed to ping GW");
            return DHCP_WD_RESULT_PING_LOST;
        } else {
            return DHCP_WD_RESULT_OK;
        }
    } else {
        ESP_LOGW(TAG, "No GW IP to ping");
        return DHCP_WD_RESULT_NO_GATEWAY;
    }
}