#pragma once

#include <ESP8266WiFi.h>
#include <ESP8266mDNS.h>
#include <ESP8266WebServer.h>
#include <ArduinoOTA.h>
#include <PubSubClient.h>
#include <cstdlib>
#include <list>

typedef std::function<bool(String)> IOTActionHandlerFunction;
typedef std::function<String()> IOTReportHandlerFunction;

struct IOTActionHandler {
	String topic;
	IOTActionHandlerFunction function;
};

struct IOTReportHandler {
	String topic;
	IOTReportHandlerFunction function;
	bool cache_only;
	unsigned long update_interval;
	unsigned long next_update_at;
	String last_result;
};

class SimpleIOT {
private:
	WiFiClient _wifi;
	PubSubClient _mqtt_client;
	ESP8266WebServer _http_server;
	const char* _wifi_ssid;
	const char* _wifi_pass;
	String _hostname;
	std::list<IOTActionHandler> _action_handlers;
	std::list<IOTReportHandler> _report_handlers;
	bool _mqtt_enabled = false;
	long _mqtt_last_reconnect_attempt = 0;
	const char* _mqtt_host;
	uint16_t _mqtt_port;
	bool _mqtt_auth = false;
	const char* _mqtt_user;
	const char* _mqtt_pass;
	const char* _mqtt_topic;
	unsigned long _startup_complete_at = 0;
	bool _initialized = false;

	void _setup();
	void _wifi_setup();
	void _ota_setup();
	void _ota_loop();
	bool _mqtt_connect();
	void _mqtt_setup();
	void _mqtt_loop();
	void _mqtt_callback(char* topic, byte* pl, unsigned int len);
	void _http_setup();
	void _http_loop();
	void _mqtt_publish_report(String topic, String report);

	bool _handle_action_request(String topic, String payload);
	String _handle_report_request(String topic);
	bool _topic_in_use(String topic);
	bool _has_action_handler(String topic);
	bool _has_report_handler(String topic);
	void _check_report_handlers();
	String _get_report_handler_result(IOTReportHandler* h, bool force_update);
public:
	SimpleIOT(const char* wifi_ssid, const char* wifi_pass, String hostname, bool append_chip_id=false);
	void setStartupDelay(uint16_t delay);
	void setMQTTData(const char* mqtt_host, uint16_t mqtt_port,
	                 const char* mqtt_user, const char* mqtt_pass,
	                 const char* mqtt_topic);
	void setMQTTData(const char* mqtt_host, uint16_t mqtt_port,
	                 const char* mqtt_topic);
	void begin();
	void loop();
	bool act_on(String topic, IOTActionHandlerFunction f);
	bool report_on(String topic, IOTReportHandlerFunction f, unsigned long update_interval, bool use_cache);
	void log(const char* fmt, ...) __attribute__((format (printf, 2, 3)));
	void publish(String topic, String payload, bool retain=false);
};

/**
 * Constructs a new instance of SimpleIOT.
 *
 * @param wifi_ssid SSID of the WiFi to join.
 * @param wifi_pass Password of the WiFi to join.
 * @param hostname The hostname this device will use to identify itself.
 *                 It will be used as WiFi client name and for the mDNS stuff
 *                 of the OTA daemon.
 * @param append_chip_id If true, the ESP's chipID will be appended to the hostname.
 *                       This can be useful in networks with more than one ESPs with
 *                       the given hostname. `test` becomes `test-13D93B0A`.
 */
SimpleIOT::SimpleIOT(const char* wifi_ssid, const char* wifi_pass, String hostname, bool append_chip_id) {
	Serial.begin(74880);
	_wifi_ssid = wifi_ssid;
	_wifi_pass = wifi_pass;
	if (append_chip_id) {
		char* temp = new char[9];
		snprintf(temp, 9, "-%08X", ESP.getChipId());
		hostname.concat(temp);
		delete temp;
	}
	_hostname = hostname;
}

/**
 * Defines a startup delay to give a chance of doing an OTA update.
 * During development there are situations where a bug in code leads to quickly
 * rebooting devices. Since this can be quite a nuisance when working with non-
 * local devices, you can set this to delay the startup a bit to give yourself a
 * chance of starting an OTA upgrade at the right moment.
 * Startup will be more or less paused with only OTA being active. This happens
 * in begin().
 *
 * @param delay Delay in seconds to wait during startup. Set to 0 to disable
                  (default).
 *
 * @warning This method has to be called before calling begin().
 */
