#include <stdlib.h>
#include <stdint.h>
#include <string.h>

#include <freertos/FreeRTOS.h>
#include <freertos/task.h>

#include <esp_system.h>
#include <esp_event.h>
#include <esp_log.h>
#include <esp_err.h>
#include <esp_wifi.h>
#include <nvs_flash.h>

#include "eos.h"
#include "net.h"
#include "wifi.h"

// XXX: WiFi fail due to no DHCP server

#define WIFI_MAX_CONNECT_ATTEMPTS   3

static const char *TAG = "EOS WIFI";

static esp_netif_t *wifi_netif;

static void wifi_event_handler(void *arg, esp_event_base_t event_base, int32_t event_id, void *event_data) {
    static int reconnect_cnt = 0;
    static int stop = 0;

    if (event_base == WIFI_EVENT) {
        switch (event_id) {
            case WIFI_EVENT_SCAN_DONE: {
                ESP_LOGI(TAG, "Event scan done");
                eos_wifi_send_scan();
                break;
            }

            case WIFI_EVENT_STA_START:
            case WIFI_EVENT_STA_STOP:
            case WIFI_EVENT_STA_CONNECTED: {
                ESP_LOGI(TAG, "Event start/stop/conneced");
                stop = 0;
                reconnect_cnt = 0;
                eos_wifi_send_status();
                break;
            }

            case WIFI_EVENT_STA_DISCONNECTED: {
                wifi_event_sta_disconnected_t *sta_disconnected;

                sta_disconnected = (wifi_event_sta_disconnected_t *)event_data;
                ESP_LOGI(TAG, "Event disconnected - reason: %d", sta_disconnected->reason);

                if (sta_disconnected->reason == WIFI_REASON_ASSOC_LEAVE) {
                    stop = 0;
                    reconnect_cnt = 0;
                    eos_wifi_send_status();
                    break;
                }
                if ((reconnect_cnt == 0) && (sta_disconnected->reason == WIFI_REASON_BEACON_TIMEOUT)) {
                    stop = 1;
                }
                if (reconnect_cnt < WIFI_MAX_CONNECT_ATTEMPTS) {
                    esp_wifi_connect();

                    reconnect_cnt++;
                } else {
                    eos_wifi_send_status();

                    reconnect_cnt = 0;
                    if (stop) {
                        esp_wifi_stop();
                        stop = 0;
                    }
                }
                break;
            }

            default:    // Ignore the other event types
                break;
        }
    } else if (event_base == IP_EVENT) {
        switch (event_id) {
            case IP_EVENT_STA_GOT_IP: {
                ip_event_got_ip_t *got_ip;

                got_ip = (ip_event_got_ip_t *)event_data;
                ESP_LOGI(TAG, "Event got IP - addres: " IPSTR, IP2STR(&got_ip->ip_info.ip));
                eos_wifi_send_status();

                /*
                if (got_ip->ip_changed) {
                    recreate all sockets
                }
                */
                break;
            }

            default:    // Ignore the other event types
                break;
        }
    }

}

static void wifi_handler(unsigned char _mtype, unsigned char *buffer, uint16_t buf_len) {
    esp_err_t ret = ESP_OK;
    uint8_t mtype;

    if (buf_len < 1) return;

    mtype = buffer[0];
    switch (mtype) {
        case EOS_WIFI_MTYPE_STATUS: {
            int reply;
            ssize_t rv;

            reply = _mtype & EOS_NET_MTYPE_FLAG_REPL;

            if (reply) {
                rv = eos_wifi_get_status(buffer + 1);
                if (rv < 0) break;

                eos_net_reply(EOS_NET_MTYPE_WIFI, buffer, rv + 1);
            } else {
                unsigned char *buf;

                buf = eos_net_alloc();
                buf[0] = EOS_WIFI_MTYPE_STATUS;
                rv = eos_wifi_get_status(buf + 1);
                if (rv < 0) break;

                eos_net_send(EOS_NET_MTYPE_WIFI, buf, rv + 1);
            }
            break;
        }

        case EOS_WIFI_MTYPE_START: {
            ret = esp_wifi_start();
            break;
        }

        case EOS_WIFI_MTYPE_STOP: {
            ret = esp_wifi_stop();
            break;
        }

        case EOS_WIFI_MTYPE_SCAN: {
            ret = esp_wifi_scan_start(NULL, 0);
            break;
        }

        case EOS_WIFI_MTYPE_CONNECT: {
            wifi_config_t wifi_sta_config;
            size_t ssid_size = sizeof(wifi_sta_config.sta.ssid);
            size_t pwd_size = sizeof(wifi_sta_config.sta.password) - 1;
            size_t ssid_len, pwd_len;
            char *ssid, *pwd;

            buffer++;
            buf_len--;

            ssid = (char *)buffer;
            ssid_len = strnlen(ssid, buf_len);
            if (ssid_len == buf_len) break;
            if (ssid_len > ssid_size) break;
            buffer += ssid_len + 1;
            buf_len -= ssid_len + 1;

            pwd = (char *)buffer;
            pwd_len = strnlen(pwd, buf_len);
            if (pwd_len == buf_len) break;
            if (pwd_len >= pwd_size) break;
            buffer += pwd_len + 1;
            buf_len -= pwd_len + 1;

            memset(&wifi_sta_config, 0, sizeof(wifi_sta_config));
            if (ssid_len < ssid_size) {
                strcpy((char *)wifi_sta_config.sta.ssid, ssid);
            } else {
                memcpy(wifi_sta_config.sta.ssid, ssid, ssid_size);
            }
            strcpy((char *)wifi_sta_config.sta.password, pwd);
            wifi_sta_config.sta.password[pwd_size] = '\0';

            ret = esp_wifi_set_config(WIFI_IF_STA, &wifi_sta_config);
            if (ret) break;

            ret = esp_wifi_connect();
            break;
        }

        case EOS_WIFI_MTYPE_DISCONNECT: {
            ret = esp_wifi_disconnect();
            break;
        }
    }

    if (ret) {
        ESP_LOGE(TAG, "HANDLER ERR:%d MTYPE:%d", ret, mtype);
        eos_wifi_send_status();
    }
}

