#include "common.h"
#include "WiFi.h"
#include "wifi.h"
#include "config.h"
#include "httpd.h"
#include "led.h"

#include <lwip/dns.h>
#include <lwip/inet.h>
#include <lwip/apps/sntp.h>
#include <esp_sntp.h>
#include <esp_wifi.h>

static const char *ssid, *password, *hostname, *dnsserver;
static TimerHandle_t sta_failure_timer;

enum connected {
    CON_STA = 1,
    CON_ETH = 2,
    CON_AP  = 4
};
static volatile bool sta_timeout_enabled;
static volatile unsigned int sta_timeout_count;
static unsigned int connected;

static void sta_bounce(void)
{
    if (!ssid)
	return;

    if (WiFi.status() != WL_CONNECTED) {
	WiFi.disconnect();
	WiFi.begin();
    }
}

static void sta_timeout(void)
{
    unsigned int count = ++sta_timeout_count;

    if (connected & CON_STA)
	return;

    // Try to ping the STA even if there are AP clients every
    // 15 seconds (if an SSID is configured)
    sta_bounce();

    if (!(connected & (CON_AP|CON_STA)) && (count >= 2)) {
	// Enable the AP if the STA doesn't connect after 30s
	WiFi.enableAP(true);
    }
}

static void sta_timeout_enable(void)
{
    if (!sta_failure_timer || sta_timeout_enabled)
	return;

    sta_timeout_enabled = xTimerStart(sta_failure_timer, 1);
    sta_timeout_count = 0;
}

static void sta_timeout_disable(void)
{
    if (!sta_failure_timer || !sta_timeout_enabled)
	return;

    sta_timeout_enabled = !xTimerStop(sta_failure_timer, 1);
}

static void sntp_sync_cb(struct timeval *tv)
{
    static uint8_t prev_sync_status = SNTP_SYNC_STATUS_RESET;
    uint8_t sync_status = sntp_get_sync_status();

    switch (sync_status) {
    case SNTP_SYNC_STATUS_RESET:
	sntp_set_sync_mode(SNTP_SYNC_MODE_IMMED); // Until first sync
	time_net_sync(NULL);
	break;
    case SNTP_SYNC_STATUS_COMPLETED:
	sntp_set_sync_mode(SNTP_SYNC_MODE_SMOOTH); // After successful sync
	time_net_sync(tv);
	break;
    default:
	break;
    }

    prev_sync_status = sync_status;
}

static void my_sntp_start(void)
{
    setenv_bool("status.net.sntp.sync", false);

    if (getenv_bool("sntp.enabled")) {
	sntp_set_time_sync_notification_cb(sntp_sync_cb);

	sntp_setoperatingmode(SNTP_OPMODE_POLL);
	sntp_servermode_dhcp(!getenv_bool("ip4.dhcp.nosntp"));
	sntp_set_sync_mode(SNTP_SYNC_MODE_IMMED); // Until first sync
	sntp_init();
    } else {
	sntp_stop();
    }
}

static bool services_started;

static void stop_services(void)
{
    if (services_started) {
	esp_unregister_shutdown_handler(stop_services);
	printf("[WIFI] Stopping network services\n");
	my_httpd_stop();
	services_started = false;
    }
}

static inline bool invalid_ip(const ip_addr_t *ip)
{
    return !memcmp(ip, &ip_addr_any, sizeof *ip);
}

static void sntp_server_show(void)
{
    const ip_addr_t *sntp_ip = sntp_getserver(0);

    if (!invalid_ip(sntp_ip)) {
	const char *sntp_server = inet_ntoa(*sntp_ip);
	printf("[SNTP] Time server: %s\n", sntp_server);
	setenv_cond("status.net.sntp.server", sntp_server);
    } else {
	setenv_cond("status.net.sntp.server", NULL);
    }
}

static void sntp_server_found(const char *name, const ip_addr_t *addr,
			      void *arg)
{
    (void)name;
    (void)arg;

    if (invalid_ip(addr))
	return;

    sntp_setserver(0, addr);
    sntp_server_show();
}

static void sntp_set_server(const char *name)
{
    if (!name || !*name)
	return;

    ip_addr_t addr;
    err_t err = dns_gethostbyname(name, &addr, sntp_server_found, NULL);
    if (err == ERR_OK)
	sntp_server_found(name, &addr, NULL);
}