void SimpleIOT::setStartupDelay(uint16_t delay) {
	if (_initialized) return;
	if (delay > 0) {
		_startup_complete_at = millis() + delay * 1000;
	} else {
		_startup_complete_at = 0;
	}
}

/**
 * Sets connection data for use with an MQTT server. This also *enables* using
 * an MQTT server - not calling this method will lead to SimpleIOT working
 * happily without MQTT.
 * This is for connecting with authentication.
 *
 * @param host Hostname or IP address of an MQTT server.
 * @param port Port of the MQTT server (usually 1883 for unencrypted connections).
 * @param user Username for authentication.
 * @param pass Password for authentication.
 * @param topic Base topic to publish and subscribe to. Must end with a slash.
 *
 * @warning This method has to be called before calling begin().
 */
 
void SimpleIOT::setMQTTData(const char* host, uint16_t port,
                            const char* user, const char* pass,
                            const char* topic) {
	if (_initialized) return;
	setMQTTData(host, port, topic);
	_mqtt_user = user;
	_mqtt_pass = pass;
	_mqtt_auth = true;
}

/**
 * Sets connection data for use with an MQTT server. This also *enables* using
 * an MQTT server - not calling this method will lead to SimpleIOT working
 * happily without MQTT.
 * This is for connecting anonymously.
 *
 * @param host Hostname or IP address of an MQTT server.
 * @param port Port of the MQTT server (usually 1883 for unencrypted connections).
 * @param topic Base topic to publish and subscribe to. Must end with a slash.
 *
 * @warning This method has to be called before calling begin().
 */
 
void SimpleIOT::setMQTTData(const char* host, uint16_t port,
                            const char* topic) {
	if (_initialized) return;
	_mqtt_host = host;
	_mqtt_port = port;
	_mqtt_topic = topic;
	_mqtt_enabled = true;
}

/**
 * Starts connecting to services.
 * 1. A Wifi connection is attempted. If it fails, the ESP is restarted.
 * 2. OTA is set up.
 * 3. If setStartupDelay() was used to set a startup delay, this delay will be
 *    spent waiting for a possible upgrade via OTA.
 * 4. If MQTT was enabled using setMQTTData(), a connection to the server will
 *    be initialized.
 * 5. A HTTP server will be configured and started on port 80.
 *
 * @note You should call this method in your setup() block.
 */
void SimpleIOT::begin() {
	log("Core * Setting up Wifi...");
	_wifi_setup();
	log("Core * Setting up OTA...");
	_ota_setup();
	
	while (_startup_complete_at > millis()) {
		log("Core * Startup delay remaining: %ld ms", _startup_complete_at - millis());
		_ota_loop();
		delay(100);
	}
	
	if (_mqtt_enabled) {
		log("Core * Setting up MQTT...");
		_mqtt_setup();
	} else {
		log("Core * Not setting up MQTT since setMQTTData wasn't called");
	}
	log("Core * Setting up HTTP...");
	_http_setup();
	_initialized = true;
}

/**
 * Checks for updates.
 * 1. OTA, MQTT (if enabled) and HTTP are checked for incoming data.
 * 2. Methods defined using report_on() are checked if an update is in order.
 *
 * @note This method should be called from your loop() block.
 */
void SimpleIOT::loop() {
	_ota_loop();

	if (_mqtt_enabled) {
		_mqtt_loop();
	}
	
	_http_loop();
	_check_report_handlers();
}

void SimpleIOT::_wifi_setup() {
	WiFi.mode(WIFI_STA);
	WiFi.hostname(_hostname);
	WiFi.begin(_wifi_ssid, _wifi_pass);
	while (WiFi.waitForConnectResult() != WL_CONNECTED) {
		Serial.println("WiFi * Connection Failed! Rebooting...");
		delay(5000);
		ESP.restart();
	}
	Serial.println("WiFi * Ready");
	Serial.print("WiFi * IP address: ");
	Serial.println(WiFi.localIP());
}