void eos_wifi_init(void) {
    esp_err_t ret;
    wifi_init_config_t wifi_config = WIFI_INIT_CONFIG_DEFAULT();

    wifi_netif = esp_netif_create_default_wifi_sta();

    ret = esp_wifi_init(&wifi_config);
    assert(ret == ESP_OK);

    ret = esp_event_handler_instance_register(WIFI_EVENT, ESP_EVENT_ANY_ID, &wifi_event_handler, NULL, NULL);
    assert(ret == ESP_OK);

    ret = esp_event_handler_instance_register(IP_EVENT, IP_EVENT_STA_GOT_IP, &wifi_event_handler, NULL, NULL);
    assert(ret == ESP_OK);

    ret = esp_wifi_set_storage(WIFI_STORAGE_RAM);
    assert(ret == ESP_OK);

    ret = esp_wifi_set_mode(WIFI_MODE_STA);
    assert(ret == ESP_OK);

    ret = esp_wifi_stop();
    assert(ret == ESP_OK);

    eos_net_set_handler(EOS_NET_MTYPE_WIFI, wifi_handler);
    ESP_LOGI(TAG, "INIT");
}

ssize_t eos_wifi_get_status(unsigned char *buffer) {
    unsigned char *p;
    wifi_ap_record_t ap_info;
    esp_err_t ret;

    p = buffer;
    ret = esp_wifi_sta_get_ap_info(&ap_info);
    switch (ret) {
        case ESP_OK: {  // Station connected
            size_t len;
            esp_netif_ip_info_t ip_info;

            len = strnlen((char *)ap_info.ssid, sizeof(ap_info.ssid));
            if (len == sizeof(ap_info.ssid)) {
                ESP_LOGE(TAG, "STATUS BAD SSID");
                return EOS_ERR;
            }
            ret = esp_netif_get_ip_info(wifi_netif, &ip_info);
            if (ret) {
                ESP_LOGE(TAG, "STATUS GET IP INFO ERR:%d", ret);
                return EOS_ERR;
            }

            if (ip_info.ip.addr) {
                *p = EOS_WIFI_STATUS_GOT_IP;
                p++;

                memcpy(p, &ip_info.ip.addr, sizeof(uint32_t));
                p += sizeof(uint32_t);
            } else {
                *p = EOS_WIFI_STATUS_CONNECTED;
                p++;
            }

            strcpy((char *)p, (char *)ap_info.ssid);
            p += len + 1;

            break;
        }

        case ESP_ERR_WIFI_NOT_CONNECT: {
            *p = EOS_WIFI_STATUS_DISCONNECTED;
            p++;
            break;
        }

        case ESP_ERR_WIFI_STOP_STATE:
        case ESP_ERR_WIFI_CONN: {
            *p = EOS_WIFI_STATUS_OFF;
            p++;
            break;
        }

        default: {
            ESP_LOGE(TAG, "STATUS GET AP INFO ERR:%d", ret);
            return EOS_ERR;
        }
    }

    return p - buffer;
}

void eos_wifi_send_status(void) {
    unsigned char *rbuf;
    ssize_t rv;

    rbuf = eos_net_alloc();
    rbuf[0] = EOS_WIFI_MTYPE_STATUS;
    rv = eos_wifi_get_status(rbuf + 1);
    if (rv < 0) {
        eos_net_free(rbuf);
        return;
    }

    eos_net_send(EOS_NET_MTYPE_WIFI, rbuf, rv + 1);
}

void eos_wifi_send_scan(void) {
    static wifi_ap_record_t scan_r[EOS_WIFI_MAX_SCAN_RECORDS];
    static uint16_t scan_n;
    unsigned char *rbuf, *p;
    int i;
    size_t len;
    esp_err_t ret;

    scan_n = EOS_WIFI_MAX_SCAN_RECORDS;
    memset(scan_r, 0, sizeof(scan_r));
    ret = esp_wifi_scan_get_ap_records(&scan_n, scan_r);
    if (ret) {
        ESP_LOGE(TAG, "SCAN GET RECORDS ERR:%d", ret);
        return;
    }

    ESP_LOGI(TAG, "Scan done: %d", scan_n);

    rbuf = eos_net_alloc();
    rbuf[0] = EOS_WIFI_MTYPE_SCAN;
    p = rbuf + 1;

    for (i=0; i<scan_n; i++) {
        len = strnlen((char *)scan_r[i].ssid, sizeof(scan_r[i].ssid));
        if (len == sizeof(scan_r[i].ssid)) continue;
        if (p - rbuf + len + 1 > EOS_NET_MTU) break;

        strcpy((char *)p, (char *)scan_r[i].ssid);
        p += len + 1;
    }
    eos_net_send(EOS_NET_MTYPE_WIFI, rbuf, p - rbuf);
}