#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 String ssid, password, hostname, dnsserver, sntpserver;
static TimerHandle_t sta_failure_timer;

static bool sta_timeout_enabled;

static void sta_timeout(void)
{
    // Enable the AP if the STA doesn't connect after a timeout
    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);
}

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()
{
    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);
	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 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.c_str(), &addr)) {
	    if (memcmp(dns_ip, &addr, sizeof addr))
		dns_setserver(0, &addr);
	}
    }

    dns_ip = dns_getserver(0);
    printf("[DNS]  DNS server: %s\n", inet_ntoa(*dns_ip));
    setenv("status.net.dns.server", inet_ntoa(*dns_ip), 1);

    // If Arduino supported both of these at the same that would be
    // awesome, but it requires ESP-IDF reconfiguration...
    if (sntp_enabled()) {
	const ip_addr_t *sntp_ip = sntp_getserver(0);

	if (invalid_ip(sntp_ip)) {
	    if (sntpserver != "") {
		sntp_setservername(0, sntpserver.c_str());
		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("status.net.sntp.server", sntp_server, 1);
	} else {
	    unsetenv("status.net.sntp.server");
	}
    }

    /* Only run on first start */
    if (!services_started) {
	services_started = true;
	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;
    enum connected {
	CON_STA = 1,
	CON_ETH = 2,
	CON_AP  = 4
    };
    static unsigned int connected;
    unsigned int prev_connected = connected;
    static int ap_clients = 0;
    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 started\n");
	break;
    case ARDUINO_EVENT_WIFI_STA_STOP:
	printf("[WIFI] Clients 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++;
	break;
    case ARDUINO_EVENT_WIFI_AP_STADISCONNECTED:
	printf("[WIFI] Client disconnected\n");
	ap_clients--;
	break;
    case ARDUINO_EVENT_WIFI_AP_STAIPASSIGNED:
	printf("[WIFI] Assigned IP address %s to client\n",
	       inet_ntoa(info.wifi_ap_staipassigned.ip));
	break;
    case ARDUINO_EVENT_WIFI_AP_PROBEREQRECVED:
	printf("[WIFI] Received probe request\n");
	break;
    case ARDUINO_EVENT_WIFI_AP_GOT_IP6:
	printf("[WIFI] AP IPv6 is preferred\n");
	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);

    if (retry_sta) {
	WiFi.disconnect();
	WiFi.begin();
    }
}

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(var, mac_str, 1);
}

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;	   // Pseudo-random
    uint8_t mac[6];
    char mac_str[6*3];
    char ap_ssid[64];

    WiFi.softAPmacAddress(mac);
    setenv_mac("status.net.ap.mac", mac);
    /* The last two bytes of the MAC */
    snprintf(ap_ssid, sizeof ap_ssid, "MAX80_%02X%02X", mac[4], mac[5]);

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

    WiFi.softAP(ap_ssid, NULL, channel+1, 0, 4, true);
    WiFi.softAPConfig(AP_IP, AP_Gateway, AP_Netmask);
    WiFi.softAPsetHostname(ap_ssid);

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

    // Enable unconditionally if no SSID
    WiFi.enableAP(ssid == "");
}

static void wifi_config_sta(void)
{
    uint8_t mac[6];
    WiFi.macAddress(mac);
    setenv_mac("status.net.sta.mac", mac);

    setenv("status.net.sta.ssid", ssid.c_str(), 1);
    if (ssid == "") {
	WiFi.enableSTA(false);
	return;
    }
    sta_failure_timer = xTimerCreate("wifi_sta", configTICK_RATE_HZ*30,
				     pdFALSE, NULL,
				     (TimerCallbackFunction_t)sta_timeout);

    sta_timeout_enable();

    WiFi.begin(ssid.c_str(), password.c_str());
    WiFi.setAutoConnect(true);
    WiFi.setAutoReconnect(true);
    WiFi.enableSTA(true);
}

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

    force_conn_update = true;

    WiFi.persistent(false);
    WiFi.setSleep(false);

    if (hostname != "")
	WiFi.hostname(hostname);

    wifi_config_sta();
    wifi_config_ap();
}

void SetupWiFi() {
    services_started = false;

    WiFi.onEvent(WiFiEvent);

    my_sntp_start();

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