///
/// © MiroZ 2024
///
/// Created on: Apr 17, 2019
/// Modified: 2024
///

#include <DNSServer.h>
#include "WiFi.h"
#include "Arduino.h"
#include <list>
#include <iterator>
#include "Settings.h"
#include <functional>

#include <ProvisionSoftAP.h>

static const char *TAG = "soft_ap";

using namespace std::placeholders;

extern const uint8_t provision_html_start[] asm("_binary_provision_html_start");
extern const uint8_t provision_html_end[] asm("_binary_provision_html_end");
extern const uint8_t logo_png_start[] asm("_binary_logo_png_start");
extern const uint8_t logo_png_end[] asm("_binary_logo_png_end");

const uint8_t sta_ip[] = {172, 68, 4, 1};
const uint8_t net_mask[] = {255, 255, 255, 0};
const uint8_t first_ip[] = {172, 68, 4, 10};
const uint8_t last_ip[] = {172, 68, 4, 100};


ProvisionSoftAP::ProvisionSoftAP(int port)
{
	m_rwriter = new ReaderWriter(1024);
	m_webServer = new AsyncWebServer(port);
	m_webSocket = new AsyncWebSocket("/provision");
}

void ProvisionSoftAP::wifiEvent(arduino_event_id_t event, arduino_event_info_t info)
{
	switch (event)
	{
	case ARDUINO_EVENT_WIFI_AP_STACONNECTED:
		ESP_LOGI(TAG, "station " MACSTR " joined, AID=%d", MAC2STR(info.wifi_ap_staconnected.mac), info.wifi_ap_staconnected.aid);
		break;

	case ARDUINO_EVENT_WIFI_AP_STADISCONNECTED:
		ESP_LOGI(TAG, "station " MACSTR " left", MAC2STR(info.wifi_ap_stadisconnected.mac));
		m_webSocket->closeAll();
		break;

	default:
		break;
	}
}

void ProvisionSoftAP::websocketEvent(AsyncWebSocket * server, AsyncWebSocketClient * client, 
			AwsEventType type, void * arg, uint8_t *data, size_t len)
{
	ESP_LOGI(TAG, "type:%d", type);

	if (type == WS_EVT_CONNECT)
	{
		ESP_LOGW(TAG, "Websocket client id:%d connection received", client->id());
		m_clients.push_back(client->id());
	}
	else if (type == WS_EVT_DISCONNECT)
	{
		ESP_LOGW(TAG, "Websocket client id:%d disconnected", client->id());
		m_clients.remove(client->id());
	}
	else if (type == WS_EVT_DATA)
	{
		ESP_LOGW(TAG, "Websocket cata client %d received: ", client->id());
		parseWebsocketCommand((char*) data, len);
	}
}

void ProvisionSoftAP::handlePortal(AsyncWebServerRequest* request)
{
	ESP_LOGI(TAG, "Serving html, requested url: %s", request->url().c_str());
	AsyncWebServerResponse * response = request->beginResponse_P(200, "text/html", (PGM_P)provision_html_start);
	request->send(response);
}

void ProvisionSoftAP::handleLogo(AsyncWebServerRequest* request)
{
	AsyncWebServerResponse * response = request->beginResponse_P(200, "image/png", logo_png_start, logo_png_end-logo_png_start);
	response->addHeader("Cache-Control", "no-cache, no-store, must-revalidate");
	response->addHeader("Pragma", "no-cache");
	response->addHeader("Expires", "-1");
	request->send(response);
}

/**
 * Initializes WiFi as AP, configures web, dhcp and dns servers and sets up captive portal.
 * When provisioning is done, credentials are saved and device is reset.
 *
 * ssid[in] ssid of the AP
 * password[in] password for the AP
 */
