#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 "eos.h"
#include "net.h"
#include "tun.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) {
                    eos_wifi_send_status();
                    eos_tun_portmap_remove();
                    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();
                    eos_tun_portmap_remove();

                    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));
                ESP_LOGI(TAG, "Event got IP - addres: %x", (unsigned int)got_ip->ip_info.ip.addr);
                eos_wifi_send_status();
                eos_tun_portmap_remove();
                eos_tun_portmap_add(got_ip->ip_info.ip.addr);

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

                break;
            }

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

}

static void wifi_handler(unsigned char _mtype, EOSMessage *msg, uint16_t len) {
    unsigned char *buffer = msg->buffer;
    uint16_t buf_size = msg->size;
    esp_err_t ret = ESP_OK;
    unsigned char mtype;

    if (len < 1) return;

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

            if (!(eos_msg_flags(msg) & EOS_MSG_FLAG_RPLY_REQ)) break;

            rv = eos_wifi_get_status(buffer + 1, buf_size - 1);
            if (rv < 0) break;

            eos_net_reply(EOS_NET_MTYPE_WIFI, msg, 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);
            unsigned char ssid_len, pwd_len;
            char *ssid, *pwd;

            buffer++;

            ssid_len = *buffer;
            buffer++;

            ssid = (char *)buffer;
            if (ssid_len > ssid_size) break;
            buffer += ssid_len;

            pwd_len = *buffer;
            buffer++;

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

            memset(&wifi_sta_config, 0, sizeof(wifi_sta_config));
            memcpy(wifi_sta_config.sta.ssid, ssid, ssid_len);
            if (ssid_len < ssid_size) {
                wifi_sta_config.sta.ssid[ssid_len] = '\0';
            }
            memcpy((char *)wifi_sta_config.sta.password, pwd, pwd_len);
            wifi_sta_config.sta.password[pwd_len] = '\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) {
    wifi_init_config_t wifi_config = WIFI_INIT_CONFIG_DEFAULT();
    esp_err_t ret;

    wifi_netif = esp_netif_create_default_wifi_sta();

    ret = esp_wifi_init(&wifi_config);
    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");
}

void eos_wifi_run(void) {
    esp_err_t ret;

    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);

    ESP_LOGI(TAG, "RUN");
}

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

    /* 1 + sizeof(uint32_t) + sizeof(ap_info.ssid) - 1 */
    if (size < sizeof(uint32_t) + sizeof(ap_info.ssid)) return EOS_ERR_SIZE;

    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++;
            }

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

            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) {
    EOSMessage msg;
    ssize_t _rv;
    int rv;

    eos_net_alloc(&msg);
    if (msg.size < 1) {
        eos_net_free(&msg);
        return;
    }
    msg.buffer[0] = EOS_WIFI_MTYPE_STATUS;
    _rv = eos_wifi_get_status(msg.buffer + 1, msg.size - 1);
    if (_rv < 0) {
        eos_net_free(&msg);
        return;
    }

    rv = eos_net_send(EOS_NET_MTYPE_WIFI, &msg, _rv + 1);
    if (rv) ESP_LOGE(TAG, "NET SEND ERR:%d", rv);
}

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

    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);

    eos_net_alloc(&msg);
    if (msg.size < 1) {
        eos_net_free(&msg);
        return;
    }
    msg.buffer[0] = EOS_WIFI_MTYPE_SCAN;
    p = msg.buffer + 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 - msg.buffer + len + 1 > msg.size) break;

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

    rv = eos_net_send(EOS_NET_MTYPE_WIFI, &msg, p - msg.buffer);
    if (rv) ESP_LOGE(TAG, "NET SEND ERR:%d", rv);
}