void SimpleIOT::_ota_setup() {
	ArduinoOTA.onStart([]() {
		String type;
		if (ArduinoOTA.getCommand() == U_FLASH) {
			type = "sketch";
		} else {
			type = "filesystem";
		}
		Serial.println("OTA * Start updating " + type);
	});
	ArduinoOTA.onEnd([]() {
		Serial.println("\nOTA * End");
	});
	ArduinoOTA.onProgress([](unsigned int progress, unsigned int total) {
		Serial.printf("OTA * Progress: %u%%\r", (progress / (total / 100)));
	});
	ArduinoOTA.onError([](ota_error_t error) {
		Serial.printf("OTA * Error[%u]: ", error);
		if (error == OTA_AUTH_ERROR) Serial.println("Auth Failed");
		else if (error == OTA_BEGIN_ERROR) Serial.println("Begin Failed");
		else if (error == OTA_CONNECT_ERROR) Serial.println("Connect Failed");
		else if (error == OTA_RECEIVE_ERROR) Serial.println("Receive Failed");
		else if (error == OTA_END_ERROR) Serial.println("End Failed");
	});

	ArduinoOTA.setHostname(_hostname.c_str());
	ArduinoOTA.begin();
}

void SimpleIOT::_ota_loop() {
	ArduinoOTA.handle();
}



/***** MQTT stuff *****/
bool SimpleIOT::_mqtt_connect() {
	String topic(_mqtt_topic);
	topic.concat("status");
	bool connection_successful = false;
	if (_mqtt_auth) {
		connection_successful = _mqtt_client.connect(_hostname.c_str(), _mqtt_user, _mqtt_pass, topic.c_str(), 0, true, "OFFLINE");
	} else {
		connection_successful = _mqtt_client.connect(_hostname.c_str(), topic.c_str(), 0, true, "OFFLINE");
	}
	if (connection_successful) {
		char buffer[60];
		#ifdef DEBUG
		snprintf(buffer, 60, "ONLINE %s %s %s", _hostname.c_str(), WiFi.localIP().toString().c_str(), DEBUG);
		#else
		snprintf(buffer, 60, "ONLINE %s %s", _hostname.c_str(), WiFi.localIP().toString().c_str());
		#endif
		_mqtt_client.publish(topic.c_str(), buffer, true);

		topic = String(_mqtt_topic);
		topic.concat("#");
		_mqtt_client.subscribe(topic.c_str());
	} else {
		log("Could not connect to MQTT server. Error was %d", _mqtt_client.state());
	}
	return _mqtt_client.connected();
}

void SimpleIOT::_mqtt_setup() {
	_mqtt_client.setClient(_wifi);
	_mqtt_client.setServer(_mqtt_host, _mqtt_port);
	_mqtt_client.setCallback([=](char* t, byte* b, uint l){this->_mqtt_callback(t, b, l);});
	_mqtt_last_reconnect_attempt = 0;
}

void SimpleIOT::_mqtt_loop() {
	if (!_mqtt_client.connected()) {
		long now = millis();
		if (now - _mqtt_last_reconnect_attempt > 5000) {
			_mqtt_last_reconnect_attempt = now;
			if (_mqtt_connect()) {
				_mqtt_last_reconnect_attempt = 0;
			}
		}
	} else {
		_mqtt_client.loop();
	}
}

void SimpleIOT::_mqtt_callback(char* top, byte* pl, uint len) {
	pl[len] = '\0';
	String payload((char*)pl);
	String topic(top);

	if (topic.startsWith(_mqtt_topic)) {
		topic.remove(0, strlen(_mqtt_topic));
		_handle_action_request(topic, payload);
	}
}

void SimpleIOT::_mqtt_publish_report(String topic, String report) {
	publish(topic, report, true);	
}

/**
 * Publishes a message via MQTT.
 *
 * @param topic The topic to publish to. Will be appended to base_topic.
 * @param payload The payload to publish.
 * @param retain Whether to ask the broker to retain the message.
 */
void SimpleIOT::publish(String topic, String payload, bool retain) {
	if (!_mqtt_enabled) return;
	String final_topic = String(_mqtt_topic);
	final_topic.concat(topic);
	_mqtt_client.publish(final_topic.c_str(), payload.c_str(), retain);
}

/***** End of MQTT stuff *****/