void ProvisionSoftAP::start(const char * ssid, const char * password)
{
	// start AP

	ESP_LOGW(TAG, "Starting provision...");

	/* Soft AP network parameters */
	IPAddress apIP(sta_ip);
	IPAddress netMsk(net_mask);
	IPAddress firstIp(first_ip);
	WiFi.onEvent(std::bind(&ProvisionSoftAP::wifiEvent, this, _1, _2));

	WiFi.mode(WIFI_MODE_AP);
	WiFi.softAPConfig(apIP, apIP, netMsk, firstIp);
	WiFi.softAP(ssid, password);

	delay(200);

	// setup the dns
	m_dnsServer.setErrorReplyCode(DNSReplyCode::NoError);
	m_dnsServer.setTTL(30);
	m_dnsServer.start(53, "*", apIP);

	// setup the http server
	m_webSocket->onEvent(std::bind(&ProvisionSoftAP::websocketEvent, this, _1, _2, _3, _4, _5, _6));
	m_webServer->addHandler(m_webSocket);
	m_webServer->on("/", std::bind(&ProvisionSoftAP::handlePortal, this, _1));
	// m_webServer->on("/provision.html", std::bind(&ProvisionSoftAP::handlePortal, this, _1));
	// m_webServer->on("/generate_204", std::bind(&ProvisionSoftAP::handlePortal, this, _1)); //Android captive portal. Maybe not needed. Might be handled by notFound handler.
	// m_webServer->on("/fwlink", std::bind(&ProvisionSoftAP::handlePortal, this, _1)); //Microsoft captive portal. Maybe not needed. Might be handled by notFound handler.
	m_webServer->on("/logo.png", std::bind(&ProvisionSoftAP::handleLogo, this, _1));
	m_webServer->onNotFound(std::bind(&ProvisionSoftAP::handlePortal, this, _1));

	m_webServer->begin();

	while (true)
	{
		m_dnsServer.processNextRequest();

		if (m_startWifiScan)
		{
			m_startWifiScan = false;
			ESP_LOGI(TAG, "Starting scan");
			int num_networks = WiFi.scanNetworks();
			ESP_LOGI(TAG, "Found %d networks", num_networks);
			m_rwriter->reset();
			m_rwriter->append("command:scan_wifi|");
			m_rwriter->appendf("num:%d|", num_networks);
			for (int n = 0; n < num_networks; n++)
			{
				m_rwriter->appendf("%s,%d,%d", WiFi.SSID(n), WiFi.encryptionType(n) == WIFI_AUTH_OPEN ? 0 : 1, WiFi.RSSI(n));
				if(n != num_networks - 1)
					m_rwriter->append('|');
			}
			ESP_LOGI(TAG, "Sending '%s'", m_rwriter->getBuffer());
			m_webSocket->textAll(m_rwriter->getBuffer());
			vTaskDelay(550);
		}

		if(m_tryConnect)
		{
			m_tryConnect = false;
			tryConnect();
		}

		vTaskDelay(50);
	}

	for (int i = 200; i >= 0; i--)
	{
		ESP_LOGI(TAG, "sleeping %d seconds...", i);
		vTaskDelay(1000 / portTICK_PERIOD_MS);
	}
}

void ProvisionSoftAP::start()
{
	uint8_t mac[6];
	WiFi.macAddress(mac);
	char ssid[33];

	sprintf(ssid, "Wellhub-%02x%02x%02x%02x%02x%02x", mac[0], mac[1], mac[2], mac[3], mac[4], mac[5]);
	start(ssid, "12345678");
}

const char verify_wifi[] = "command:verify_wifi";

void ProvisionSoftAP::parseWebsocketCommand(char* data, size_t len)
{
	data[len] = 0;
	ESP_LOGI(TAG, "'%s'", data);

	if(strcmp(data, "command:scan_wifi") == 0)
	{
		m_startWifiScan = true;
		ESP_LOGI(TAG, "Start scan command received");
	}
	else if(strncmp(data, verify_wifi, sizeof(verify_wifi) - 1) == 0)
	{
		char buf[80], buf1[40];
		Parser p(buf, ":");

		m_rwriter->copyString(data);

		if(!m_rwriter->elementIs(1, "ssid"))
			return;

		m_rwriter->getElementAt(1, buf, 80);
		p.getElementAt(1, buf1, 40);
		strcpy(m_ssid, buf1);
		ESP_LOGI(TAG, "ssid: '%s'", m_ssid);

		if(!m_rwriter->elementIs(2, "pwd"))
			return;

		m_rwriter->getElementAt(2, buf, 80);
		p.getElementAt(1, buf1, 40);
		strcpy(m_pwd, buf1);
		ESP_LOGI(TAG, "pwd: '*****'");

		ESP_LOGI(TAG, "got complete network credentials!");
		m_tryConnect = true;
	}
}

void ProvisionSoftAP::waitBufferEmpty()
{
	ESP_LOGI(TAG, "waitBufferEmpty start");
	bool empty = false;
	while (!empty)
	{
		empty = true;
		for (const auto& c : m_clients)
		{
			if (m_webSocket->client(c) && !m_webSocket->client(c)->isQueueEmpty())
			{
				empty = false;
				break;
			}
		}
		vTaskDelay(10);
	}
	ESP_LOGI(TAG, "waitBufferEmpty end");
}

const char wifi_ok[] = "verify_wifi:ok";
const char wifi_fail[] = "verify_wifi:fail";

void ProvisionSoftAP::tryConnect()
{
	ESP_LOGI(TAG, "trying to connect to %s", m_ssid);

	WiFi.begin(m_ssid, m_pwd);
	int connRes = WiFi.waitForConnectResult(5000);
	if(connRes == WL_CONNECTED)
	{
		// all is gud!
		ESP_LOGI(TAG, "we're connected, sending confirmation");
		m_webSocket->textAll(wifi_ok, sizeof(wifi_ok)-1);

		strcpy(SETTINGS.wifi.entry[0].ssid, m_ssid);
		strcpy(SETTINGS.wifi.entry[0].pwd, m_pwd);
		SETTINGS.wifi.selected = 0;
		SETTINGS.wifi.num = 1;
		
		SETTINGS_SAVE;

		vTaskDelay(3000 / portTICK_PERIOD_MS);
		waitBufferEmpty();

		ESP_LOGI(TAG, "restarting...");
		esp_restart();
	}
	WiFi.mode(WIFI_MODE_AP);
	ESP_LOGE(TAG, "nuh-uh peppernip, error %d", connRes);
	m_webSocket->textAll(wifi_fail, sizeof(wifi_fail)-1);
	waitBufferEmpty();

	m_webSocket->closeAll();
}

ProvisionSoftAP::~ProvisionSoftAP()
{
	delete m_webServer;
}