static void start_services(void)
{
    /* Always run after (re)connect */
    const ip_addr_t *dns_ip = dns_getserver(0);
    if (invalid_ip(dns_ip) || getenv_bool("ip4.dhcp.nodns")) {
	/* Static DNS server configuration */
	ip_addr_t addr;
	if (dnsserver && inet_aton(dnsserver, &addr)) {
	    if (memcmp(dns_ip, &addr, sizeof addr))
		dns_setserver(0, &addr);
	}
    }

    {
	dns_ip = dns_getserver(0);
	const char *dns_server_str = inet_ntoa(*dns_ip);
	printf("[DNS]  DNS server: %s\n", dns_server_str);
	setenv_cond("status.net.dns.server", dns_server_str);
    }

    // If Arduino supported both of these at the same that would be
    // awesome, but it requires ESP-IDF reconfiguration...
    if (getenv_bool("sntp.enabled")) {
	if (!invalid_ip(sntp_getserver(0))) {
	    sntp_server_show();
	} else {
	    sntp_set_server(getenv("sntp.server"));
	}
    }

    /* Only run on first start */
    if (!services_started) {
	services_started = true;
	printf("[WIFI] Starting network services\n");
	my_httpd_start();
	esp_register_shutdown_handler(stop_services);
    }
}

static const char *ip_str(const IPAddress &ip)
{
    static char ip_buf[4*4];
    const IPAddress ip_none(0,0,0,0);
    return strcpy(ip_buf, ip == ip_none ? "" : ip.toString().c_str());
}

static void setenv_ip(const char *var, const IPAddress &ip)
{
    setenv_config(var, ip_str(ip));
}

static bool force_conn_update;