/***** HTTP stuff *****/
void SimpleIOT::_http_setup() {
	_http_server.on("/", HTTP_GET, [=]() {
		String response("<html><head><title>");
		response.concat(_hostname);
		response.concat("</title></head><body><h1>");
		response.concat(_hostname);
		response.concat("<h2>Available report endpoints:</h2><table><tr><th>Endpoint</th><th>Last value</th></tr>");
		for (std::list<IOTReportHandler>::iterator it = _report_handlers.begin(); it!=_report_handlers.end(); ++it) {
			response.concat("<tr><td><a href='/");
			response.concat(it->topic);
			response.concat("'>/");
			response.concat(it->topic);
			response.concat("</a></td><td>");
			response.concat(it->last_result);
			response.concat("</td></tr>");
		}
		response.concat("</table>");

		response.concat("<h2>Available action endpoints:</h2><table>");
		for (std::list<IOTActionHandler>::iterator it = _action_handlers.begin(); it!=_action_handlers.end(); ++it) {
			response.concat("<form method='post' action='/");
			response.concat(it->topic);
			response.concat("'><tr><td>/");
			response.concat(it->topic);
			response.concat("</td><td><input type='text' placeholder='command' name='command' /><input type='submit' value='Send' /></td></tr></form>");
		}
		response.concat("</table></body></html>");
		_http_server.send(200, "text/html", response);
	});
	_http_server.onNotFound([=]() {
		String uri = _http_server.uri();
		uri.remove(0, 1); // Strip leading slash
		if (_http_server.method() == HTTP_POST) {
			if (_http_server.args()!=1 || !_has_action_handler(uri)) {
				_http_server.send(404, "text/plain", "Not found");
				return;
			}
			if (_handle_action_request(uri, _http_server.arg(0))) {
				_http_server.send(200, "text/plain", "OK");
				return;
			}
		} else if (_http_server.method() == HTTP_GET) {
			if (!_has_report_handler(uri)) {
				_http_server.send(404, "text/plain", "Not found");
				return;
			}
			String result = _handle_report_request(uri);
			if (result.length() > 0) {
				_http_server.send(200, "text/plain", result);
				return;
			}
		}

		_http_server.send(404, "text/plain", "Not found");
	});

	_http_server.begin(80);
}

void SimpleIOT::_http_loop() {
	_http_server.handleClient();
}
/***** End of HTTP stuff *****/

bool SimpleIOT::_topic_in_use(String topic) {
	return _has_action_handler(topic) || _has_report_handler(topic);
}

bool SimpleIOT::_has_action_handler(String topic) {
	for (std::list<IOTActionHandler>::iterator it = _action_handlers.begin(); it!=_action_handlers.end(); ++it) {
		if (topic.compareTo(it->topic)==0) {
			return true;
		}
	}
	return false;
}

bool SimpleIOT::_has_report_handler(String topic) {
	for (std::list<IOTReportHandler>::iterator it = _report_handlers.begin(); it!=_report_handlers.end(); ++it) {
		if (topic.compareTo(it->topic)==0) {
			return true;
		}
	}
	return false;
}

/**
 * Defines a topic to act on incoming data.
 * Actions can be anything you want: Switching on a relay, rebooting the ESP, ...
 * You can trigger this by (up to) two ways:
 * 1. By sending a POST request to `/<TOPIC>` containing the payload, or
 * 2. (If you enabled MQTT) by publishing your payload to `<BASE_TOPIC>/<TOPIC>`,
 *    with BASE_TOPIC being the topic you set with setMQTTData().
 *
 * @param topic The topic to listen on. May not start or end with a slash and may
 *              also not be `status` or `log`.
 * @param f A function to call when data for this topic is received. This function
 *          will be given the payload of the request.
 *
 * @returns true if the handler function could be added.
 */
bool SimpleIOT::act_on(String topic, IOTActionHandlerFunction f) {
	if (topic.startsWith("/") || topic.endsWith("/") || topic.compareTo("status")==0 || topic.compareTo("log")==0) {
		return false;
	}
	if (_topic_in_use(topic)) return false;
	_action_handlers.push_back({topic, f});
	return true;
}

/**
 * Defines a topic to report data to.
 * Data is reported using (up to) two ways:
 * 1. As answer to a GET request to `/<TOPIC>`, or
 * 2. (If you enabled MQTT) by publishing the data to `<BASE_TOPIC>/<TOPIC>`,
 *    with BASE_TOPIC being the topic you set with setMQTTData().
 *
 * @param topic The topic to listen on. May not start or end with a slash and may
 *              also not be `status` or `log`.
 * @param f A function to call to receive the current data for this topic. This
 *          function has to return the data as a String object.
 * @param interval Interval in ms to report this data. If MQTT is enabled, this
 *                 method will be called in this interval and its result will be
 *                 published.
 * @param cache_only If set to `true`, requesting the data via a GET request to
 *                   the HTTP server will return a cached value, unless at least
 *                   `interval` has passed. Use this for functions that take a
 *                   long time to process.
 *
 * @returns true if the handler function could be added.
 */
