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

#include <WiFi.h>

#include <mdns.h>

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

#include "IP4.h"

WiFiUDP UDP;

static const char *ssid, *password, *hostname;
static IP4 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 inline bool setvar_ip4(enum sysvar_enum var, const IP4 &ip)
{
    return setvar_ip(var, static_cast<uint32_t>(ip));
}
static inline IP4 getvar_ip4(enum sysvar_enum var)
{
    return IP4(getvar_ip(var));
}

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

static void my_sntp_start(void)
{
    setvar_bool(status_net_sntp_sync, false);

    if (getvar_bool(config_sntp_enabled)) {
	sntp_set_time_sync_notification_cb(sntp_sync_cb);

	sntp_setoperatingmode(SNTP_OPMODE_POLL);
	sntp_servermode_dhcp(!getvar_bool(config_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 void sntp_server_show(void)
{
    IP4 sntp_ip = *sntp_getserver(0);

    if (!sntp_ip) {
	printf("[SNTP] Time server: %s\n", sntp_ip.cstr());
	setvar_ip4(status_net_sntp_server, sntp_ip);
    } else {
	setvar_ip4(status_net_sntp_server, null_ip);
    }
}

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

    if (!IP4(*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 dns_setup(void)
{
    IP4 dns_ip = *dns_getserver(0);

    if (!dns_ip || getvar_bool(config_ip4_dhcp_nodns)) {
	/* Static DNS server configuration */
	if (dns_ip != dnsserver) {
	    ip_addr_t addr = dnsserver;
	    dns_setserver(0, &addr);
	}
    }

    dns_ip = *dns_getserver(0);
    printf("[DNS]  DNS server: %s\n", dns_ip.cstr());
    setvar_ip4(status_net_dns_server, dns_ip);
}

static void mdns_setup(void)
{
    static bool mdns_started;
    static const struct mdns_service {
	const char *type, *proto;
	uint16_t port;
    } mdns_services[] = {
	{ "_http", "_tcp", 80 },
	{ NULL, NULL, 0 }
    };
    char unique_name[32];
    esp_err_t unique_mdns;

    if (mdns_started)
	mdns_free();

    if (!getvar_bool(config_mdns_enabled))
	return;

    mdns_started = mdns_init() == ESP_OK;
    if (!mdns_started)
	return;

    mdns_hostname_set(hostname);
    mdns_instance_name_set(hostname);
    printf("[MDNS] mDNS hostname: %s\n", hostname);

    unique_mdns = ESP_ERR_INVALID_STATE;
    snprintf(unique_name, sizeof unique_name, "MAX80-%s", serial_number);
    if (connected & CON_STA) {
	mdns_ip_addr_t iplist;
	iplist.addr = IP4(WiFi.localIP());
	iplist.next = NULL;
	unique_mdns = mdns_delegate_hostname_add(unique_name, &iplist);
	printf("[MDNS] mDNS unique hostname: %s\n", unique_name);
    }

    for (const struct mdns_service *svc = mdns_services; svc->type; svc++) {
	mdns_service_add(NULL, svc->type, svc->proto, svc->port, NULL, 0);
	if (unique_mdns == ESP_OK) {
	    mdns_service_add_for_host(NULL, svc->type, svc->proto,
				      unique_name, svc->port, NULL, 0);
	}
    }
}

static void start_services(void)
{
    /* Always run after (re)connect */
    dns_setup();
    mdns_setup();

    // If Arduino supported both of these at the same that would be
    // awesome, but it requires ESP-IDF reconfiguration...
    if (getvar_bool(config_sntp_enabled)) {
	IP4 sntp_ip = *sntp_getserver(0);
	if (sntp_ip) {
	    sntp_server_show();
	} else {
	    sntp_set_server(getvar_str(config_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 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;
    IP4 wifi_local_ip = WiFi.localIP();
    const char *local_ip = wifi_local_ip.cstr();

    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) {
	    setvar_bool(status_net_sta_conn, connected & CON_STA);
	    setvar_ip4(status_net_sta_ip4,
		      connected & CON_STA ? wifi_local_ip : null_ip);
	    setvar_ip4(status_net_sta_ip4_mask, WiFi.subnetMask());
	    setvar_ip4(status_net_sta_ip4_gw, WiFi.gatewayIP());
	}
	if (conn_change & CON_AP)
	    setvar_bool(status_net_ap_conn, connected & CON_AP);
	if (conn_change & CON_ETH) {
	    setvar_bool(status_net_eth_conn, connected & CON_ETH);
	    setvar_ip4(status_net_eth_ip4,
		      connected & CON_STA ? wifi_local_ip : null_ip);
	    setvar_ip4(status_net_eth_ip4_mask, WiFi.subnetMask());
	    setvar_ip4(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)
	setvar_uint(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 wifi_config_ap(void)
{
    /* No network configuration set */
    IP4 AP_IP      = IP4(192,168,0,1);
    IP4 AP_Netmask = IP4(255,255,255,0);
    IP4 AP_Gateway = IP4(0,0,0,0); // No gateway
    unsigned int channel = (time(NULL) % 11) + 1;	   // Pseudo-random
    uint8_t mac[6];
    static char ap_ssid[64];

    // The default SoftAP MAC is totally useless, so try to
    // set it to something more sensible...
    WiFi.macAddress(mac);
    uint8_t add = !!(mac[0] & 2);
    mac[0] |= 2;		// Set local bit
    mac[5] += add;		// Increment last byte if already local
    esp_wifi_set_mac(WIFI_IF_AP, mac);

    // Read it back to check what we got...
    WiFi.softAPmacAddress(mac);
    setvar_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 ", ap_ssid, AP_IP.cstr());
    printf("netmask %s channel %u\n", AP_Netmask.cstr(), channel);
    setvar_str(status_net_ap_ssid, ap_ssid);
    setvar_ip4(status_net_ap_ip4, AP_IP);
    setvar_ip4(status_net_ap_ip4_mask, AP_Netmask);
    setvar_uint(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(hostname);

    // 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
    WiFi.enableAP(!ssid);
}

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

    printf("WiFi.macAddress\n");
    WiFi.macAddress(mac);
    printf("setenv_mac\n");
    setvar_mac(status_net_sta_mac, mac);

    printf("setenv ssid\n");
    setvar_str(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 void wifi_config(void)
{
    ssid         = dupstr(getvar_str(config_wifi_ssid));
    password     = dupstr(getvar_str(config_wifi_psk));
    hostname     = dupstr(getvar_str(config_hostname));
    dnsserver    = getvar_ip4(config_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);

    setvar_str(status_hostname, 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();
}