static void WiFiEvent(WiFiEvent_t event, WiFiEventInfo_t info)
{
    bool retry_sta  = false;
    bool is_connect = false;
    unsigned int prev_connected = connected;
    static int ap_clients;
    int prev_ap_clients = ap_clients;
    IPAddress wifi_local_ip = WiFi.localIP();
    const char *local_ip = ip_str(wifi_local_ip);

    switch (event) {
    case ARDUINO_EVENT_WIFI_READY:
	printf("[WIFI] Interface ready\n");
	break;
    case ARDUINO_EVENT_WIFI_SCAN_DONE:
	printf("[WIFI] Completed scan for access points\n");
	break;
    case ARDUINO_EVENT_WIFI_STA_START:
	printf("[WIFI] Client mode started\n");
	break;
    case ARDUINO_EVENT_WIFI_STA_STOP:
	printf("[WIFI] Client mode stopped\n");
	connected &= ~CON_STA;
	break;
    case ARDUINO_EVENT_WIFI_STA_CONNECTED:
	printf("[WIFI] Connected to access point\n");
	break;
    case ARDUINO_EVENT_WIFI_STA_DISCONNECTED:
	printf("[WIFI] Disconnected from WiFi access point\n");
	connected &= ~CON_STA;
	retry_sta = true;
	break;
    case ARDUINO_EVENT_WIFI_STA_AUTHMODE_CHANGE:
	printf("[WIFI] Authentication mode of access point has changed\n");
	break;
    case ARDUINO_EVENT_WIFI_STA_GOT_IP:
    {
	printf("[WIFI] Obtained IP address: %s\n", local_ip);
	connected |= CON_STA;
	is_connect = true;
	break;
    }
    case ARDUINO_EVENT_WIFI_STA_LOST_IP:
    {
	printf("[WIFI] Lost IP address\n");
	connected &= ~CON_STA;
	retry_sta = true;
	break;
    }
    case ARDUINO_EVENT_WPS_ER_SUCCESS:
	printf("[WIFI] WiFi Protected Setup (WPS): succeeded in enrollee mode\n");
	break;
    case ARDUINO_EVENT_WPS_ER_FAILED:
	printf("[WIFI] WiFi Protected Setup (WPS): failed in enrollee mode\n");
	break;
    case ARDUINO_EVENT_WPS_ER_TIMEOUT:
	printf("[WIFI] WiFi Protected Setup (WPS): timeout in enrollee mode\n");
	break;
    case ARDUINO_EVENT_WPS_ER_PIN:
	printf("[WIFI] WiFi Protected Setup (WPS): pin code in enrollee mode\n");
	break;
    case ARDUINO_EVENT_WIFI_AP_START:
	printf("[WIFI] Access point started\n");
	ap_clients  = 0;
	connected  |= CON_AP;
	is_connect  = true;
	break;
    case ARDUINO_EVENT_WIFI_AP_STOP:
	printf("[WIFI] Access point stopped\n");
	connected  &= ~CON_AP;
	ap_clients  = 0;
	break;
    case ARDUINO_EVENT_WIFI_AP_STACONNECTED:
	printf("[WIFI] Client connected\n");
	ap_clients = WiFi.softAPgetStationNum();
	break;
    case ARDUINO_EVENT_WIFI_AP_STADISCONNECTED:
	printf("[WIFI] Client disconnected\n");
	ap_clients = WiFi.softAPgetStationNum();
	break;
    case ARDUINO_EVENT_WIFI_AP_STAIPASSIGNED:
	printf("[WIFI] Assigned IP address %s to client\n",
	       inet_ntoa(info.wifi_ap_staipassigned.ip));
	ap_clients = WiFi.softAPgetStationNum();
	is_connect = true;
	break;
    case ARDUINO_EVENT_WIFI_AP_PROBEREQRECVED:
	printf("[WIFI] Received probe request\n");
	ap_clients = WiFi.softAPgetStationNum();
	break;
    case ARDUINO_EVENT_WIFI_AP_GOT_IP6:
	printf("[WIFI] AP IPv6 is preferred\n");
	ap_clients = WiFi.softAPgetStationNum();
	is_connect = true;
	break;
    case ARDUINO_EVENT_WIFI_STA_GOT_IP6:
	printf("[WIFI] STA IPv6 is preferred\n");
	is_connect = true;
	break;
    case ARDUINO_EVENT_ETH_GOT_IP6:
	printf("[ETH]  Ethernet IPv6 is preferred\n");
	is_connect = true;
	break;
    case ARDUINO_EVENT_ETH_START:
	printf("[ETH]  Ethernet started\n");
	break;
    case ARDUINO_EVENT_ETH_STOP:
	printf("[ETH]  Ethernet stopped\n");
	connected &= ~CON_ETH;
	break;
    case ARDUINO_EVENT_ETH_CONNECTED:
	printf("[ETH]  Ethernet connected\n");
	break;
    case ARDUINO_EVENT_ETH_DISCONNECTED:
	printf("[ETH]  Ethernet disconnected\n");
	connected &= ~CON_ETH;
	retry_sta = true;
	break;
    case ARDUINO_EVENT_ETH_GOT_IP:
	printf("[ETH]  Obtained IP address: %s\n", local_ip);
	connected |= CON_ETH;
	is_connect = true;
	break;
    default:
	break;
    }

    if (connected & ~CON_AP) {
	sta_timeout_disable();
	if (!ap_clients) {
	    WiFi.enableAP(false);
	}
    } else if (ssid) {
	sta_timeout_enable();
    }

    unsigned int conn_change = force_conn_update ?
	-1U : (connected ^ prev_connected);

    if (conn_change) {
	force_conn_update = false;

	if (conn_change & CON_STA) {
	    setenv_bool("status.net.sta.conn", connected & CON_STA);
	    setenv_config("status.net.sta.ip4",
			  connected & CON_STA ? local_ip : "");
	    setenv_ip("status.net.sta.ip4.mask", WiFi.subnetMask());
	    setenv_ip("status.net.sta.ip4.gw", WiFi.gatewayIP());
	}
	if (conn_change & CON_AP)
	    setenv_bool("status.net.ap.conn", connected & CON_AP);
	if (conn_change & CON_ETH) {
	    setenv_bool("status.net.eth.conn", connected & CON_ETH);
	    setenv_config("status.net.eth.ip4",
			  connected & CON_ETH ? local_ip : "");
	    setenv_ip("status.net.eth.ip4.mask", WiFi.subnetMask());
	    setenv_ip("status.net.eth.ip4.gw", WiFi.gatewayIP());
	}

	if (!ssid) {
	    // No network configured
	    led_set(LED_GREEN, connected & CON_AP ? LED_FLASH_SLOW : LED_OFF);
	} else {
	    led_set(LED_GREEN, connected & CON_AP ? LED_FLASH_FAST : LED_ON);
	}
    }

    if (is_connect) {
	start_services();
    }

    if (ap_clients != prev_ap_clients)
	setenv_ul("status.net.ap.clients", ap_clients);

    /*
     * Don't keep retrying if there are AP clients - makes the AP
     * too unreliable to use.
     */
    if (retry_sta && !ap_clients) {
	sta_bounce();
    }
}