bool SimpleIOT::report_on(String topic, IOTReportHandlerFunction f, unsigned long interval, bool cache_only) {
	if (topic.startsWith("/") || topic.endsWith("/") || topic.compareTo("status")==0 || topic.compareTo("log")==0) {
		return false;
	}
	if (_topic_in_use(topic)) return false;
	IOTReportHandler h = {topic, f, cache_only, interval, 0, String()};
	_report_handlers.push_back(h);
	return true;
}

bool SimpleIOT::_handle_action_request(String topic, String payload) {
	for (std::list<IOTActionHandler>::iterator it = _action_handlers.begin(); it!=_action_handlers.end(); ++it) {
		if (it->topic.compareTo(topic)==0) {
			return it->function(payload);
		}
	}
	return false;
}

String SimpleIOT::_handle_report_request(String topic) {
	for (std::list<IOTReportHandler>::iterator it = _report_handlers.begin(); it!=_report_handlers.end(); ++it) {
		if (it->topic.compareTo(topic)==0) {
			return _get_report_handler_result(&(*it), false);
		}
	}
	return String();
}

String SimpleIOT::_get_report_handler_result(IOTReportHandler* handler, bool force_update) {
	if (handler->cache_only && !force_update) {
		return handler->last_result;
	} else {
		String result = handler->function();
		handler->last_result = result;
		return result;
	}
}

void SimpleIOT::_check_report_handlers() {
	unsigned long now = millis();
	for (std::list<IOTReportHandler>::iterator it = _report_handlers.begin(); it!=_report_handlers.end(); ++it) {
		bool update_needed = it->next_update_at <= now;
		bool millis_overflowed = it->next_update_at - it->update_interval > now;
		if (it->update_interval > 0 && (update_needed || millis_overflowed)) {
			log("Updating %s...", it->topic.c_str());
			String result = _get_report_handler_result(&(*it), true);
			if (_mqtt_enabled) {
				_mqtt_publish_report(it->topic, result);
			}
			it->next_update_at = now + it->update_interval;
		}
	}
}

/**
 * Log a message using Serial and (if enabled) MQTT.
 * Message is processed by sprintf with a maximum length of 128.
 * If MQTT is enabled, log messages are published at `<BASE_TOPIC>/log`.
 *
 * @param fmt Message or format string in printf syntax.
 * @param ... Data to insert into the format string.
 */
void SimpleIOT::log(const char* fmt, ...) {
	va_list arg;
	va_start(arg, fmt);
	char buffer[128];
	vsnprintf(buffer, 128, fmt, arg);
	va_end(arg);
	String topic(_mqtt_topic);
	topic.concat("log");
	if (_mqtt_enabled && _mqtt_client.connected()) {
		_mqtt_client.publish(topic.c_str(), buffer);
	}
	Serial.println(buffer);
}


/**
 * @example 01_actions.ino
 * 
 * This example shows how to use actions.
 * Three actions are defined:
 * * "on" and "off" set Pin 13 to HIGH or LOW, respectively. They are provided as lambda functions
 *   and don't care about the payload they're given.
 * * "set" is defined as function reference triggering only if the payload is "ON" or "OFF".
 *
 * You can access them via HTTP by sending a POST request to `/on`, `/off` or `/set`. There's a
 * nice-ish website at `/`, listing all registered handlers. If you kept the suggested hostname,
 * this link should lead you there: http://basic-test/.
 *
 * If you enabled MQTT, cou can also publish data via MQTT to trigger those actions. Again, if you
 * didn't change the topic below, it would be `basic-test-mqtt-topic/on`, `basic-test-mqtt-topic/off`
 * and `basic-test-mqtt-topic/set`.
 */

/**
 * @example 02_reports.ino
 * 
 * This example shows how to use reports.
 * Two reports are defined:
 * * "free_heap" returns the amount of free heap space. The method is provided as lambda function.
 *   Since `interval` ist set to 0, it will not be reported via MQTT, making it accessible only
 *   via HTTP.
 * * "uptime" returns the current uptime in seconds. It will be reported via MQTT (as long as you
 *   keep MQTT enabled) every 10 seconds. Since `cache_only` is set to true, the method won't be
 *   called any more often than that.
 *
 * You can access them via HTTP GET request to `/free_heap` or `/uptime`. There's a
 * nice-ish website at `/`, listing all registered reports along with (cached) data from them.
 * If you kept the suggested hostname, this link should lead you there: http://basic-test/.
 *
 * If you enabled MQTT, "free_heap" will be published every 10 seconds. Again, if you
 * didn't change the topic below, it would be `basic-test-mqtt-topic/free_heap`.
 */