static void setenv_mac(const char *var, const uint8_t mac[6])
{
    char mac_str[3*6];

    snprintf(mac_str, sizeof mac_str, "%02x:%02x:%02x:%02x:%02x:%02x",
	     mac[0], mac[1], mac[2], mac[3], mac[4], mac[5]);
    setenv_cond(var, mac_str);
}

static void wifi_config_ap(void)
{
    /* No network configuration set */
    IPAddress AP_IP      = IPAddress(192,168,0,1);
    IPAddress AP_Netmask = IPAddress(255,255,255,0);
    IPAddress AP_Gateway = IPAddress(0,0,0,0); // No gateway
    unsigned int channel = (time(NULL) % 11) + 1;	   // Pseudo-random
    uint8_t mac[6];
    char mac_str[6*3];
    static char ap_ssid[64];

    WiFi.softAPmacAddress(mac);
    setenv_mac("status.net.ap.mac", mac);

    /* The last two bytes of the efuse MAC */
    snprintf(ap_ssid, sizeof ap_ssid, "MAX80_%02X%02X",
	     efuse_default_mac[4], efuse_default_mac[5]);

    printf("[WIFI] AP SSID %s IP %s netmask %s channel %u\n",
	       ap_ssid, AP_IP.toString(), AP_Netmask.toString(), channel);
    setenv_cond("status.net.ap.ssid", ap_ssid);
    setenv_ip("status.net.ap.ip4", AP_IP);
    setenv_ip("status.net.ap.ip4.mask", AP_Netmask);
    setenv_ul("status.net.ap.clients", 0);

    printf("WiFi.softAP\n");
    WiFi.softAP(ap_ssid, NULL, channel, 0, 4, true);
    printf("WiFi.softAPConfig\n");
    WiFi.softAPConfig(AP_IP, AP_Gateway, AP_Netmask);
    printf("WiFi.softAPsetHostname\n");
    WiFi.softAPsetHostname("max80");

    // Conservative setting: 20 MHz (single channel) only; this is for
    // reliability, not performance.
    printf("esp_wifi_set_bandwidth\n");
    esp_wifi_set_bandwidth((wifi_interface_t)ESP_IF_WIFI_AP, WIFI_BW_HT20);

    // Enable AP immediately if no SSID configured
    printf("WiFi.enableAP\n");
    WiFi.enableAP(!ssid);
}

static void wifi_config_sta(void)
{
    uint8_t mac[6];

    printf("WiFi.macAddress\n");
    WiFi.macAddress(mac);
    printf("setenv_mac\n");
    setenv_mac("status.net.sta.mac", mac);

    printf("setenv ssid\n");
    setenv_cond("status.net.sta.ssid", ssid);
    if (!ssid) {
	WiFi.enableSTA(false);
	return;
    }

    printf("xTimerCreate\n");
    sta_failure_timer = xTimerCreate("wifi_sta", configTICK_RATE_HZ*15,
				     pdTRUE, NULL,
				     (TimerCallbackFunction_t)sta_timeout);
    printf("sta_timeout_enable\n");
    sta_timeout_enable();

    printf("WiFi.setAutoReconnect\n");
    WiFi.setAutoReconnect(false); // We are doing this "ourselves"
    printf("WiFi.enableSTA\n");
    WiFi.enableSTA(true);
    printf("WiFi.begin(%s)\n", ssid);
    WiFi.begin(ssid, password);
}

static const char *getenv_notempty(const char *env)
{
    const char *str = getenv(env);
    if (str && !*str)
	str = NULL;
    return str;
}

static void wifi_config(void)
{
    ssid         = getenv_notempty("wifi.ssid");
    password     = getenv_notempty("wifi.psk");
    hostname     = getenv_notempty("hostname");
    dnsserver    = getenv_notempty("ip4.dns");

    force_conn_update = true;

    printf("WiFi.persistent\n");
    WiFi.persistent(false);
    printf("WiFi.setSleep\n");
    WiFi.setSleep(false);

    WiFi.setTxPower(WIFI_POWER_19_5dBm);

    if (hostname)
	WiFi.hostname(hostname);

    printf("wifi_config_ap\n");
    wifi_config_ap();
    printf("wifi_config_sta\n");
    wifi_config_sta();
    printf("wifi_config done\n");
}

void SetupWiFi() {
    services_started = false;

    WiFi.onEvent(WiFiEvent);

    my_sntp_start();

    printf("[INFO] Setting up WiFi\n");
    wifi_config();
}